Multi-Task Learning

WIP

• Multi task Learning deals with the goal of trying to jointly learn all the tasks (eg.walking, running, jumping etc., ) instead of learning them independently. The overaching goal is to not just learn these tasks but to do it in an efficient manner by trying to exploit the commonalities across the tasks. The shared dynamics underlying the tasks dictates the efficacy of multi task learning.

Why Multi Task Learning?

• Imagine if you are trying to learn how to draw, assuming that you already know how to write, you bring the knowledge of how to manipulate a pen and how to do certain kind of strokes to produce a pattern when learning to draw. On the other hand it would be very hard and sad if you had to learn from scratch how to do everything. Here you are trying to exploit the shared structure across tasks and use it to learn efficently when performing second task. In similar vein, we can try to train a system that can utilize this underlying structure to learn fast and efficiently. It also helps in the case of low resource problem settings when there isn’t enough data available for independent tasks, but could amount to enough data when tasks are combined. One thing to keep in mind is that, we are trying to learn the different tasks simultaneously, but when we relearn for a specific task after learning on macro set of tasks, it falls into the territory of meta learning.

Definitions

• To begin with, lets start with a more formal definition of what a task is and some other notations along with it. This is for the case of Supervised Learning. For multi task setting in Reinforcement Learning the definitions are listed here, after which the multi-task RL formulation is discussed.

• In a single task setting, we have dataset $\mathcal{D}=\{(x,y{)}_k\}$ with the objective given as ${\min }_{\theta }\mathcal{L}(\theta ,\mathcal{D})$ , where $\theta$ are the parameters of the model.

  • In multi task setting, we define a task as the one consisting of a dataset, given as dataset
  $${\mathcal{T}}_i\triangleq \{p_i(x),p_i(y\Vert x),{\mathcal{L}}_i)\}$$
  • Each Task ${\mathcal{T}}_i$ has a train ${\mathcal{D}}_i^{tr}$ and a test ${\mathcal{D}}_i^{test}$ dataset corresponing to it.

Naive Summation

The naive way to do multi task learning when there are multiple datasets is to create a unified dataset and metrics

$$\begin{array}{r}\mathcal{D}=\bigcup {\mathcal{D}}_i\\ \\ \mathcal{L}=\sum {\mathcal{L}}_i\end{array}$$

thus converting the multi task setting to single task and call it done.

• The problem with this approach is that we are assuming that all the tasks are independent of each other and are failing to exploit the shared substructure of these tasks by not creating better models that can capture this idea.

Mutli-Head Architecture

One of the most straight forward and simplest way to build multi-task architectures is to, maintain a common network for most of the layers from the beginning and then split off into multiple networks after a certain point.

Single Level Conditioning

To indicate task specificity we can use a variable $z_i$ as the descriptor of the task $i$, which then turns our Network $$\begin{gathered} f_{\theta }(y \\ |x) \end{gathered}$$ into $$\begin{gathered} f_{\theta }(y \\ |x,z_i) \end{gathered}$$. So now to predict the output $y$ we use both the input $x$ and task description given by $z_i$.

Minimal Shared Conditioning

To keep conditioning at minimum, we can only intorduce the task specific information at the very later stages of the network. One example might be to do Multiplicative Gating. This is bit similar to Multi-Head architecture. In this setting we train a desired number of neural networks independently of each other and the outputs $y_{\mathrm{1..}n}$ are aggregated with task based conditioning $z_i$. All the outputs $y_{\mathrm{1..}n}$ are aggregated into $y$ as

$$y=\sum _{j}1(z_i=j)y_j$$

where $1$ is an indicator variable that is valid when the one hot encoding provided by $z_i$ is lined up with $y_i$.

Shared Conditioning

Instead of introducing the task specific information at the final stages, we can also introduce it at an earlier stage in the network. This also has the added benefit of increasing parameter sharing as we also replace multiple networks with a single one.

In this setting, all the parameters before the introduction of $z_i$ are shared and the parameters followed by the conditioning aren't. This also brings us to the idea we introduced in the beginning that Multi Task Learning is about exploiting the shared sub-structure. We can envision that the shared part of the network before $z_i$ is learning the common factors and the part after $z_i$ is learning some task specific information.

A new Multi Task Objective

We can control the amount of information being shared across tasks by choosing to introduce task specific information at various stages of the network. This gives us an alternative view of Multi Task Learning.

We can now split the network parameters $\theta$ into shared parameters ${\theta }^{sh}$ and task-specific parameters ${\theta }^i$, a kind of lateralization. This gives us the objective

$$\underset{{\theta }^{sh},{\theta }^1,\dots ,{\theta }^T}{\min }\sum _{i=1}^T{\mathcal{L}}_i(\{{\theta }^{sh},{\theta }^i\},{\mathcal{D}}_i^{tr})$$

Types of Shared Conditioning

