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

Initial State

In principle, without any task priors or biases, the most natural choice is to treat each task equally, i.e., \alpha_i = 1/n. However, in reality, tasks may differ significantly. For example, mixing classification tasks with different numbers of categories, mixing classification and regression tasks, or mixing classification and generative tasks. From a physical perspective, the dimensions and magnitudes of each loss function are different, so direct addition is meaningless.

If we view each loss function as a physical quantity with different dimensions, then starting from the idea of "non-dimensionalization," we can use the reciprocal of the initial value of the loss function as the weight, namely: \mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{init})}} \label{eq:init} where \mathcal{L}_i^{(\text{init})} represents the initial loss value of task i. This formula is "homogeneous" with respect to each \mathcal{L}_i, so one obvious advantage is scale invariance. That is, if the loss of task i is multiplied by a constant, the result remains unchanged. Furthermore, since each loss is divided by its own initial value, larger losses are scaled down and smaller losses are scaled up, allowing each loss to be roughly balanced.

So, how do we estimate \mathcal{L}_i^{(\text{init})}? The most direct way 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, we can assume the initial model output (before the activation function) is a zero vector. If softmax is applied, it results in a uniform distribution. Thus, for a "K-class classification + cross-entropy" problem, its initial loss is \log K. For a "regression + L2 loss" problem, the initial loss can be estimated using a zero vector, i.e., \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 set of all labels in the training set.

Prior State

One problem with using the initial loss is that the initial state may not necessarily reflect the learning difficulty of the current task. A better approach would be to change the "initial state" to a "prior state": \mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{prior})}} \label{eq:prior} For example, if the frequencies of each class in a K-class classification are [p_1, p_2, \dots, p_K] (prior distribution), then although the predicted distribution at the initial state is uniform, we can reasonably assume the model can easily learn to predict [p_1, p_2, \dots, p_K] for every sample. In this case, the model’s loss is the entropy: \mathcal{L}_i^{(\text{prior})} = \mathcal{H} = -\sum_{i=1}^K p_i \log p_i In a sense, the "prior distribution" reflects the essence of "initial" better than the "initial distribution." It represents the idea that "even if the model learns nothing, it knows to output results randomly according to the prior distribution." Therefore, the loss value at this point better represents the initial difficulty of the current task. Thus, replacing \mathcal{L}_i^{(\text{init})} with \mathcal{L}_i^{(\text{prior})} should be more reasonable. Similarly, for the "regression + L2 loss" problem, its prior result should be the expectation of all labels \mu = \mathbb{E}_{y\sim \mathcal{D}}[y], so we use \mathcal{L}_i^{(\text{prior})} = \mathbb{E}_{y\sim \mathcal{D}}[\Vert y-\mu\Vert^2] to replace \mathcal{L}_i^{(\text{init})} = \mathbb{E}_{y\sim \mathcal{D}}[\Vert y\Vert^2], which is expected to yield more reasonable results.

Real-time State

Looking back at the previous sections, the core idea of formulas [eq:init] and [eq:prior] is to use the reciprocal of the loss value as the task weight. Can we simply use the "real-time" reciprocal of the loss value to achieve dynamic weight adjustment? That is: \mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\mathcal{L}_i^{(\text{sg})}} \label{eq:sg} Here \mathcal{L}_i^{(\text{sg})} is shorthand for \text{stop\_gradient}(\mathcal{L}_i). In this scheme, the loss function of each task is adjusted to be constantly 1, so they are consistent in both dimension and magnitude. Due to the presence of the \text{stop\_gradient} operator, although the loss is constant at 1, the gradient is not constantly 0: \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} Simply put, once a function is wrapped by the \text{stop\_gradient} operator, it becomes a new function whose value is identical to the original function, but its derivative is forced to 0. Thus, the final result is a real-time adjustment of the gradient ratio using the dynamic weight 1/\mathcal{L}_i. Many "informal experiments" indicate that formula [eq:sg] can indeed serve as a quite good baseline in most cases.

