Multi-Task Learning Musings (I): In the Name of Loss

多任务学习漫谈（一）：以损失之名 苏剑林 · 2022-01-18 · Machine Learning deep learning loss function multi-task

There are many ways to improve model performance, and Multi-Task Learning is one of them. Put simply, multi-task learning hopes to jointly train multiple related tasks, so that the different tasks can complement and reinforce one another, ultimately achieving better single-task performance (accuracy, robustness, etc.). However, multi-task learning is not as simple as just piling all the tasks together and letting it work — how to balance the training of each task so that every task gets a beneficial boost is still a topic worth studying.

Recently, through a series of coincidences, I've had occasion to try out some multi-task learning myself, and along the way I've studied the related literature. Here I'd like to pick out some of the results to share and discuss with everyone.

Weighted Sum

From the perspective of the loss function, multi-task learning means having multiple loss functions $\mathcal{L}_1,\mathcal{L}_2,\cdots,\mathcal{L}_n$, which in general share the vast majority of parameters while each having a small number of independent parameters, and our goal is to make every loss function as small as possible. To this end, we introduce weights $\alpha_1,\alpha_2,\cdots,\alpha_n\geq 0$, converting it into single-task learning with the following weighted-sum loss function:

\begin{equation}\mathcal{L} = \sum_{i=1}^n \alpha_i \mathcal{L}_i\label{eq:w-loss}\end{equation}

From this perspective, the main difficulty of multi-task learning is how to determine each $\alpha_i$. more

Initial State

In principle, in the absence of any task prior or bias, the most natural choice is to treat every task equally, i.e., $a_i=1/n$. In reality, however, tasks can differ enormously — for instance, mixing classification tasks with different numbers of classes, mixing classification with regression, mixing classification with generation, and so on. From a physical point of view, each loss function has a different dimension and magnitude, so simply adding them together is meaningless.

If we regard each loss function as a physical quantity with a different dimension, then, following the idea of "non-dimensionalization," we can use the reciprocal of the loss function's initial value as the weight, i.e.,

\begin{equation}\mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{init})}}\label{eq:init}\end{equation}

where $\mathcal{L}_i^{(\text{init})}$ denotes the initial loss value of task $i$. This expression is "homogeneous" with respect to each $\mathcal{L}_i$, so one obvious advantage is scale invariance: if we multiply the loss of task $i$ by a constant, the result does not change. Moreover, since each loss is divided by its own initial value, larger losses get shrunk and smaller losses get amplified, so that the losses become roughly balanced.

So how do we estimate $\mathcal{L}_i^{(\text{init})}$? The most direct method, of course, is to estimate it using a few batches of data. Alternatively, we can derive a theoretical value based on certain assumptions. For instance, under mainstream initialization schemes, we can assume that the initial model's output (before the activation function) is a zero vector; if a softmax is applied, this corresponds to a uniform distribution. So for an "$K$-way classification + cross-entropy" problem, the initial loss is $\log K$; for a "regression + L2 loss" problem, we can estimate the initial loss using a zero vector, giving $\mathbb{E}_{y\sim \mathcal{D}}[\Vert y-0\Vert^2] = \mathbb{E}_{y\sim \mathcal{D}}[\Vert y\Vert^2]$, where $\mathcal{D}$ is the full set of labels in the training set.

Prior State

One problem with using the initial loss is that the initial state doesn't necessarily reflect the current difficulty of learning the task very well. A better scheme would be to replace the "initial state" with a "prior state":

\begin{equation}\mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{prior})}}\label{eq:prior}\end{equation}

For example, if in a $K$-way classification problem the frequency of each class is $[p_1,p_2,\dots,p_K]$ (the prior distribution), then although the predicted distribution at initialization is uniform, we can reasonably assume that the model could easily learn to predict $[p_1,p_2,\dots,p_K]$ for every sample based on the prior alone, at which point the model's loss equals the entropy

\begin{equation}\mathcal{L}_i^{(\text{prior})}=\mathcal{H} = -\sum_{i=1}^K p_i\log p_i\end{equation}