An extreme case of Shared Conditioning is to introduce the task specific information at the input This takes out the guesswork and the assumptions we might need to make in selecting the position of conditioning. So, Let’s use this type of conditioning to explore some common choices. We’ve already seen concatenation based conditioning where we extend one of the hidden states of the network with task specific information.

Another way to do conditioning is to sum the task specific information with one of the hidden layers. This can be extended with a linaer operation performed on task embedding and then adding this to the network. This is called conditional biasing. The output of the linear operation over the task embedding is considered as bias, as it is being added to the input.

A similar way to introduce conditioning is through scaling. Similar to additive conditioning, but the network layer or the input is scaled based on the task specific information. This is called conditional scaling or multiplicative conditioning. Mutliplicative conditioning cam be seen as more general version of multiplicative gating.

One of the intgersting things is that both conditional biasing and and concatenation based conditioning can be viewed as one and the same. If we take the concatenation of input $x$ and the task specific information $x_i$ and passed them through a linear layer, i.e., multiply them with a weight matrix $W$, we end up with a similar result.

One way to eliminate the choise of choosing between different types of task conditiong is to combine both conditional biasing and conditional scaling through an affine transformation.

Mutli level Conditioning

In multi level conditioning, we introduce task specific information and also share information across tasks at multiple levels and multiple times unlike the laterlized way we’ve seen in single level conditioning.

Cross-Stitching

Cross-Stitching takes two networks runnning in parallel and introduces Sross-Stitch units that enables parameter sharing and information exchange at certain points in the network architecture. The two networks are intended to solve two tasks independently after training.

Cross-Stitch units considers any of the hidden layers outputs in the Network as inputs and outputs a linear combination of those inputs, which are then fed into the next hidden layer.

There are many more complex architctures to consider, that can be built in a multi-task setting, but all these are task and domain dependent unlike the single level conditioning which are more broadly applicable. But the general Objective is basically

$$\underset{\theta }{\min }\mathcal{L}(\theta ,\mathcal{D})=\underset{\theta }{\min }\sum _{i=1}^{\mathcal{T}}{\mathcal{L}}_i(\theta ,{\mathcal{D}}_i)$$

The training process is :

• Sample mini-batch of task $\mathcal{B}\sim \{{\mathcal{T}}_i\}$
• Sample mini-batch datapoints for each task ${\mathcal{D}}_i^b\sim {\mathcal{D}}_i$
• Compute loss for each mini-batch

$$\hat{\mathcal{L}}(\theta ,\mathcal{B})=\sum _{{\mathcal{T}}_kϵ\mathcal{B}}{\mathcal{L}}_k(\theta ,{\mathcal{D}}_k^b)$$

• Backprop, compute gradient ${\nabla }_{\theta }\hat{\mathcal{L}}$ and update weights.

Challenges

One of the main challenges we see is when performance on one task hurts the performance on the other, called negative transfer. This doesn’t necessearly mean that there isn’t an shared strcuture across tasks; it is most likely an issue with the distribution of the data across datasets. Here, sometimes it is better to train models independently. The other issues involve optimization, where the gradients from one task may hurt the other task. Same learning rate across tasks might impede performance if the tasks may be learnt at different rates.

Another reason for negative transfer might be from limited representational capacity of the models, leading to underfit models. For multi-task learning the networks often need to much larger than the single task learning models, to actually be anywhere comparable. Multiplicative gating works best for these kind of issues, as the multi-task network is comprised of independent networks, by sharing less across tasks.

Instead of a binary decision of where and how much of the parameters are being shared, we can use a technique called soft-parameter sharing wher the parameters across the network are kind of constrained on each other. This is a bit similar to the Cross-Stitched networks but more general in that we don’t need to design specific architectures. We modify our multi-task obejctive to

$$\stackrel{\text{original multi-task objective}}{\overbrace{\underset{{\theta }^{sh},{\theta }^1,\dots ,{\theta }^T}{\min }\sum _{i=1}^T{\mathcal{L}}_i(\{{\theta }^{sh},{\theta }^i\},{\mathcal{D}}_i^{tr})}}+\underset{\text{soft parameter sharing}}{\underbrace{\sum _{t^{{}^{\prime }}=1}^T||{\theta }^t={\theta }^{t^{{}^{\prime }}}||}}$$

This new objective allows for a more fluid sharing of parameters across networks,but comes with additional set of hyperparameters to tune.

Another issue we see is the issue of overfitting. Here we observe overfitting for individual unshared networks and the overfit on their respective task.The most likely casue is when enough information is not being shared. The idea of soft-parameter sharing could alleviate some of the performance issues caused by overfitting.

# **Reinforcement Learning** > WIP

References

1. https://web.stanford.edu/class/cs330/
2. Dumoulin, et al., Feature-wise transformations
3. Ishan Misra, et al., Cross-stitch Networks for Multi-task Learning