Multi-task neural networks by learned contextual inputs
This paper explores learned-context neural networks. It is a multi-task
learning architecture based on a fully shared neural network and an augmented
input vector containing trainable task parameters. The architecture is
interesting due to its powerful task adaption mechanism, which facilitates a
low-dimensional task parameter space. Theoretically, we show that a scalar task
parameter is sufficient for universal approximation of all tasks, which is not
necessarily the case for more common architectures. Evidence towards the
practicality of such a small task parameter space is given empirically. The
task parameter space is found to be well-behaved, and simplifies workflows
related to updating models as new data arrives, and training new tasks when the
shared parameters are frozen. Additionally, the architecture displays
robustness towards cases with few data points. The architecture's performance
is compared to similar neural network architectures on ten datasets.
Anders T. Sandnes, Bjarne Grimstad, Odd Kolbjørnsen
# Multi-task neural networks by learned contextual inputs

###### Abstract

This paper explores learned-context neural networks. It is a multi-task learning architecture based on a fully shared neural network and an augmented input vector containing trainable task parameters. The architecture is interesting due to its powerful task adaption mechanism, which facilitates a low-dimensional task parameter space. Theoretically, we show that a scalar task parameter is sufficient for universal approximation of all tasks, which is not necessarily the case for more common architectures. Evidence towards the practicality of such a small task parameter space is given empirically. The task parameter space is found to be well-behaved, and simplifies workflows related to updating models as new data arrives, and training new tasks when the shared parameters are frozen. Additionally, the architecture displays robustness towards cases with few data points. The architecture's performance is compared to similar neural network architectures on ten datasets.

## 1 Introduction

A remarkable feat of nature is its ability to create endless variations on concepts that seem to be rigorously defined. Across all domains, we find classes of objects, clusters of phenomena, and groups of individuals, all of which seem to follow some overarching set of rules. And still, each separate instance puts a unique spin on the outcome. While this is fascinating to observe, it can be frustrating to model. One can quickly run into both performance issues and operational
challenges. An instance may have too few data points to produce a satisfying model. Or, there are just too many models to train and maintain.

Multi-task learning (MTL), as presented by Caruana (1997), is a learning strategy where multiple related models are trained simultaneously. The models share a subset of their parameters, which allows them to capture general domain knowledge. The purpose is to improve model generalization, by effectively training the shared parameters on more data.

Model frameworks based on the multi-task philosophy appear in many sciences. In the statistical literature they are, among others, referred to as mixed-, hierarchical-, multi-level-, or random-effect models. These methods see frequent use in sociology, economics, biometrics, and medicine (Demidenko, 2004; Raudenbush and Bryk, 2002). These models are often of moderate size and complexity. In the machine learning domain, there is broad diversity within transfer-learning and MTL (Lu et al., 2015; Zhang and Yang, 2021). Methods that combine transfer-learning or MTL with deep learning have seen success in several domains, with extensive research going into areas such as image analysis (Zamir et al., 2018; Kokkinos, 2017; Morid et al., 2021) and natural language processing (Raffel et al., 2020; Devlin et al., 2018; Brown et al., 2020). Engineering domains, such as solar and wind power (Wu et al., 2021; Dorado-Moreno et al., 2020), have also seen some MTL applications. However, in many cases, research is still needed to create satisfying solutions (Curreri et al., 2021).

Our focus is on non-linear regression problems. We are particularly interested in problems with many similar tasks that have complex input-output relationships. The tasks may have limited data. This class of problems captures many modeling challenges within engineering and industrial systems. Examples are cases with multiple instances of the same object, such as turbines in a wind farm, or batch processes, such as biomass accumulation in a fish farm or a field of crops. Here, each task typically has few observations, but they are, by design, quite similar. Solving such problems requires an architecture with a high degree of flexibility, where task adaptation can be done with few data points. Additionally, the architecture must facilitate the many operations needed to keep machine-learning models running in practice. Examples may be model updates in case of time-varying conditions, or the identification of new tasks that arrive at a later time.

We study _learned-context neural networks_. The architecture consists of two components. A neural network where all parameters are shared, and a set of task-specific parameter vectors. The task parameter vector serves as additional inputs to the neural network, which alters the network computation. This results in a powerful task adaptation mechanism, that can reduce the task parameter dimension significantly. The task parameter input provides contextual information about each task. They are referred to as a learned context because they are found during training. Learned contextual inputs facilitate the identification of a well-behaved task parameter space. By this, we mean a space that captures continuous latent properties of the tasks themselves, rather than just being a task encoding. A well-behaved task parameter space is desirable, because it enables us to train the shared network once, and then only
be concerned with the task parameter space for day-to-day operations. This is especially useful if a complete re-training takes significant manual effort, is computationally expensive, or new data points arrive frequently.

Variations of learned-context neural networks have, to our knowledge, only been applied in a meta-learning study by Zintgraf et al. (2019) and a virtual flow metering application by Sandnes et al. (2021). Zintgraf et al. (2019) proposed a model agnostic meta-learning method, which is concerned with context parameters in general. Learned-context neural networks are used as one of the concrete examples in an experimental study. Sandnes et al. (2021) use a variation of the architecture to study the benefit of multi-task learning for a soft sensing application. The focus is on the benefits of a multi-task architecture over the single-task learners traditionally used within the domain. The learned-context neural network itself has never been thoroughly analyzed.

### Contributions

We provide a deep dive into theoretical and practical aspects of the learned-context neural network architecture. We explore its task adaptation mechanism and relate it to existing methodology within statistics and machine learning, and we prove that a scalar task parameter is sufficient for the universal approximation of a set of tasks. The viability of such a low dimensional task parameter is studied empirically. The performance of the architecture is compared to similar architectures on ten datasets. Comparisons are made on predictive performance, sensitivity to dataset size, and the effect of the task parameter dimension.

### Outline

The content of the paper is organized as follows. Section 2 presents notation and gives an introduction to mixed models, multi-task learning, and methods related to the architecture of interest. Section 3 gives a theoretical description of learned-context neural networks. Section 4 dives deeper into the task adaptation mechanism of the learned-context network. Section 5 presents an empirical analysis of the architecture and compares is to related methods. Section 6 summarizes and discusses the findings, and Section 7 concludes the paper.

## 2 Background

Here we present a minimal background required to study learned-context neural networks. For a broader picture, see Demidenko (2004) for more statistical methods or Zhang and Yang (2021) for machine learning applications.

### Notation

We consider a set of \(m\) tasks \(\left\{(\mathcal{D}_{j},f_{j})\right\}_{j=1}^{m}\). Each task consists of a set of observations \(\mathcal{D}_{j}=\left\{(x_{ij},y_{ij})\right\}_{i=1}^{n_{j}}\) and a function \(f_{j}:\mathbf{R}^{d_{x}}\mapsto\mathbf{R}\). These are
related by \(y_{ij}=f_{j}(x_{ij})+\epsilon_{ij}\), where \(\epsilon_{ij}\sim\mathcal{N}(0,\sigma_{\epsilon}^{2})\). The tasks are homogeneous, with the elements of \(x\) and \(y\) representing the same quantities across all tasks. The indicator variable \(c_{j}\in\left\{0,1\right\}^{m}\) is a one-hot task encoding where the \(j\)th element is one and the rest zeros.

We are interested in the case of knowledge transfer through hard parameter sharing, and task adaptation is achieved using a small set of task parameters. We use \(\alpha\) to denote a set of shared parameters, and \(\beta_{j}\in\mathbf{R}^{d_{\beta}}\) to denote a vector of task specific parameters. The general form of the problem of interest is \(y_{ij}=f(x_{ij};\alpha,\beta_{j})+\epsilon_{ij}\).

### Mixed models and multi-task neural networks