In a sense, the "prior distribution" captures the essence of "initial" better than the "initial distribution" does — it reflects the idea that "even if the model learns nothing at all, it still knows to guess according to the prior distribution." So this loss value better represents the initial difficulty of the current task, and thus using $\mathcal{L}_i^{(\text{prior})}$ in place of $\mathcal{L}_i^{(\text{init})}$ should be more reasonable. Similarly, for the "regression + L2 loss" problem, the prior result should be the expectation of all the labels, $\mu = \mathbb{E}_{y\sim \mathcal{D}}[y]$, so we replace $\mathcal{L}_i^{(\text{init})}=\mathbb{E}_{y\sim \mathcal{D}}[\Vert y\Vert^2]$ with $\mathcal{L}_i^{(\text{prior})}=\mathbb{E}_{y\sim \mathcal{D}}[\Vert y-\mu\Vert^2]$, which is expected to yield more sensible results.

Dynamic Adjustment

Whether we use the initial-state formula $\eqref{eq:init}$ or the prior-state formula $\eqref{eq:prior}$, the task weights, once determined, remain fixed, and the method of determining them does not depend on the learning process. However, although we can roughly gauge task difficulty from information such as the prior distribution, how difficult a task actually is can only be known through the process of learning it. So a more sensible scheme should dynamically adjust the weights as training progresses.

Real-Time State

Looking back over the previous sections, the core idea behind both formula $\eqref{eq:init}$ and formula $\eqref{eq:prior}$ is to use the reciprocal of the loss value as the task weight. So why not simply use the reciprocal of the "real-time" loss value to dynamically adjust the weights? That is,

\begin{equation}\mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{sg})}}\label{eq:sg}\end{equation}

Here $\mathcal{L}_i^{(\text{sg})}$ is shorthand for $\text{stop_gradient}(\mathcal{L}_i)$. In this scheme, the loss function for every task is adjusted to be identically 1, so it is consistent in both dimension and magnitude. Because of the $\text{stop_gradient}$ operator, even though the loss is always 1, the gradient is not identically 0:

\begin{equation}\nabla_{\theta}\left(\frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{sg})}}\right) = \frac{\nabla_{\theta}\mathcal{L}_i}{\mathcal{L}_i^{(\text{sg})}} = \frac{\nabla_{\theta}\mathcal{L}_i}{\mathcal{L}_i}\label{eq:sg-grad}\end{equation}

In simple terms, once a function is wrapped by the $\text{stop_gradient}$ operator, it becomes a new function whose value is identical to the original function's, but whose derivative is forced to be 0. The net effect is that the dynamic weight $1/\mathcal{L}_i$ scales the gradient in real time. Many "folk experiments" have shown that formula $\eqref{eq:sg}$ indeed serves as quite a decent baseline in most cases.

Equivalent Gradient

We can look at this scheme from another angle. From formula $\eqref{eq:sg-grad}$ we obtain

\begin{equation}\nabla_{\theta}\left(\frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{sg})}}\right) = \frac{\nabla_{\theta}\mathcal{L}_i}{\mathcal{L}_i} = \nabla_{\theta} \log \mathcal{L}_i\end{equation}

So, in terms of the gradient, formula $\eqref{eq:sg}$ is no different from $\mathcal{L} = \sum\limits_{i=1}^n \log \mathcal{L}_i$, and furthermore we have

\begin{equation}\mathcal{L} = \sum_{i=1}^n \log \mathcal{L}_i = n\log \sqrt[n]{\prod_{i=1}^n\mathcal{L}_i}\end{equation}

Since $\log$ is monotonically increasing, formula $\eqref{eq:sg}$ agrees in gradient direction with the following:

\begin{equation}\mathcal{L} = \sqrt[n]{\prod_{i=1}^n\mathcal{L}_i}\end{equation}

Generalized Mean

Clearly, the expression above is exactly the "geometric mean" of $\mathcal{L}_1,\mathcal{L}_2,\cdots,\mathcal{L}_n$, and if we stipulate that $a_i$ is identically equal to $1/n$, then the original formula $\eqref{eq:w-loss}$ is the "arithmetic mean" of $\mathcal{L}_1,\mathcal{L}_2,\cdots,\mathcal{L}_n$. In other words, we find that this whole chain of derivations conceals a transition from the arithmetic mean to the geometric mean, which suggests that we might consider a "generalized mean":