Equivalent Gradient

We can look at this scheme from another perspective. From formula [eq:sg-grad], we get: \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 Therefore, in terms of gradients, formula [eq:sg] has no substantial difference from \mathcal{L} = \sum_{i=1}^n \log \mathcal{L}_i. Furthermore, we have: \mathcal{L} = \sum_{i=1}^n \log \mathcal{L}_i = n \log \sqrt[n]{\prod_{i=1}^n \mathcal{L}_i} Since \log is monotonically increasing, formula [eq:sg] is consistent in gradient direction with the following: \mathcal{L} = \sqrt[n]{\prod_{i=1}^n \mathcal{L}_i}

Generalized Mean

Obviously, the above formula is the "geometric mean" of \mathcal{L}_1, \mathcal{L}_2, \dots, \mathcal{L}_n. If we fix \alpha_i to 1/n, then the original formula [eq:w-loss] is the "arithmetic mean" of \mathcal{L}_1, \mathcal{L}_2, \dots, \mathcal{L}_n. In other words, we find that this series of derivations actually hides a transition from the arithmetic mean to the geometric mean. This inspires us to perhaps consider the "generalized mean": \mathcal{L}(\gamma) = \sqrt[\gamma]{\frac{1}{n}\sum_{i=1}^n \mathcal{L}_i^{\gamma}} This involves taking the \gamma-th power of each loss function, averaging them, and then taking the \gamma-th root. Here \gamma can be any real number. The arithmetic mean corresponds to \gamma=1, and the geometric mean corresponds to \gamma=0 (taking the limit). It can be proven that \mathcal{L}(\gamma) is a monotonically increasing function of \gamma, and: \min(\mathcal{L}_1, \dots, \mathcal{L}_n) = \lim_{\gamma\to-\infty} \mathcal{L}(\gamma) \leq \dots \leq \mathcal{L}(\gamma) \leq \dots \leq \lim_{\gamma\to+\infty} \mathcal{L}(\gamma) = \max(\mathcal{L}_1, \dots, \mathcal{L}_n) This means that as \gamma increases, the model becomes more concerned with the maximum value among the losses; conversely, it becomes more concerned with the minimum value. In this way, although there is still a hyperparameter \gamma to adjust, compared to the original formula [eq:w-loss], the number of hyperparameters has been reduced from n to just 1, simplifying the tuning process.

Ideal Goal

On one hand, we use the reciprocal of the loss function (or some state of it) as the weight for the current task, but the derivative of the loss function does not possess translation invariance. On the other hand, the loss function can be understood as the distance between the current model and the target state, and gradient descent is essentially looking for points where the gradient is 0. Therefore, the magnitude of the gradient can actually play a similar role. Thus, we can replace the loss function with the magnitude of the gradient, turning formula [eq:sg] into: \mathcal{L} = \sum_{i=1}^n \frac{\mathcal{L}_i}{\Vert\nabla_{\theta}\mathcal{L}_i\Vert^{(\text{sg})}} \label{eq:grad} A clear difference from the loss function is that the gradient magnitude obviously possesses translation invariance. Furthermore, the numerator and denominator are still homogeneous with respect to \mathcal{L}_i, so the above formula also retains scale invariance. Therefore, this is an ideal goal that can simultaneously possess translation and scale invariance.

Gradient Normalization

Taking the gradient of formula [eq:grad], we get: \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} As we can see, formula [eq:grad] essentially normalizes the gradient of each task’s loss and then accumulates the gradients. It also suggests an implementation scheme: we can train each task sequentially, training only one task at a time, and then accumulate the normalized gradients of each task before updating. This eliminates the trouble of having to calculate gradients while defining the loss function.

Regarding gradient normalization, a related work I found is "GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks". It is essentially a mixture of formula [eq:init] and formula [eq:grad-norm], and it also includes the idea of rescaling gradient magnitudes, but it requires additional optimization to determine task weights, which I personally find somewhat cumbersome and redundant.