The concept of multi-task learning has been studied extensively for linear models under a variety of names and contexts. The simplest form is a _varying-intercept_ linear model,

\[y_{ij}=a^{\top}x_{ij}+b+\beta_{j}+\epsilon_{ij}, \tag{1}\]

where each task is a linear function of the observed variables and tasks differ only in the individual bias terms \(\beta_{j}\)(Gelman et al., 2013). With such a setup it is common to assume task parameters are distributed according to \(\beta_{j}\sim\mathcal{N}(0,\sigma_{\beta}^{2})\). Task-specific slope parameters are also common. These models, known as multilevel-, hierarchical-, or mixed effect models, are extensively described in the statistical literature (Demidenko, 2004; Raudenbush and Bryk, 2002). An extension to the linear model is to factorize the parameters into a task component and a shared component, a notable example being the group-and-overlap model of Kumar and Daume III (2012), \(y_{ij}=a_{j}^{\top}x_{ij}+b_{j}+\epsilon_{ij}\), where the parameters are found as

\[\begin{bmatrix}a_{j}\\ b_{j}\end{bmatrix}=L\beta_{j}. \tag{2}\]

Task-specific parameters are linear combinations of latent basis tasks, given as the columns of \(L\). A tailored algorithm is used to jointly optimize the latent task matrix \(L\) and the individual task parameters. The optimization is controlled by the desired number of latent tasks and their sparsity. This structure allows the degree of sharing between tasks to be found during training. The mechanism introduces a robustness towards outlier tasks, because tasks that do not share any aspects with the others can be isolated to separate columns of \(L\). This allows the main tasks to be differentiated without inference from potential outliers.

Moving beyond linear model structures introduces several design choices. The simplest is fixed nonlinear analytical expressions where parameters can be partially shared, task-specific, or found through expressions such as Equation 2. A classic example is nonlinear growth curve models, commonly formulated as logistic curves (Demidenko, 2004).

Alternatively, the nonlinearities can be selected from a predefined set of candidates as proposed by Argyriou et al. (2008). This produces the expression
\(y_{ij}=\sum_{k=1}^{d_{\beta}}\beta_{j,k}\phi_{k}(x_{ij})+\epsilon_{ij}\), where the knowledge sharing mechanism lies in the shared choice of basis functions \(\phi_{k}\).

If the space of basis functions is not predefined, a more flexible solution is to allow them to be learned by a multi-task neural network (Caruana, 1997). This can be expressed as

\[y_{ij}=\beta_{j}^{\top}h(x_{ij};\alpha)+\epsilon_{ij}, \tag{3}\]

where \(h\) is a neural network with \(d_{\beta}\) outputs, parametrized by \(\alpha\). The task-parameters \(\beta_{j}\) represents a task specific output layer.

A different neural network strategy is the context-sensitive networks of Silver et al. (2008),

\[y_{ij}=h(z_{ij};\alpha)+\epsilon_{ij},\ z_{ij}=\begin{bmatrix}x_{ ij}\\ c_{j}\end{bmatrix}, \tag{4}\]

where all parameters of the neural networks are shared. Tasks are differentiated through the additional input \(c_{j}\), which is a one-hot task encoding. This leads to networks where the input dimension grows with the number of tasks. Context-adaptation, introduced by Zintgraf et al. (2019), replaces the one-hot encoding with a trainable vector of task parameters. This allows the dimension of the input space to remain fixed while the number of tasks grows. The learned-context networks presented here are based on this idea.

Other, more complex neural network architectures allow all network weights to have some level of task adaptation. Examples are tensor factorization (Yang and Hospedales, 2017), and masked networks (Mallya et al., 2018). Alternatively, one can introduce layers that facilitate sharing between otherwise disjoint neural networks (Misra et al., 2016). We are, however, focusing our attention on simple feed-forward architectures with few task parameters.

### Related learning paradigms

Multi-task learning resembles other learning paradigms, in particular _transfer-learning_ and _meta-learning_. The distinctions between these can be difficult, and their nuances have changed over time. We base our discussion on the recent survey of Hospedales et al. (2022).

Transfer-learning attempts to improve the performance of a task using information from related source tasks. This can, for instance, be achieved by copying an existing model trained on the source tasks and applying optional fine-tuning steps to the parameters. This tuning has no concern for the performance of the source tasks. In contrast, multi-task learning attempts to jointly optimize the performance of all tasks.

Meta-learning also attempts to jointly optimize all tasks, but the objective of the optimization is different. Multi-task learning is concerned with a fixed set of provided tasks and produces a single joint model. Meta-learning, on the other hand, attempts to improve learning for the whole distribution of tasks, which can include unseen future tasks. The output is a _learning procedure_ that
can be applied to all tasks in the task distribution. As such, meta-learning can draw parallels to _hyper-parameter optimization_. However, hyper-parameter optimization only considers performance on the current learning problem, while meta-learning considers the performance over a family of learning problems that may occur.

Special cases of these optimization problems may be closely linked in practice. Grant et al. (2018) connects a particular meta-learning algorithm to hierarchical Bayesian models and hyper-parameter optimization. Zintgraf et al. (2019) develops this meta-learning algorithm further with context-adaptation. The multi-task learned-context neural networks studied in this paper is similar to one of the context-adaptation mechanisms.

## 3 Model description

We explore a particular version the context-adaptation architecture, denoted by \(y_{ij}=f(x_{ij};\beta_{j},\alpha)+\epsilon_{ij}\). The architecture consists of a residual feedforward neural network, \(y_{ij}=h(z_{ij};\alpha)+\epsilon_{ij}\), which is shared between all tasks. Tasks are differentiated by trainable parameters \(\beta_{j}\), which is used as input to the network,

\[z_{ij}=\begin{bmatrix}x_{ij}\\ \beta_{j}\end{bmatrix}. \tag{5}\]

We refer to this architecture as a learned-context neural network.

### Residual feedforward networks

We consider feedforward neural network architectures (Goodfellow et al., 2016), with residual skip connections (He et al., 2016). In the following we omit the \(ij\) subscripts and noise terms to simplify the notation.

We denote a residual feedforward network by \(y=h(z;\alpha)\). It consists of \(K\) layers parametrized by \(\alpha=\left\{(W_{k},b_{k})|k=1,\ldots,K\right\}\). The network has three main elements. The first linear layer,

\[z^{(2)}=W_{1}z+b_{1}, \tag{6}\]

a sequence of residual blocks

\[z^{(k+1)} =W_{k}g\left(z^{(k)}\right)+b_{k}\] \[z^{(k+2)} =z^{(k)}+W_{k+1}g\left(z^{(k+1)}\right)+b_{k+1}\Bigg{\}}\,k=2,4, \ldots,K-2, \tag{7}\]

and the final linear layer

\[y=W_{K}z^{(K)}+b_{K}. \tag{8}\]

The residual blocks in Equation 7 is illustrated in Figure 1. We let the residual skip connection span two hidden layers, and use a are pre-activated structure (He
et al., 2016). All hidden layers have the same dimension. The ReLU function, \(g(z)=\max(0,z)\), where the max operation is performed element-wise, is used for activation.

### Comparable architectures

Throughout the analysis of the learned-context neural network, we will make comparisons with two closely related architectures. The first is the context-sensitive networks as presented by Silver et al. (2008), given in Equation 4. The second is the classic MTL architecture of Caruana (1997), which we will refer to as last-layer MTL. This is described by Equation 3.

## 4 Task adaptation mechanism

This section explores the theoretical aspects of the task adaptation mechanism in the learned-context neural network. First, the relationship between learned-context neural networks, linear mixed models, and context-sensitive networks are discussed. Then, a simple example illustrates how the task parameters can influence the output of the neural network. Finally, the universal approximation power of learned-context networks are investigated and shown to be equal to that of using a separate neural network for each task.