\begin{equation}\mathcal{L}(\gamma) = \sqrt[\gamma]{\frac{1}{n}\sum_{i=1}^n\mathcal{L}_i^{\gamma}}\end{equation}

That is, we raise each loss function to the power of $\gamma$, average them, and then take the $\gamma$-th root; here $\gamma$ can be any real number, with the arithmetic mean corresponding to $\gamma=1$ and the geometric mean corresponding to $\gamma=0$ (taken as a limit). It can be shown that $\mathcal{L}(\gamma)$ is a monotonically increasing function of $\gamma$, and moreover

\begin{equation}\min(\mathcal{L}_1,\cdots,\mathcal{L}_n)=\lim_{\gamma\to-\infty} \mathcal{L}(\gamma) \leq\cdots\leq \mathcal{L}(\gamma) \leq\cdots\leq \lim_{\gamma\to+\infty} \mathcal{L}(\gamma)=\max(\mathcal{L}_1,\cdots,\mathcal{L}_n)\end{equation}

This means that as $\gamma$ increases, the model pays more and more attention to the largest loss, while as it decreases, it pays more attention to the smallest loss. Thus, although there is still a hyperparameter $\gamma$ to tune, compared with the original formula $\eqref{eq:w-loss}$, the number of hyperparameters has been reduced from $n$ down to just 1, simplifying the tuning process.

Translation Invariance

Let's revisit formulas $\eqref{eq:init}$, $\eqref{eq:prior}$, and $\eqref{eq:sg}$ — all of them adjust the weights by dividing each task's loss by some state of itself, thereby achieving scale invariance. However, even though all of them possess scale invariance, they lose the more fundamental property of "translation invariance." That is to say, if we add a constant to every loss, the gradient directions of $\eqref{eq:init}$, $\eqref{eq:prior}$, and $\eqref{eq:sg}$ may change — which is not good news for optimization, since in principle a constant carries no meaningful information, and the optimization result should not change because of it.

The Ideal Target

On the one hand, we use the reciprocal of some state of the loss function as the weight for the current task, but the derivative of a loss function does not have translation invariance. On the other hand, a loss function can be understood as the distance between the current model and the target state, and gradient descent is essentially a search for points where the gradient is 0 — so the norm of the gradient can serve a similar purpose. Thus, we can replace the loss function with the norm of the gradient, turning formula $\eqref{eq:sg}$ into

\begin{equation}\mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\Vert\nabla_{\theta}\mathcal{L}_i\Vert^{(\text{sg})}}\label{eq:grad}\end{equation}

A clear difference from the loss function is that the gradient norm obviously has translation invariance, and both numerator and denominator are still homogeneous with respect to $\mathcal{L}_i$, so the expression above also retains scale invariance. Hence, this is an ideal target that possesses both translation and scale invariance simultaneously.

Gradient Normalization

Taking the gradient of formula $\eqref{eq:grad}$, we get

\begin{equation}\nabla_{\theta}\mathcal{L} = \sum_{i=1}^n \frac{\nabla_{\theta}\mathcal{L}_i}{\Vert\nabla_{\theta}\mathcal{L}_i\Vert}\label{eq:grad-norm}\end{equation}

We can see that formula $\eqref{eq:grad}$ is essentially normalizing the gradient of each task's loss and then summing the normalized gradients. This also suggests an implementation scheme: each task can be trained in turn, one task at a time per step, and the normalized gradients of each task can be accumulated before performing the update. This avoids the trouble of having to compute gradients already at the point of defining the loss function.

Regarding gradient normalization, the related work I was able to find is GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks, which is essentially a hybrid of formulas $\eqref{eq:init}$ and $\eqref{eq:grad-norm}$ and also incorporates the idea of rescaling the gradient norm — but it requires an additional optimization step to determine the task weights, which personally I find rather cumbersome and redundant.

Summary

From the perspective of the loss function, the key issue in multi-task learning is how to adjust the weight of each task to balance their respective losses. This post introduced some reference approaches from the angles of scale invariance and translation invariance, and supplemented these with the concept of the "generalized mean," which converts the problem of adjusting multiple task weights into the problem of tuning a single parameter, simplifying the tuning process.