Multi-Task Learning Notes (II): Acting on Gradients

Descent Direction

In fact, we answered the first question in "Optimization Algorithms from a Dynamics Perspective (III): A More Holistic View". Suppose the loss function is \mathcal{L} and the current parameter is \boldsymbol{\theta}. We want to design a parameter increment \Delta\boldsymbol{\theta} that makes the loss function smaller, i.e., \mathcal{L}(\boldsymbol{\theta}+\Delta\boldsymbol{\theta}) < \mathcal{L}(\boldsymbol{\theta}). To this end, we consider the first-order Taylor expansion: \mathcal{L}(\boldsymbol{\theta}+\Delta\boldsymbol{\theta})\approx \mathcal{L}(\boldsymbol{\theta}) + \langle \nabla_{\boldsymbol{\theta}}\mathcal{L}, \Delta\boldsymbol{\theta}\rangle \label{eq:approx-1} Assuming the accuracy of this approximation is sufficient, \mathcal{L}(\boldsymbol{\theta}+\Delta\boldsymbol{\theta}) < \mathcal{L}(\boldsymbol{\theta}) implies \langle \nabla_{\boldsymbol{\theta}}\mathcal{L}, \Delta\boldsymbol{\theta}\rangle < 0. This means the angle between the update amount and the gradient must be at least greater than 90 degrees. The most natural choice is: \Delta\boldsymbol{\theta} = -\eta \nabla_{\boldsymbol{\theta}}\mathcal{L} This is gradient descent, where the update is taken in the opposite direction of the gradient, and \eta > 0 is the learning rate.

No Exceptions

Returning to multi-task learning, if we assume each task is equally important, we can interpret this assumption as requiring \mathcal{L}_1, \mathcal{L}_2, \dots, \mathcal{L}_n to all decrease or remain unchanged at each update step. If, after the parameters reach \boldsymbol{\theta}^*, any further change leads to an increase in at least one \mathcal{L}_i, then \boldsymbol{\theta}^* is said to be a Pareto Optimal solution. Simply put, Pareto optimality means we cannot improve one task by sacrificing another; it means there is no "internal competition" (involution) between tasks.

Assuming the approximation [eq:approx-1] still holds, seeking Pareto optimality means we need to find \Delta\boldsymbol{\theta} such that: \left\{\begin{aligned} &\langle \nabla_{\boldsymbol{\theta}}\mathcal{L}_1, \Delta\boldsymbol{\theta}\rangle \leq 0\\ &\langle \nabla_{\boldsymbol{\theta}}\mathcal{L}_2, \Delta\boldsymbol{\theta}\rangle \leq 0\\ &\quad \vdots \\ &\langle \nabla_{\boldsymbol{\theta}}\mathcal{L}_n, \Delta\boldsymbol{\theta}\rangle \leq 0\\ \end{aligned}\right. Note that a trivial solution \Delta\boldsymbol{\theta}=\boldsymbol{0} exists, so the feasible region of the above system of inequalities is certainly non-empty. We are primarily concerned with whether non-zero solutions exist in the feasible region. If they do, we find one to use as the update direction; if not, we may have reached Pareto optimality (necessary but not sufficient), and we call this state a Pareto Stationary point.

Problem Transformation

First, we further transform the problem. Note that: \forall i, \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle \geq 0 \quad \Leftrightarrow \quad \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle \geq 0 \label{eq:q-0} Therefore, we only need to maximize the minimum \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle as much as possible to find the ideal \boldsymbol{u}. The problem becomes: \max_{\boldsymbol{u}}\min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle However, this is somewhat dangerous because if a non-zero \boldsymbol{u} exists such that \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle > 0, then letting the norm of \boldsymbol{u} tend to infinity would cause the maximum value to tend to infinity. For stability, we add a regularization term: \max_{\boldsymbol{u}}\min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle - \frac{1}{2}\Vert \boldsymbol{u}\Vert^2 \label{eq:q-1} In this way, a \boldsymbol{u} with an infinite norm cannot be the optimal solution. Note that substituting \boldsymbol{u}=0 yields \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle - \frac{1}{2}\Vert \boldsymbol{u}\Vert^2 = 0. Thus, if the optimal solution for \max is \boldsymbol{u}^*, it must satisfy: \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}^*\rangle - \frac{1}{2}\Vert \boldsymbol{u}^*\Vert^2 \geq 0 \quad \Leftrightarrow \quad \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}^*\rangle \geq \frac{1}{2}\Vert \boldsymbol{u}^*\Vert^2 \geq 0 So the solution to problem [eq:q-1] must satisfy condition [eq:q-0], and if it is a non-zero solution, its opposite direction must be a direction that decreases the loss of all tasks.

Smooth Approximation