### Task parameter input yield varying-bias layers

Recall the residual neural network described in Equations 6-8. The augmented input to the first layer, given in Equation 5, consists of observations \(x_{ij}\) and task parameters \(\beta_{j}\). Inspecting the first linear layer, \(z_{ij}^{(2)}=W_{1}z_{ij}+b_{1}\), we note

Figure 1: Residual block spanning two hidden layers, each layer consisting of a activation function and a linear transformation. The block output is the sum of the block input \(z^{(k)}\) and the residual component \(\Delta^{(k)}\).
that partitioning the weight matrix according to \(W_{1}=\begin{bmatrix}A&L\end{bmatrix}\) yields

\[z_{ij}^{(2)}=Ax_{ij}+L\beta_{j}+b_{1}. \tag{9}\]

This is recognized as the linear varying-intercept model from Equation 1, where \(b_{1}\) is the population average intercept and \(\tilde{\beta}_{j}=L\beta_{j}\) is the individual intercept. However, linear varying-intercept models are usually concerned with scalar, or low dimensional, response variables. The hidden layer \(z_{ij}^{(2)}\) may be of high dimension, and the term \(L\beta_{j}\) allows for a low-rank approximation of the task intercept. This is the same mechanism utilized in the group-and-overlap method given in Equation 2, and the method inherits the benefits this provides in terms of outlier robustness.

On the other hand, if the number of task parameters equals the number of tasks, the task mechanism is similar to the one-hot encoding of context-sensitive networks. To see this, repeat the exercise using the augmented input from Equation 4. This leads to \(z_{ij}^{(2)}=Ax_{ij}+Bc_{j}+b_{1}\), where the contextual input selects the \(j\)th column of \(B\) as the task intercept. With the task parameter dimension equal to the number of tasks, each of the first layer biases can be set individually for each task. This essentially makes each task responsible for a number of parameters equal to the hidden layer size, which can be a significant number even for moderately sized networks.

### Task adaptation examples

Considering Equation 9 alone, the task adaptation mechanism may seem limited. We study a simple example to gain insight into how changing the biases affects the network response. The example is adapted from Telgarsky (2016). A network with three layers is constructed to linearly interpolate the points \((x,y)=(0,0),(0.5,1),(1,0),(1.5,1),(2,0)\), given task parameters equal to zero. We refer to this as the base shape. It is illustrated as the bold dark line in Figure 2. A detailed description of the network weights is given in Appendix A.

We now explore the effect of different values for the task parameter component \(L\) from Equation 9. Starting with the simple case where \(L=A\), which yields a translation of the base shape. This can be seen by manipulating Equation 6, \(z_{ij}^{(2)}=Ax_{ij}+A\beta_{j}+b_{1}=A\left(x_{ij}+\beta_{j}\right)+b_{1}\), which is the expression for translation \(y_{ij}=f(x_{ij}+\beta_{j})\). This is illustrated in the top part of Figure 2. A different transform is obtained by setting \(L=b_{1}\). For this particular network, it yields dilation by a factor \(1+\beta\). This changes the base shape of the output to a dilated shape according to \((x,y)\mapsto((1+\beta)x,(1+\beta)y)\). This is illustrated in the bottom part of Figure 2. The derivation of this result is given in Appendix A.1.

These simple examples illustrate how the contextual inputs can produce different transformations by changing the first layer weights. While looking limiting at first glance, the ability to influence the first layer bias creates changes that propagate through the entire network, enabling powerful nonlinear transformations.
Figure 2: Evaluation of the simple learned-context network described in Appendix A. The two examples only differ in their specification of the matrix \(L\) from Equation 9. Top is a translation case, with \(L=A\). Bottom is a dilation case, with \(L=b_{1}\). Both have been evaluated with task parameter \(\beta\) fixed at three different values, with the bold black line representing the base network.
### Universal task approximation

The previous section illustrates how a combination of contextual inputs and deep neural networks can produce complex task adaptations. Here we make this notion more precise by proving a multi-task equivalent to the universal approximation property. The main result is that any task adaption is achievable by learned-context neural networks regardless of how many task parameters are used, provided that the shared neural network is sufficiently wide and deep.

Universal approximation refers to the ability of neural network architectures to approximate a function with arbitrary precision, given a set of qualifiers on the function and the network (Cybenko, 1989; Hornik et al., 1989). Several variations of universal approximation properties have been shown for different neural network architectures, including ReLU-based feedforward networks (Lu et al., 2017; Kidger and Lyons, 2020). Here we explore how universal approximation properties carry over to the learned-context networks, and compare it to the properties of context-sensitive networks and last-layer MTL networks. We only consider traditional feedforward networks without the residual skip connections, as these do not change the results presented here.

We adapt the definitions from (Kidger and Lyons, 2020) and rely upon their main result for universal approximation. Our interest is to study the number of task parameters required, so the dimensions of the neural networks themselves are omitted.

**Definition 1**.: _Let \(\mathcal{F}_{k,l}\) be the class of functions described by feedforward neural networks with input dimension \(k\), output dimension \(l\), an arbitrary number of hidden layers of arbitrary size, where the ReLU activation function is applied to all hidden units._

**Definition 2**.: _Let \(C(K)\) be the space of continuous functions \(f:K\rightarrow\mathbf{R}\), with domain \(K\subseteq\mathbf{R}^{d_{x}}\), where \(K\) is compact._

Consider \(m\) tasks given by \(y=f_{j}(x),\ j=1,\ldots,m\), where \(f_{j}\in C(K)\). Individually, the functions \(f_{j}\) can be approximated arbitrarily close by a sufficiently sized neural network \(\hat{f}_{j}\in\mathcal{F}_{d_{x},1}\)(Kidger and Lyons, 2020). Extending this notion to MTL architectures means their ability to approximate the composed function

\[y=f(x,j)=\sum_{k=1}^{m}I(j=k)f_{k}(x). \tag{10}\]

Here, \(I(j=k)\) is the indicator function, tanking the value one if \(j=k\) and zero otherwise.

In Equation 10, the inputs of \(f(x,j)\) closely resembles those of context-sensitive networks, making this a natural starting point.

**Proposition 1**.: _There exists a context-sensitive neural network from the class \(\mathcal{F}_{d_{x}+m,1}\) that is arbitrarily close to the multi-task function \(f\) of Equation 10 with respect to the uniform norm._
Proof.: An equivalent computation to Equation 10 is \(f(x,c_{j})=\sum_{k=1}^{m}c_{j,k}f_{k}(x)\), where \(c_{j,k}\) is the \(k\)th element of \(c_{j}\). Relaxing the indicator variable domain to \(c_{j}\in[0,1]^{m}\), allows the context-sensitive input vector, \(z\), to exist in a compact subspace \(K^{+}\subseteq\mathbf{R}^{d_{x}+m}\). The relaxation gives \(f\in C(K^{+})\), a space of which the class \(\mathcal{F}_{d_{x}+m,1}\) is dense with respect to the uniform norm. It follows that context-sensitive networks can approximate the set of task functions arbitrarily well. 

This result means that the context-sensitive architecture is sufficiently flexible to achieve the same approximation power as using an individual neural network for each task. However, it does not say anything about the number of parameters required to achieve this.

Learned-context neural networks can achieve the same universal approximation power as context-sensitive networks, using only a scalar task parameter.

**Theorem 2**.: _There exists a learned-context neural network from the class \(\mathcal{F}_{d_{x}+1,1}\) and a set of task parameters \(\beta_{j}\in\mathbf{R},\ j=1,\ldots,m\), that is arbitrarily close to the multi-task function \(f\) of Equation 10 with respect to the uniform norm._

Proof.: The proof is to construct the mapping \((x,\beta_{j})\mapsto(x,c_{j})\), as the first two hidden layers of a neural network. The remainder of the network could then be taken as a context-sensitive network.

First, we construct a similar triangle wave as in Figure 2 for the task parameters. Let the task parameters be assigned values \(\beta_{j}=j\), which is a trivial task encoding. The first layer is assigned the weights

\[W_{1}=1_{2m},\ b_{1}^{\top}=-\begin{bmatrix}1-\delta&1&2-\delta&2&\ldots&m- \delta&m\end{bmatrix},\]

where \(1_{2m}\) is a vector of \(2m\) ones and \(\delta\in(0,0.5)\) is a number determining the slope and spacing of the triangles. After ReLU activation, this gives a shifted sequence of the task parameter. The second layer is assigned the weights

\[W_{2}=\frac{1}{\delta}I_{m}\otimes\begin{bmatrix}1&-2\end{bmatrix},\ b_{2}=0_{m}\]

where \(\otimes\) denotes the Kronecker product and \(0_{m}\) is a vector of \(m\) zeros. This leads to the \(j\)th entry of \(g(z^{(3)})\) to be given by

\[g(z^{(3)}_{j})=\begin{cases}(\beta-j+\delta)/\delta&\text{if }\beta\in[j- \delta,j),\\ (-\beta+j+\delta)/\delta&\text{if }\beta\in[j,j+\delta),\\ 0&\text{otherwise }.\end{cases} \tag{11}\]

Only one of the entries of \(g(z^{(3)}_{j})\) will be non-zero at a time, and when evaluated at \(\beta=j\) the \(j\)th entry will be one and the rest zero. The sharpness of the transition between zero and one is adjusted by \(\delta\). The weights (\(W_{1}\), \(b_{1}\), \(W_{2}\), and \(b_{2}\)) can now be augmented to store a copy of the other input values \(x\), which makes the hidden layer \(z^{(3)}\) the same as the input layer of a context-sensitive network. Therefore, a learned-context network with a scalar task parameter share approximation properties with the context-sensitive networks from Proposition 1.
This result only considers the number of task parameters and does not put any bounds on the number of shared parameters. While a scalar task parameter is theoretically sufficient, it does not indicate that this is the ideal choice. Indeed, the proof is based on the scalar task parameter taking an indicator role, which is counterproductive to our desire to learn the latent properties of the tasks. This is useful to bear in mind when selecting hyperparameters, as too few task parameters may force an indicator behavior.

The last-layer MTL architecture is less flexible in its task adaptation mechanism. As such, it may require more task parameters than the learned-context networks.

**Proposition 3**.: _Last-layer MTL networks with base network from the class \(\mathcal{F}_{d_{x},k}\) and \(\beta_{j}\in\mathbf{R}^{k}\) require \(k\geq m\) to guarantee the existence of a function arbitrarily close to the multi-task function \(f\) in Equation 10._

Proof.: The last-layer network is a linear combination of \(n_{\beta}\) basis functions, \(y=\beta_{j}^{\top}h(x)\). Let the tasks be given as \(m\) sine waves \(y=\sin(\omega_{j}x)\), with frequency \(\omega_{j}=j\) and \(x\in[-\pi,\pi]\). These cannot be constructed by a set of basis functions of dimension lower than \(m\), because they are orthogonal to each other. Hence, in this case the shared neural network cannot have an output dimension less than \(m\). 

Task adaption in the last-layer MTL architecture is through linear combination of features produced by the fully shared neural network. This means that it is limited to scaling the given nonlinearities. Contrary, in the learned-context architecture the task parameters can influence the nonlinearities themselves, which leads to a much broader range of adaptions it can achieve.

## 5 Empirical analysis

This section investigates the properties of learned-context neural networks by empirically. First, the learned-context neural network is compared to similar architectures. The architectures are evaluated on predictive performance, training robustness, sensitivity to the number of data points, and the effect of the number of task parameters. The task parameter space produced by learned-context neural networks is then explored by estimating task parameters for new tasks after the shared parameters are trained and fixed.

### Benchmark models

Learned-context neural networks, described in Section 3, are compared to learned-context networks are compared to last-layer multi-task network described by Equation 3, and context-sensitive network described by Equation 4. All three network models use the residual structure described in Equations 6-8.

A linear mixed-effect model acts as a baseline for performance. The linear models have a set of shared slope and intercept parameters. Additionally, each
task is has its own intercept. This structure is given in Equation 1. Discrete features are one-hot encoded for this model.

### Datasets

The models are compared on ten datasets. Two of the datasets are synthetically created to highlight differences between the architectures. Three datasets, Schools, Parkinson, and Sarcos, are common to the MTL literature (Zhang and Yang, 2021), while the last five are selected from a range of open access data sources. All datasets correspond to regression problems with a scalar response variable. The datasets are described below and summarized in Table 1.

**Frequency** is a synthetic dataset where tasks are sine waves of different frequency. Data is generated according to \(y_{ij}=0.5\sin(2\pi\omega_{j}x_{ij})+0.5+\epsilon_{ij}\). The task frequencies are sampled as \(\omega_{j}\sim\text{Uniform}(0.5,4.0)\). We take \(x_{ij}\sim\text{Uniform}(0,1)\), and \(\epsilon_{ij}\sim\mathcal{N}(0,\sigma^{2})\), \(\sigma=0.1\).

**Sine and line** is a synthetic dataset proposed by Finn et al. (2018). Tasks are taken from two classes, affine functions, \(y_{ij}=a_{j}x_{ij}+b_{j}+\epsilon_{ij}\) or sine waves, \(y_{ij}=c_{j}\sin(x_{ij}+d_{j})+\epsilon_{ij}\). We generate an equal number of tasks from each class, with parameters sampled as \(a_{j},b_{j}\sim\text{Uniform}(-3,3)\), \(c_{j}\sim\text{Uniform}(0.1,5.0)\), and \(d_{j}\sim\text{Uniform}(0,\pi)\). For both task types we use \(\epsilon_{ij}\sim\mathcal{N}(0,\sigma^{2})\), \(\sigma=0.3\), and \(x_{ij}\sim\text{Uniform}(-5,5)\). We note that this problem can be represented as linear regression, \(y=\beta^{\top}z\), on a set of nonlinear basis functions \(z^{\top}=\begin{bmatrix}1&x&\sin(x)&\cos(x)\end{bmatrix}\) by applying the trigonometric identity \(\sin(x+d)=\sin(d)\cos(x)+\cos(d)\sin(x)\).

**Schools** is a dataset of examination results for students from different schools over a three year period (Nuttall et al., 1989). The goal is to map features of the schools and students to student performance, in order to study school effectiveness. We treat each school as a task. The dataset is provided by Center of Multilevel Modelling (1987).

**Parkinson** telemonitoring relates patient observations to a disease symptom score (Tsanas et al., 2010). Each patient is considered a task. It is provided by UCI Machine Learning Repository (2009).

**Bike sharing** relates weather and calendar features to the number of trips of a bike sharing system over two years (Fanaee-T and Gama, 2014). Each month is considered a task, which allows tasks to capture seasonal effects and potential changes to the bike sharing system itself. Data is provided by UCI Machine Learning Repository (2013).

**Inflation** is a dataset from OECD (2022). It considers the development of Consumer Price Index (CPI) in 45 countries. CPI is taken quarterly from 1970 to 2020, normalized with 2015 being equal to 100% for each country. Each country is a task. Time is the only input variable.

**Obesity** is a dataset that describes the development of mean body-mass index from 1985 to 2019 in 200 countries (NCD Risk Factor Collaboration, 2020a). Input variables are age, sex, and year. The response variable is the percentage considered to be obese. Each country is a task. Data is provided by NCD Risk Factor Collaboration (2020c).
**Height** is a dataset with the same structure as Obesity, but the response variable is the mean height within the groups. Data is provided by NCD Risk Factor Collaboration (2020b).