Now we introduce the first solution for problem [eq:q-1]. It assumes the reader, like the author, is unfamiliar with solving min-max problems. We can replace the \min in the first step with a smooth approximation (refer to "Seeking a Smooth Maximum Function"): \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle \approx -\frac{1}{\lambda}\log\sum_i e^{-\lambda \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle} \quad (\text{for sufficiently large } \lambda) Thus, we can first solve: \max_{\boldsymbol{u}}-\frac{1}{\lambda}\log\sum_i e^{-\lambda \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle} - \frac{1}{2}\Vert \boldsymbol{u}\Vert^2 And then let \lambda \to \infty. This transforms the problem into an unconstrained maximization of a single function. Taking the gradient and setting it to zero gives: \frac{\sum\limits_i e^{-\lambda \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle} \boldsymbol{g}_i}{\sum\limits_i e^{-\lambda \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle}} = \boldsymbol{u} Assuming the differences between various \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle are larger than the order of \mathcal{O}(1/\lambda), then as \lambda \to \infty, the above equation is effectively: \boldsymbol{u} = \boldsymbol{g}_{\tau}, \quad \tau = \mathop{\text{argmin}}_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle However, if we iterate directly according to \boldsymbol{u}^{(k+1)} = \boldsymbol{g}_{\tau}, \tau = \mathop{\text{argmin}}\limits_i \langle \boldsymbol{g}_i, \boldsymbol{u}^{(k)}\rangle, it will likely oscillate. This is because it asks us to find the \boldsymbol{g}_i that minimizes \langle \boldsymbol{g}_i, \boldsymbol{u}^{(k)}\rangle and use it as \boldsymbol{u}^{(k+1)}. If \boldsymbol{u}^{(k+1)} = \boldsymbol{g}_{i^*}, then in the next step, the \boldsymbol{g}_i that minimizes \langle \boldsymbol{g}_i, \boldsymbol{u}^{(k+1)}\rangle = \langle \boldsymbol{g}_i, \boldsymbol{g}_{i^*}\rangle is very likely no longer \boldsymbol{g}_{i^*}; in fact, \boldsymbol{g}_{i^*} might be the largest one.

Intuitively, although the algorithm oscillates, it should oscillate around the optimal point \boldsymbol{u}^*. Therefore, if we average all results during the oscillation process, we should obtain the optimal point. This means the iteration format converging to the optimal point is: \boldsymbol{u}^{(k+1)} = \frac{k \boldsymbol{u}^{(k)} + \boldsymbol{g}_{\tau}}{k + 1}, \quad \tau = \mathop{\text{argmin}}_i \langle \boldsymbol{g}_i, \boldsymbol{u}^{(k)}\rangle \label{eq:sol-1} Note that since each added term is some \boldsymbol{g}_i, the final \boldsymbol{u}^* must be a weighted average of the \boldsymbol{g}_i’s. That is, there exist \alpha_1, \alpha_2, \dots, \alpha_n \geq 0 and \alpha_1 + \alpha_2 + \dots + \alpha_n = 1 such that: \boldsymbol{u}^* = \sum_i \alpha_i \boldsymbol{g}_i We can also interpret \alpha_1, \alpha_2, \dots, \alpha_n as the current optimal weight distribution for each \mathcal{L}_i.

Dual Problem

The advantage of the smooth approximation technique is that it is simple and intuitive, requiring little background in optimization algorithms. However, it is ultimately a "non-mainstream" approach with some lack of rigor (though the result is correct). Next, we introduce the "standard answer" based on the idea of duality.

First, define \mathbb{P}^n as the set of all n-dimensional discrete distributions: \mathbb{P}^n = \left\{(\alpha_1,\alpha_2,\dots,\alpha_n) \left| \alpha_1,\alpha_2,\dots,\alpha_n \geq 0, \sum_i \alpha_i = 1 \right.\right\} Then it is easy to verify: \min_i \langle \boldsymbol{g}_i, \boldsymbol{u}\rangle = \min_{\alpha\in\mathbb{P}^n} \langle \tilde{\boldsymbol{g}}(\alpha), \boldsymbol{u} \rangle, \quad \tilde{\boldsymbol{g}}(\alpha) = \sum_i \alpha_i \boldsymbol{g}_i Thus, problem [eq:q-1] is equivalent to: \max_{\boldsymbol{u}}\min_{\alpha\in\mathbb{P}^n} \langle \tilde{\boldsymbol{g}}(\alpha), \boldsymbol{u} \rangle - \frac{1}{2}\Vert \boldsymbol{u}\Vert^2 \label{eq:q-2} The above function is concave with respect to \boldsymbol{u} and convex with respect to \alpha, and the feasible regions for \boldsymbol{u} and \alpha are both convex sets. Therefore, according to von Neumann’s Minimax Theorem, the \min and \max in [eq:q-2] can be swapped: \min_{\alpha\in\mathbb{P}^n}\max_{\boldsymbol{u}} \langle \tilde{\boldsymbol{g}}(\alpha), \boldsymbol{u} \rangle - \frac{1}{2}\Vert \boldsymbol{u}\Vert^2 = \min_{\alpha\in\mathbb{P}^n} \frac{1}{2}\Vert \tilde{\boldsymbol{g}}(\alpha) \Vert^2 \label{eq:q-3} The right side of the equation follows because the \max part is just an unconstrained maximum problem for a quadratic function, which yields \boldsymbol{u}^* = \tilde{\boldsymbol{g}}(\alpha). Finally, only the \min remains, and the problem becomes finding a weighted average of \boldsymbol{g}_1, \boldsymbol{g}_2, \dots, \boldsymbol{g}_n that minimizes its norm.