**Death rate** describes the rate of death in different age groups. Input variables are age and sex. Each pair of year and country is considered a task. There are 183 countries and five years. This results in a greater number of tasks than Obesity and Height, but fewer input variables. Data is provided by World Health Organization (2020)

**Sarcos** is data from a robotic arm with seven joints. The goal is to map joint position, velocity, and acceleration to motor torque (Vijayakumar and Schaal, 2000). Each joint is one task, taking all 21 joint measurements as input to predict torque at the joint. As a consequence, each task has exactly the same input data, differing only in their response. It is hosted by Gaussian Processes for Machine Learning (2000).

### Training Procedure

Model parameters are found by minimizing a loss function of mean squared prediction error and parameter regularization (Hastie et al., 2009),

\[\min_{\alpha,\beta_{1},\ldots,\beta_{m}}\ \frac{1}{n}\sum_{j=1}^{m}\sum_{i=1}^{ n_{j}}\left(y_{ij}-f(x_{ij};\beta_{j},\alpha)\right)^{2}+\lambda_{\alpha}l( \alpha)+\lambda_{\beta}\sum_{j=1}^{m}l(\beta_{j}). \tag{12}\]

The prediction error is divided by the total number of data points \(n\). Parameters are regularized by the L2 norm, scaled by factors \(\lambda_{\alpha}\) and \(\lambda_{\beta}\) for shared and task parameters respectively. The task parameter term is ignored for context-sensitive networks. All data points from all tasks are weighted equally. More complex loss functions with individual weights for each task could potentially improve performance in some cases (Kendall et al., 2018; Gong et al., 2019), but this is not considered here.

\begin{table}
\begin{tabular}{l r r r r r} \hline Name & Num. feat. & Num. tasks & Data train & Data test & Std. \(y\) \\ \hline Frequency & 1 & 250 & 30000 & 25000 & 0.36 \\ Sine and line & 1 & 100 & 6000 & 10000 & 3.93 \\ Schools & 7 & 139 & 12339 & 3023 & 12.72 \\ Parkinson & 16 & 42 & 4717 & 1158 & 10.70 \\ Bike sharing & 6 & 24 & 13915 & 3464 & 181.38 \\ Inflation & 1 & 45 & 6684 & 1344 & 34.61 \\ Obesity & 3 & 200 & 126000 & 126000 & 5.22 \\ Height & 3 & 200 & 105000 & 105000 & 19.64 \\ Death rate & 2 & 915 & 16470 & 16470 & 0.12 \\ Sarcos & 21 & 7 & 311388 & 31143 & 18.84 \\ \hline \end{tabular}
\end{table}
Table 1: Summary of dataset properties. Listed is the number of features in the dataset, the number of tasks, the number of data points used for model training and testing, and the standard deviation of the response variable \(y\).
Optimization of Equation 12 is done by stochastic gradient decent (Botto et al., 2018), with a two stage learning rate schedule. Hyperparameters are optimized by the global optimizer LIPO (Malherbe and Vayatis, 2017), which is run for a predetermined number of iterations. Further details is given in Appendix B for the training procedure and Appendix C for hyperparameters.

### Base performance

The three neural network architectures and the linear mixed model are first compared in terms of test set errors. The results are summarized in Table 2. The neural network models have similar performance on all datasets, with learned-context and last-layer MTL being slightly favoured over context-sensitive networks in general. In some cases, such as Obesity and Height, we see that the context-sensitive architecture has close to twice the error of the other architectures. While this difference is significant, we note that when the errors are compared to the standard deviation of the response variables, given in Table 1, these errors are all quite small. For the synthetic datasets, learned-context network perform better on the Frequency data, and last-layer MTL network perform better on the Sine and line data. This is expected, because these architectures match the data generating processes. The linear mixed model performed well on two of the datasets, Schools and Parkinson. In these cases, the neural network models are able to achieve similar performance, which is reassuring.

To compare the robustness of the training procedure we re-run the training of all neural network models. We reuse the hyperparameters and run the training nine additional times. The best and worst RMSE values for the ten runs combined are given in Table 3. It also lists the number of times training runs diverged and had to be restarted. Overall the results are quite consistent. However, there are cases where the training appears less stable. An example

\begin{table}
\begin{tabular}{l l l l l} \hline \hline Dataset & LC & CS & LL & LME \\ \hline Frequency & **0.106** & 0.136 & 0.122 & 0.360 \\ Sine and line & 0.325 & 0.34 & **0.316** & 3.843 \\ Schools & 10.203 & 10.363 & 10.314 & **10.108** \\ Parkinson & 2.914 & 2.856 & **2.670** & 2.776 \\ Bike sharing & 53.869 & 80.208 & **45.043** & 142.087 \\ Inflation & **1.833** & 2.526 & 2.501 & 11.095 \\ Obesity & **0.123** & 0.281 & 0.210 & 2.512 \\ Height & 0.394 & 0.61 & **0.347** & 5.055 \\ Death rate & **0.011** & 0.017 & 0.026 & 0.078 \\ Sarcos & **2.176** & 2.188 & 2.482 & 10.653 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Root mean squared error on test data for all models. The best score on each dataset is highlighted. The column headers are learned-context neural networks (LC), context-sensitive networks (CS), last-layer MTL networks (LL), and linear mixed effect models (LME).
being last-layer MTL on the Height dataset, which yield a large span in relative errors. Again we note that these errors are small in absolute value, which makes such comparisons sensitive to the randomness in the training procedure.

### Effect of dataset size

We now compare the sensitivity to dataset size, by training the neural network architectures on reduced datasets. The training datasets are cut to 50% and 10% of their original size. The same fraction of data is removed from each task, keeping the data balance from the original dataset. Test sets are kept at full size. Training and hyperparameter searches are conducted in the same way as for the full data case. The results are summarized in Table 4. Context-sensitive networks are on average slightly behind the others on all data sizes. A reduction in data size naturally leads to a reduction in performance for all models, but the learned-context architecture is less sensitive to this than the others.

### Effect of task parameter dimension

Section 4.3 established theoretical differences in the number of task parameters required by learned-context networks and last-layer MTL networks. In this section we explore the practical impact of the task parameter dimension. To this end, learned-context networks and last-layer networks are trained on all datasets with different number of task parameters. All other hyperparameters are fixed. The models are trained on the full dataset. Figure 3 illustrates the results. Additional task parameters generally improve performance of both models on all datasets, up to a certain point. There does not seem to be a significant downside of excessively large task parameter dimensions. This means the hyperparameter searches will likely arrive at a larger than necessary value,

\begin{table}
\begin{tabular}{l l l l} \hline Dataset & LC & CS & LL \\ \hline Frequency & 1.00, 1.01 & 1.10, 1.29 (1) & 1.16, 1.21 \\ Sine and line & 1.00, 1.00 & 1.04, 1.06 & 0.97, 0.97 \\ Schools & 1.00, 1.00 & 1.01, 1.02 & 1.01, 1.02 \\ Parkinson & 0.95, 1.00 & 0.92, 0.98 & 0.90, 0.94 \\ Bike sharing & 0.94, 1.01 & 1.29, 1.51 & 0.82, 0.85 \\ Inflation & 0.98, 1.09 & 1.33, 1.58 & 1.35, 1.38 \\ Obesity & 0.94, 1.08 (1) & 2.10, 2.45 & 1.30, 1.86 \\ Height & 0.99, 1.14 & 1.53, 1.65 & 0.82, 1.37 (1) \\ Death rate & 0.99, 1.04 & 1.52, 1.62 & 2.24, 2.30 \\ Sarcos & 1.00, 1.04 & 0.99, 1.05 (1) & 1.11, 1.49 \\ \hline \end{tabular}
\end{table}
Table 3: Results from ten repeated training runs with the same hyperparameter settings. Reported are the min and max relative RMSE value in each case. The values are normalized by the performance of the learned-context neural network in Table 2. The number of diverging training runs, if any, are given in brackets.
unless additional model selection penalty terms are introduced. Overall, the learned-context neural networks achieve better performance with fewer task parameters, but the gap is closed as the number increases.

As noted in Section 5.2, the sine and line dataset is easily represented as regression on four nonlinear basis functions. This is reflected by both models performing identically for four or more task parameters. The frequency dataset does not map to linear regression in the same way. As a consequence, the last-layer MTL network requires a significantly higher number of task parameters to achieve the same performance as the learned-context neural network.

A benefit of a low-dimensional task parameter space is that it simplifies visualization and interpretation of the task parameters. Examples of this are given in Appendix D, where the models for Inflation, Bike sharing, and Death rate datasets are studied in detail. The task parameters are given sensible interpretations related to properties of the model domain.

### Hold-out task performance

A key aspect differentiating the learned-context from a fixed encoding, such as the context-sensitive neural networks, is that similar tasks can be assigned similar task parameter values. A desirable feature would be that the parameters capture latent properties of the tasks. For this to happen, the task parameter space must in some sense be well-behaved. A qualitative study, given in Appendix D, supports the hypothesis that learned-context neural networks facilitate such behavior, and the task parameters capture properties fundamental to the domain, as opposed to separating tasks by memorization.

This section attempts to quantify this behavior by studying the viability of estimating parameters for a new task after the shared neural network parameters are trained and fixed. This is similar to meta-learning. However, as discussed in Section 2.3, meta-learning takes this generalization into account in the original

\begin{table}
\begin{tabular}{l|l l l|l l l|l l l} \multicolumn{4}{c}{100\% training data} & \multicolumn{2}{c}{50\% training data} & \multicolumn{2}{c}{10\% training data} \\ \hline Dataset & LC & CS & LL & LC & CS & LL & LC & CS & LL \\ \hline Frequency & **0.29** & 0.37 & 0.33 & **0.31** & 0.36 & 0.37 & **0.49** & 0.87 & 0.53 \\ Sine and line & 0.08 & 0.09 & **0.08** & 0.09 & 0.11 & **0.09** & **0.25** & 0.40 & 0.50 \\ Schools & **0.80** & 0.81 & 0.81 & **0.84** & 0.84 & 0.93 & **0.88** & 0.98 & 1.09 \\ Parkinson & 0.27 & 0.27 & **0.25** & 0.28 & 0.27 & **0.26** & 0.37 & **0.35** & 0.42 \\ Bike sharing & 0.30 & 0.44 & **0.25** & 0.37 & 0.56 & **0.28** & **0.54** & 0.74 & 0.70 \\ Inflation & **0.05** & 0.07 & 0.07 & **0.06** & 0.08 & 0.10 & **0.10** & 0.13 & 0.13 \\ Obesity & **0.02** & 0.05 & 0.04 & 0.04 & 0.05 & **0.02** & 0.04 & 0.10 & **0.02** \\ Height & 0.02 & 0.03 & **0.02** & 0.01 & 0.03 & **0.01** & 0.03 & 0.04 & **0.01** \\ Death rate & **0.10** & 0.15 & 0.22 & **0.12** & 0.17 & 0.19 & **0.20** & 0.30 & 0.58 \\ Sarcos & **0.12** & 0.12 & 0.13 & **0.12** & 0.12 & 0.16 & 0.13 & 0.12 & **0.10** \\ \hline \end{tabular}
\end{table}
Table 4: Relative test errors for models trained on reduced datasets. Errors are normalized by the response standard deviation from Table 1.
Figure 3: Average test data error as function of the number of task parameters. Learned-context neural networks is given in blue circles and last-layer MTL networks in green diamonds. Task parameter dimensions are set to 1, 2, 4, 8, and 16. Errors are normalized using the learned-context error from Table 2. A dotted line at one marks the base performance.
optimization objective. We only study these properties as a consideration _after_ the model has been trained using conventional multi-task learning.

For these experiments, task are divided into two groups. The base group are the tasks that participate in training of the shared parameters, by minimizing Equation 12. The hold-out group are tasks that arrive after the shared parameters are fixed. Hold-out tasks are considered one at a time. The loss function for hold-out task \(j\) is

\[\min_{\beta_{j}}\ \frac{1}{s^{2}}\sum_{i=1}^{n_{j}}\left(y_{ij}-f(x_{ij};\beta_ {j},\alpha)\right)^{2}+\beta_{j}^{\top}D^{-1}\beta_{j}. \tag{13}\]

This is equivalent to the maximum a posteriori estimate of the task parameters. The task parameter prior is \(\beta_{j}\sim\mathcal{N}(0,D)\), where \(D\) is found from the distribution of task parameters in the base group. The log-likelihood term is scaled by \(s^{2}\), which is the test error variance for the base group.

A critical part of parameter estimation is the shape of the likelihood function. Using the Frequency dataset, we compare the hold-out task likelihood with the true data-generating process, which in this case is known. The result is given in Figure 4. Both parameter spaces display the same behavior for the different sets of data points. The modes of the likelihood develop predictably with additional data. This is evidence of a well-behaved task parameter space, and that the task parameters can capture properties that are fundamental to the problem, rather than just being a complex task encoding. For the synthetic Frequency case, we were free to generate a sufficiently large number of tasks and data points for this to be possible. The extent to which such a relationship between task parameters and underlying task properties can be identified is naturally dependent on the particular problem being studied.

To further investigate the hold-out task viability, an experiment is contucted on all datasets. For each dataset, tasks are randomly divided into three approximately equal groups. One group is selected as the hold-out group, while the other two constitute the base group. First, a learned-context neural network is trained using the base group tasks. Task parameters for the hold-out group is then estimated one task at the time, with the shared parameters frozen. The process is repeated using all groups as the hold-out group, and results are averaged over all iterations. The loss in Equation 13 is minimized using the LIPO optimizer (Malherbe and Vayatis, 2017).

The results are summarized in Figure 5 and Table 5. The outcome of this experiment will naturally vary with the datasets. In cases with few tasks or where tasks show a greater degree of uniqueness, it will be hard to identify sensible parameters for the hold-out tasks. This can be seen clearly for the Sarcos dataset, which only has seven tasks, and the tasks are quite different. For other cases, such as the Frequency dataset, there are many similar tasks, and the hold-out tasks are much more likely to resemble the base tasks. Inspecting Table 5, the hold-out task performance is never able to match the baseline from Table 2, but it is able to stay below a 20% error increase for half of the datasets. In real scenarios were new tasks appear, this leads to a trade-off between retraining
Figure 4: Task parameter likelihood functions for a hypothetical new task constructed for the Frequency dataset. The task has \(\omega_{j}=1.5\), and up to four data points. The left column is the likelihood of the true data-generating function, which is a sine wave parametrized by its frequency. The middle column is the likelihood of the learned-context neural network, with one task parameter. The right column is the underlying function in black, the data points in orange, and the learned-context network evaluated with task parameters sampled from the likelihood in blue. The model is taken from the experiment in Section 5.6. The rows show the likelihoods for a different number of data points.
the full model for a potential performance gain, or just estimating the new task parameters to save time and resources. As seen in Figure 5, it is beneficial to have more data, but the amount required for satisfying parameter estimation can vary greatly.