When n=2, the solution is relatively simple, equivalent to constructing the altitude of a triangle, as shown below:

When n > 2, we can use the Frank-Wolfe algorithm to transform it into multiple n=2 cases for iteration. The Frank-Wolfe algorithm can be understood as a constrained gradient descent algorithm suitable for cases where the parameter’s feasible region is a convex set. Explaining it in detail would take too much space, so I won’t elaborate here; readers are encouraged to find materials to study it. Briefly, the Frank-Wolfe algorithm first linearizes the objective to find the next update direction e_{\tau}, where \tau = \mathop{\text{argmin}}\limits_i \langle \boldsymbol{g}_i, \tilde{\boldsymbol{g}}(\alpha)\rangle and e_{\tau} is a one-hot vector with 1 at position \tau. Then it performs an interpolation search between \alpha and e_{\tau} to find the optimal result. The iteration process is: \left\{\begin{aligned} &\tau = \mathop{\text{argmin}}_i \langle \boldsymbol{g}_i, \tilde{\boldsymbol{g}}(\alpha^{(k)})\rangle\\ &\gamma = \mathop{\text{argmin}}_{\gamma} \Vert \tilde{\boldsymbol{g}}((1-\gamma)\alpha^{(k)} + \gamma e_{\tau}) \Vert^2 = \mathop{\text{argmin}}_{\gamma} \Vert (1-\gamma)\tilde{\boldsymbol{g}}(\alpha^{(k)}) + \gamma \boldsymbol{g}_{\tau} \Vert^2\\ &\alpha^{(k+1)} = (1-\gamma)\alpha^{(k)} + \gamma e_{\tau} \end{aligned}\right. The solution for \gamma is exactly the n=2 special case, which can be solved using the algorithm in the screenshot above. If \gamma is not obtained through search but fixed at 1/(k+1), the result is equivalent to [eq:sol-1], which is a simplified version of the Frank-Wolfe algorithm. In other words, the result obtained through smooth approximation is equivalent to the result of the simplified Frank-Wolfe algorithm.

De-constraint

In fact, for the solution of problem [eq:q-3], we could theoretically solve it directly via gradient descent by removing constraints. For example, we can set parameters \beta_1, \beta_2, \dots, \beta_n \in \mathbb{R} and: \alpha_i = \frac{e^{\beta_i}}{Z}, \quad Z = \sum_i e^{\beta_i} Then it can be transformed into: \min_{\beta} \frac{1}{2Z^2} \left\Vert \sum_i e^{\beta_i} \boldsymbol{g}_i \right\Vert^2 This is an unconstrained optimization problem that can be solved with conventional gradient descent. However, for some reason, the author has not seen this approach used (perhaps to avoid tuning the learning rate?).

Gradient Inner Product

As seen, in both schemes, the key step is \mathop{\text{argmin}}\limits_i \langle \boldsymbol{g}_i, \tilde{\boldsymbol{g}}(\alpha)\rangle, which means we must traverse gradients to calculate inner products. However, in deep learning scenarios, model parameters are often very large, so the gradient is a very high-dimensional vector. Calculating an inner product at every iteration is computationally expensive. We can utilize the expansion: \langle \boldsymbol{g}_i, \tilde{\boldsymbol{g}}(\alpha)\rangle = \left\langle \boldsymbol{g}_i, \sum_j \alpha_j \boldsymbol{g}_j \right\rangle = \sum_j \alpha_j \langle \boldsymbol{g}_i, \boldsymbol{g}_j \rangle In each iteration, only \alpha changes. Therefore, \langle \boldsymbol{g}_i, \boldsymbol{g}_j \rangle only needs to be calculated once and stored for each training step, avoiding repeated high-dimensional vector inner product calculations.

Shared Encoding

However, when the model is large enough, it is difficult to calculate the gradient for each task separately and then perform iterative calculations. If we assume that the various tasks share the same encoder, we can further simplify the algorithm with an approximation.

Specifically, assume the batch size is b, and the encoder output for the j-th sample is \boldsymbol{h}_j. From the chain rule, we know: \boldsymbol{g}_i = \nabla_{\boldsymbol{\theta}}\mathcal{L}_i = \sum_j (\nabla_{\boldsymbol{h}_j}\mathcal{L}_i)(\nabla_{\boldsymbol{\theta}}\boldsymbol{h}_j) = \underbrace{\big(\nabla_{\boldsymbol{h}_1}\mathcal{L}_i , \dots , \nabla_{\boldsymbol{h}_b}\mathcal{L}_i\big)}_{\nabla_{\boldsymbol{H}}\mathcal{L}_i}\underbrace{\begin{pmatrix}\nabla_{\boldsymbol{\theta}}\boldsymbol{h}_1 \\ \vdots \\ \nabla_{\boldsymbol{\theta}}\boldsymbol{h}_b\end{pmatrix}}_{\nabla_{\boldsymbol{\theta}}\boldsymbol{H}} Letting \boldsymbol{H} = (\boldsymbol{h}_1, \dots, \boldsymbol{h}_b), we get \boldsymbol{g}_i = (\nabla_{\boldsymbol{H}}\mathcal{L}_i) (\nabla_{\boldsymbol{\theta}}\boldsymbol{H}). Using the matrix norm inequality: \left\Vert\sum_i \alpha_i \boldsymbol{g}_i\right\Vert^2 = \left\Vert\sum_i \alpha_i (\nabla_{\boldsymbol{H}}\mathcal{L}_i) (\nabla_{\boldsymbol{\theta}}\boldsymbol{H})\right\Vert^2 \leq \left\Vert\sum_i \alpha_i \nabla_{\boldsymbol{H}}\mathcal{L}_i\right\Vert^2 \big\Vert \nabla_{\boldsymbol{\theta}}\boldsymbol{H}\big\Vert^2 Naturally, if we minimize \left\Vert\sum\limits_i\alpha_i \nabla_{\boldsymbol{H}}\mathcal{L}_i\right\Vert^2, the computational cost will be significantly reduced because this only requires the gradient of the final output encoding vector, not the gradients of all parameters. The above formula tells us that minimizing \left\Vert\sum\limits_i\alpha_i \nabla_{\boldsymbol{H}}\mathcal{L}_i\right\Vert^2 is actually minimizing the upper bound of equation [eq:q-3]. Like many problems that are difficult to optimize directly, we hope that minimizing the upper bound will yield similar results.

However, while this upper bound is more efficient, it has limitations. It generally only applies to multi-task learning where each sample has multiple types of labels. it is not suitable for scenarios where training data for different tasks have no intersection (i.e., each task is labeled on different samples, and a single sample has only one type of label). In the latter case, the various \nabla_{\boldsymbol{H}}\mathcal{L}_i are mutually orthogonal, meaning there is no interaction between tasks, and the upper bound becomes too loose to be meaningful.

Error in Proof

The "standard answer" mentioned earlier and the optimization of the upper bound for shared encoders both come from the paper "Multi-Task Learning as Multi-Objective Optimization". The original paper then attempts to prove that when \nabla_{\boldsymbol{\theta}}\boldsymbol{H} is full rank, optimizing the upper bound can also find a Pareto stationary point. Unfortunately, the proof in the original paper is incorrect.

The proof is located in Appendix A of the original paper, where it uses an incorrect conclusion:

If \boldsymbol{M} is a symmetric positive definite matrix, then \boldsymbol{x}^{\top}\boldsymbol{y}\geq 0 if and only if \boldsymbol{x}^{\top}\boldsymbol{M}\boldsymbol{y}\geq 0.

It is easy to provide a counterexample to prove this conclusion is wrong. For instance, let \boldsymbol{x}=\begin{pmatrix}1 \\ -2\end{pmatrix}, \boldsymbol{y}=\begin{pmatrix}1 \\ 1\end{pmatrix}, \boldsymbol{M}=\begin{pmatrix}3 & 0\\ 0 & 1\end{pmatrix}. Here \boldsymbol{x}^{\top}\boldsymbol{y} < 0 but \boldsymbol{x}^{\top}\boldsymbol{M}\boldsymbol{y} > 0.

After reflection, the author believes the proof in the original paper is difficult to fix, meaning the original paper’s conjecture does not hold. In other words, even if \nabla_{\boldsymbol{\theta}}\boldsymbol{H} is full rank, the update direction derived from the upper bound may not necessarily be a direction that prevents all task losses from rising, and thus may not find a Pareto stationary point. As for why the experimental results of the upper bound optimization in the original paper are good, it can only be said that the parameter space of deep learning models is so large that there is plenty of "room to maneuver," allowing the upper bound approximation to achieve decent results.