We note that when using fewer data points for the hold-out tasks, a lower task parameter dimension can be advantageous. This is observed in the Obesity and Height datasets, where all three task parameter dimensions yield similar performance using the full dataset, but the smaller task parameter spaces become increasingly favoured with fewer data points. Compare this to the results observed in Figure 3, where it was found that an increasing task parameter dimension was favourable when all tasks were trained simultaneously. The optimal choice of task-parameter dimension is then up to the specific use case.

## 6 Summary and discussion

The theoretical and empirical results show that the learned-context neural network is a favorable architecture for MTL problems with similar tasks, and where tasks may have few data points. Its ability to learn a small and well-behaved task parameter spaces can be particularly advantageous for such problems.

Theoretically, scalar task parameters were shown to be sufficient for a learned-context neural network to achieve universal approximation of all tasks (Section 4.3). The contextual inputs facilitates complex task adaptations even for simple constructed networks (Section 4.2). The ability to adapt to different tasks naturally increases with the size of the neural network. A potential downside to this flexibility is the possibility of over fitting to the individual tasks, which counters the desirable properties of multi-task learning. This puts emphasis on careful hyperparameter selection.

Experimentally it is seen that the ideal number of task parameters varies between problems, but the architecture can generally achieve a large degree of

\begin{table}
\begin{tabular}{l r r r} \hline Dataset & \(d_{\beta}=2\) & \(d_{\beta}=4\) & \(d_{\beta}=8\) \\ \hline Frequency & **1.07** & 1.56 & 1.4 \\ Sine and line & 1.28 & **1.26** & 1.72 \\ Schools & 1.06 & **1.04** & 1.05 \\ Parkinson & 1.11 & 1.05 & **1.01** \\ Bike sharing & 1.27 & **1.18** & 1.23 \\ Inflation & 2.08 & **1.74** & 2.9 \\ Obesity & 4.17 & **2.77** & 9.74 \\ Height & 2.18 & **1.67** & 2.4 \\ Death rate & **1.02** & 1.25 & 1.88 \\ Sarcos & **3.81** & 4.89 & 9.71 \\ \hline \end{tabular}
\end{table}
Table 5: Mean test error for hold-out tasks when trained on 100% of the training set. Errors are normalized using the learned-context network performance from Table 2.
Figure 5: Average test set errors for hold-out tasks. Errors are normalized using the learned-context network performance from Table 2. Hold-out tasks are trained on 1%, 5%, 10%, 25%, 50%, 75%, and 100% of the task data, and evaluated on the full test set. Models with 2, 4, and 8 task parameters are used.
task adaptation with only a few task parameters (Section 5.6). Increasing the task parameter dimension is observed to have a beneficial effect in cases where all tasks are trained simultaneously (Section 5.7). However, if the shared network model is to be used for new tasks, then a smaller parameter space may be preferable. The ideal task parameter dimension will likely have to be set by the practitioner in light of domain knowledge and the desired application of the model. The architecture facilitates task parameters that capture latent properties of the tasks (Section 5.7 and Appendix D), which can enable convenient workflows for updating and maintaining such models in practice.

Learned-context neural networks performed similarly to the benchmark architectures on the full datasets (Section 5.4). On the reduced datasets the learned-context neural network had less performance deterioration than the others (Section 5.5).

Training of the learned-context networks appears to be robust, and comparable to the other architectures discussed (Section 5.4). The construction used to prove Theorem 2 gives a motivation for initializing the task parameters to zero. Randomly initialized task parameters would have a higher chance to get stuck in local minima with "task-encoding" properties. Zero initializing, on the other hand, encourage similar task to follow the same trajectories during optimization, which enables a grouping of similar tasks. This likely promotes a well-behaved parameter space, and reduces the chance that multiple regions represents the same phenomena.

## 7 Conclusion

The learned-context neural network is a simple, but powerful, multi-task learning architecture. Its properties make it well suited to problems with moderate amounts of data, and in situations where a low-dimensional, well-behaved task parameter space is beneficial for application and analysis of the model. The task adaptation mechanism yields universal approximation capabilities with only a single task parameter. Empirical studies show that the ideal task parameter dimension varies with the domain and model application, but the number of required task parameters is generally lower than that of comparable architectures.

## Acknowledgements

This work was supported by Solution Seeker Inc. and The Research Council of Norway.
Example network

The following describes the neural network example used in Section 4.2. The input vector is \(z^{\top}=\begin{bmatrix}x&\beta\end{bmatrix}\), with both \(x\) and \(\beta\) scalar. The network has three layers, with weights

\[W_{1} =\begin{bmatrix}1&L_{1}\\ 1&L_{2}\\ 1&L_{3}\\ 1&L_{4}\end{bmatrix},b_{1}=-\frac{1}{2}\begin{bmatrix}0\\ 1\\ 2\\ 3\end{bmatrix},\] \[W_{2} =\begin{bmatrix}2&-4&0&0\\ 0&0&2&-4\end{bmatrix},b_{2}=0,\] \[W_{3} =\begin{bmatrix}1&1\end{bmatrix},b_{3}=0.\]

For simplicity, we ignore the residual connections in this case. Recall that we only consider the ReLU activation function. The first layer maps the inputs to a four dimensional hidden layer. Ignoring the task parameter, this mapping creates four identical unit slopes starting at 0.0, 0.5, 1.0, and 1.5. The second layer add pairs of the previous hidden layer together, creating two triangle shaped responses as the second hidden layer. These are added together in the third layer. The output is seen as the black lines in Figure 2.

### Derivation of dilation effect

We show that setting \(L=b_{1}\) in the network above creates a dilation effect. This was explored in Section 4.2. To simplify the analysis we take \(\beta\in(-1,1)\). Let \(z_{p}^{(2)}=x+L_{p}\beta+b_{1,p}\) be the \(p\)th element of the first hidden layer. Substituting \(L_{p}=b_{1,p}=-(p-1)/2\), we get

\[g(z_{p}^{(2)})=\begin{cases}x-\frac{1+\beta}{2}(p-1)&\text{ if }x\geq\frac{1+ \beta}{2}(p-1)\\ 0&\text{ otherwise.}\end{cases}\]

Continuing the notation for the second hidden layer, we get

\[g(z_{1}^{(3)}) =\begin{cases}0,&\text{ if }x<0,\\ 2x,&\text{ if }x\in\left[0,\frac{1+\beta}{2}\right),\\ -2x+2(1+\beta),&\text{ if }x\in\left[\frac{1+\beta}{2},1+\beta\right),\\ 0,&\text{ if }x\geq 1+\beta,\end{cases}\] \[g(z_{2}^{(3)}) =\begin{cases}0,&\text{ if }x<1+\beta,\\ 2x-2(1+\beta),&\text{ if }x\in\left[1+\beta,3\frac{1+\beta}{2}\right),\\ -2x+4(1+\beta),&\text{ if }x\in\left[3\frac{1+\beta}{2},2(1+\beta)\right),\\ 0,&\text{ if }x\geq 2(1+\beta),\end{cases}\]
The output is then given as \(y=g(z_{1}^{(3)})+g(z_{2}^{(3)})\), which is a piecewise linear function interpolating the points \((x,y)=(0,0),(0.5(1+\beta),1+\beta),(1+\beta,0),(1.5(1+\beta),1+\beta),(2(1+\beta ),0)\). This is equivalent with a dilation of both \(x\) and \(y\) with a factor \(1+\beta\).

## Appendix B Optimizer and learning rate schedule

All neural networks are implemented and trained with PyTorch (Paszke et al., 2019). They are trained on a single GPU. Optimization is done by stochastic gradient decent with momentum (Bottou et al., 2018) and a learning rate scheduler.

The learning rate scheduler has two stages. Starting at \(1^{-8}\), it is first increased linearly during a warm-up stage (Nakamura et al., 2021; Arpit et al., 2019) until it reaches peak learning rate.

The second stage is to train the model over several epochs until the loss converges(Chee and Toulis, 2018; Lang et al., 2019), at which point the learning rate is reduced by half. This is repeated until the learning rate is reduced back down below \(1^{-8}\), the maximum number of epochs is reached, or the loss is sufficiently close to zero. Loss convergence is determined by inspecting a window of the last epoch losses. A linear regression model with slope and intercept is fitted to the loss values. This is compared to a model with intercept only, using a likelihood ratio test (Wilks, 1938). Convergence is flagged if the test is above 0.51, which is an aggressive threshold. The test is implemented in Statsmodels (Seabold and Perktold, 2010). The new learning rate is kept for a minimum number of epochs equal to the window size.

The number of epochs and data batches vary with dataset size. For data sets with less that 100 000 data points we use 10 000 epochs of two batches, otherwise we use 1000 epochs of 20 batches. The warm-up stage is equal to 10% of the maximum allowed epochs, and the loss convergence is found with a window size equal to 1% of the epochs. Peak learning rate and momentum are treated as hyperparameters.

Neural network parameters are initialized by the He initializer (He et al., 2015). Task parameters are initialized to zero.

Linear mixed effect models are implemented and optimized with Statsmodels (Seabold and Perktold, 2010).

All data is transformed to approximately unit scale before training and evaluating models. Still, all figures and errors are given in the original units, unless stated otherwise.

## Appendix C Hyperparameters

All three neural network architectures require the number of residual blocks, hidden layer size, and network parameter regularization factor as hyperparameters. They also required the peak optimizer learning rate. Additionally, learned
context networks and last-layer MTL networks require the number of task parameters and task parameter regularization factor.

Hyperparameters are optimized using the training data in Table 1. The training data is split into two parts, using one part for training and one part for validation during the search. A final model is then trained on all training data using the hyperparameter configuration with the best validation error. Optimization is done by a variant of the LIPO solver (Malherbe and Vayatis, 2017) implemented in dlib (King, 2009). The search runs for 25 iterations in all cases.

The range of values explored in the searches varies by datasets. Details are summarized in Table 6. To limit the number of hyperparameters, momentum is fixed to 0.7 and the number of residual blocks is fixed to two for all datasets and architectures. These values are found as reasonable choices for most cases.

Due to a large number of experiments and hyperparameter searches, we observe that some searches arrive at an optimal learning rate that is too high to be used in training the final model. It can be due to randomness in weight initialization and batch samples that allowed a high learning rate to succeed during the trials, only to diverge during final training. To address this, we multiply the peak learning rate by 0.9 and retry the training if the loss diverges.

Full hyperparameter searches are conducted for the experiments in Section 5.4 and Section 5.5. For experiment in Section 5.6, the hyperparameters are copied from the final models in Section 5.4. For the experiment in Section 5.7, the hyperparameters are copied from the experiment using 50% training data in Section 5.5.

## Appendix D Qualitative study of task parameters

This section provides additional visualizations of the learned-context neural networks trained in Section 5. The intention is to give further insight into the qualitative behavior of the task parameters. We study the Inflation, Bike sharing, and Death rate datasets because they have a low dimensional input space that is convenient to visualize. The models with a scalar task parameter from Section 5.6 are used for all visualizations.

Figure 6 illustrates the Inflation data and the fitted model. The model

\begin{table}
\begin{tabular}{l l l} \hline Name & Min & Max \\ \hline Peak learning rate & \(10^{-4}\) & 1.5 \\ Hidden layer size & 50 & 500 \\ Shared parameter regularization & \(10^{-15}\) & \(10^{-5}\) \\ Task parameter regularization & \(10^{-15}\) & \(10^{-3}\) \\ Number of task parameters & 1 & min(25, \(m\)) \\ \hline \end{tabular}
\end{table}
Table 6: Summary of hyperparameter search space. The number of task parameters is limited to the number of tasks \(m\) in the dataset, up to a maximum of 25.
appears as a smoothed version of the data. The task parameters range from -0.2 to 0.2. Higher task parameter values seem to represent countries where the increase begins in the 1990s, and lower values represent countries where the increase is well underway in the 1970s. The transition between these two categories is highly non-linear. For a task parameter equal to -0.2, the curve is initially steep and flattens towards the end. The curve then gradually transitions to an almost linear trend for values around -0.1. For even higher values, there curve starts with a flat phase that eventually ramps up, with the start time increasing along with the task parameter.

The bike-sharing dataset relates the number of bike trips to weather, calendar features, and time of day, over two years. Each month is a task, yielding 24 tasks. Figure 7 illustrates the average number of bike trips and the task parameters for each month. There is an increasing number of trips over the two years, combined with a seasonal effect of more activity during the summer months. The increase in activity in the second year may be due to an increase in the capacity or popularity of the system. The task parameters nicely follow the average trips taken during peak hours. A possible interpretation is that the task parameters capture some latent state of the system, such as the number of bikes being distributed. We emphasize that while the parameters are visualized as a connected line, they are _not_ regularized to follow a time-dependent pattern. All parameters are assumed to be independent and centered around zero. However, a time-dependent regularization between the task parameters could be considered if the model were to be used in practice.

The death rate dataset studies the rate of death at different ages, as a

Figure 6: Inflation dataset. The left plot is the data points, where each line is the normalized inflation rate for a country (recall that each country is a task). The right plot is the fitted models evaluated at the data points. In both plots, the lines are colored by the value of the task parameter for that country.
function of country, year, and sex. Age is grouped in intervals of five years, and the maximum age included is 80. All ages above 80 are grouped together with a death rate of one, which has been removed from the dataset in this study. The data is given in five different years, 2000, 2005, 2010, 2015, and 2019. Each country and year is modeled as an individual task, making each country associated with five tasks. This relationship is _not known_ by the model. Figure 8 illustrates the task parameters, with six countries highlighted. The year 2010 is investigated in detail in Figure 9, which illustrates the data and fitted models for the highlighted countries in Figure 8. It seems that lower task parameter values relate to higher death rates in general. For most countries the task parameters seem to be incrementally increased with the years, indicating a decrease in death rates for ages below 80. For instance, task parameters for Ethiopia (ETH) show a steady increase over time. This correlates with the increase in life expectancy at birth observed over the same period (GBD 2019 Ethiopia Subnational-Level Disease Burden Initiative Collaborators, 2022).

Haiti (HTI) 2010 stands out in Figure 8. This is due to the devastating earthquake of January 2010, which had a large death toll (Cavallo et al., 2010). In Figure 9 we see that this leads to a curve shape that is quite unique, with an elevated death rate across all ages. In this case, the other tasks fails to supply the information missing from the training data, and the resulting model overshoots on the test data. For the other tasks, the gaps left by the test data are covered nicely by related tasks. For instance, the female death rate in Zambia (ZMB) is almost perfectly captured with only six data points available for training, of which only one represents ages 50 and above.

Figure 7: Bike sharing dataset. The top plot is the average number of trips at different times of day for each month, with the standard deviations given as transparent bands. This plot only includes data from workdays. The bottom plot is the identified task parameter for each month.
Figure 8: Death rate dataset. Visualization of task parameters. Each combination of country and year is an independent task in the model formulation, but tasks from the same country are connected by a line. Six countries are highlighted in colors.

Figure 9: Death rate dataset. Data and fitted models for six countries for the year 2010. The country and corresponding task parameters for this year are given in the titles. Data is given in black and colored markers, where black is the training data and colored is the test data. Circle markers are used for males and square markers for females. Fitted models are given in dashed lines for males and solid lines for females.
We emphasize that the task parameters are not forced into the relationships seen in Figure 6, Figure 7, and Figure 8. The continuous and interpretable task parameters are only motivated by regularization towards zero. The discovered structures are due to the information in the training data.
