Conflict-Averse Gradient Descent for Multi-task learning

1 Introduction . Multi-task learning ( MTL ) refers to learning a single model that can tackle multiple different tasks [ 11 , 28 , 44 , 38 ] . By sharing parameters across tasks , MTL methods learn more efficiently with an overall smaller model size compared to learning with separate models [ 38 , 40 , 25 ] . Moreover , it has been shown that MTL could in principle improve the quality of the learned representation and therefore benefit individual tasks [ 35 , 43 , 34 ] . For example , an early MTL result by [ 2 ] demonstrated that training a neural network to recognize doors could be improved by simultaneously training it to recognize doorknobs . However , learning multiple tasks simultaneously can be a challenging optimization problem because it involves multiple objectives [ 38 ] . The most popular MTL objective in practice is the average loss over all tasks . Even when this average loss is exactly the true objective ( as opposed to only caring about a single task as in the door/doorknob example ) , directly optimizing the average loss could lead to undesirable performance , e.g . the optimizer struggles to make progress so the learning performance significantly deteriorates . A known cause of this phenomenon is the conflicting gradients [ 41 ] : gradients from different tasks 1 ) may have varying scales with the largest gradient dominating the update , and 2 ) may point in different directions so that directly optimizing the average loss can be quite detrimental to a specific task ’ s performance . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To address this problem , previous work either adaptively re-weights the objectives of each task based on heuristics [ 3 , 15 ] or seeks a better update vector [ 30 , 41 ] by manipulating the task gradients . However , existing work often lacks convergence guarantees or only provably converges to any point on the Pareto set of the objectives . This means the final convergence point of these methods may largely depend on the initial model parameters . As a result , it is possible that these methods over-optimize one objective while overlooking the others ( See Fig . 1 ) . Motivated by the limitation of current methods , we introduce Conflict-Averse Gradient descent ( CAGrad ) , which reduces the conflict among gradients and still provably converges to a minimum of the average loss . The idea of CAGrad is simple : it looks for an update vector that maximizes the worst local improvement of any objective in a neighborhood of the average gradient . In this way , CAGrad automatically balances different objectives and smoothly converges to an optimal point of the average loss . Specifically , we show that vanilla gradient descent ( GD ) and the multiple gradient descent algorithm ( MGDA ) are special cases of CAGrad ( See Sec . 3.1 ) . We demonstrate that CAGrad can improve over prior state-of-the-art gradient manipulation methods on a series of challenging multi-task supervised , semi-supervised , and reinforcement learning problems . 2 Background . In this section , we first introduce the problem setup of multi-task learning ( MTL ) . Then we analyze the optimization challenge of MTL and discuss the limitation of prior gradient manipulation methods . 2.1 Multi-task Learning and its Challenge . In multi-task learning ( MTL ) , we are given K ≥ 2 different tasks , each of which is associated with a loss function Li ( θ ) for a shared set of parameters θ . The goal is to find an optimal θ ∈ Rm that achieves low losses across all tasks . In practice , a standard objective for MTL is minimizing the average loss over all tasks : θ∗ = argmin θ∈Rm { L0 ( θ ) ≜ 1 K K∑ i=1 Li ( θ ) } . ( 1 ) Unfortunately , directly optimizing ( 1 ) using gradient descent may significantly compromise the optimization of individual losses in practice . A major source of this phenomenon is known as the conflicting gradients [ 41 ] . Optimization Challenge : Conflicting Gradients Denote by gi = ∇Li ( θ ) the gradient of task i , and g0 = ∇L0 ( θ ) = 1K ∑K i gi the averaged gradient . With learning rate α ∈ R+ , θ ← θ − αg0 is the steepest descent update that appears to be the most natural update to follow when optimizing ( 1 ) . However , g0 may conflict with individual gradients , i.e . ∃ i , ⟨gi , g0⟩ < 0 . When this conflict is large , following g0 will decrease the performance on task i . As observed by [ 41 ] and illustrated in Fig . 1 , when θ is near a steep “ valley '' , where a specific task ’ s gradient dominates the update , manipulating the direction and magnitude of g0 often leads to better optimization . 2.2 Prior Attempts and Convergence Issues . Several methods have been proposed to manipulate the task gradients to form a new update vector and have shown improved performance on MTL . Sener et al . apply the multiple-gradient descent algorithm ( MGDA ) [ 6 ] for MTL , which directly optimizes towards the Pareto set [ 30 ] . Chen et al . dynamically re-weight each Li using a pre-defined heuristic [ 3 ] . More recently , PCGrad identifies conflicting gradients as the motivation behind manipulating the gradients and projects each task gradient to the normal plane of others to reduce the conflict [ 41 ] . While all these methods have shown success at improving the learning performance of MTL , they manipulate the gradient without respecting the original objective ( 1 ) . Therefore , these methods could in principle converge to any point in the Pareto set ( See Fig . 1 and Sec . 3.2 ) . We provide the detailed algorithms of MGDA and PCGrad in Appendix A.1 and A.2 , and a visualization of the update vector by each method in Fig . 2 . 3 Method . We introduce our main algorithm , Conflict-Averse Gradient descent in Sec . 3.1 , and then show theoretical analysis in Sec . 3.2 . 3.1 Conflict-Averse Gradient Descent . Assume we update θ by θ′ ← θ − αd , where α is a step size and d an update vector . We want to choose d to decrease not only the average loss L0 , but also every individual loss . To do so , we consider the minimum decrease rate across the losses , R ( θ , d ) = max i∈ [ K ] { 1 α ( Li ( θ − αd ) − Li ( θ ) ) } ≈ − min i∈ [ K ] ⟨gi , d⟩ , ( 2 ) where we use the first-order Taylor approximation assuming α is small . If R ( θ , d ) < 0 , it means that all losses are decreased with the update given a sufficiently small α . Therefore , R ( θ , d ) can be regarded as a measurement of conflict among objectives . With the above measurement , our algorithm finds an update vector that minimizes such conflict to mitigate the optimization challenge while still converging to an optimum of the main objective L0 ( θ ) . To this end , we introduce Conflict-Averse Gradient descent ( CAGrad ) , which on each optimization step determines the update d by solving the following optimization problem : max d∈Rm min i∈ [ K ] ⟨gi , d⟩ s.t . ∥d− g0∥ ≤ c ∥g0∥ , ( 3 ) Here , c ∈ [ 0 , 1 ) is a pre-specified hyper-parameter that controls the convergence rate ( See Sec . 3.2 ) . The optimization problem ( 3 ) looks for the best update vector within a local ball centered at the averaged gradient g0 , which also minimizes the conflict in losses measured by ( 2 ) . Since we focus on MTL and choose the average loss as the main objective , g0 is the average gradient . However , CAGrad also applies when g0 is the gradient of some other user-specified objective . We leave exploring this possibility as a future direction . Dual Objective The optimization problem ( 3 ) involves decision variable d that has the same dimension as the number of parameters in θ , which could be millions for a deep neural network . It is not practical to directly solve for d on every optimization step . However , the dual problem of Eq . ( 3 ) , as we will derive in the following , only involves solving for a decision variable w ∈ RK , which can be efficiently found using standard optimization libraries [ 7 ] . Specifically , first note that mini⟨gi , d⟩ = minw∈W⟨ ∑ i wigi , d⟩ , where w = ( w1 , . . . , wK ) ∈ RK and W denotes the probability simplex , i.e . W = { w : ∑ i wi = 1 and wi ≥ 0 } . Denote gw = ∑ i wigi and ϕ = c2 ∥g0∥2 . The Lagrangian of the objective in Eq . ( 3 ) is max d∈Rm min λ≥0 , w∈W g⊤wd− λ ( ∥g0 − d∥ 2 − ϕ ) /2 . Since the objective for d is concave with linear constraints , by switching the min and max , we reach the dual form without changing the solution by Slater ’ s condition : min λ≥0 , w∈W max d∈Rm g⊤wd− λ ∥g0 − d∥ 2 /2 + λϕ/2 . Algorithm 1 Conflict-averse Gradient Descent ( CAGrad ) for Multi-task Learning Input : Initial model parameter vector θ0 , differentiable loss functions { Li } Ki=1 , a constant c ∈ [ 0 , 1 ) and learning rate α ∈ R+ . repeat At the t-th optimization step , define g0 = 1K ∑K i=1∇Li ( θt−1 ) and ϕ = c2 ∥g0∥ 2 . Solve min w∈W F ( w ) : = g⊤wg0 + √ ϕ ∥gw∥ , where gw = 1 K K∑ i=1 wi∇Li ( θt−1 ) . Update θt = θt−1 − α ( g0 + ϕ1/2 ∥gw∥gw ) . until convergence We end up with the following optimization problem w.r.t . w after several steps of calculus , w∗ = argmin w∈W g⊤wg0 + √ ϕ ∥gw∥ , where the optimal λ∗ = ∥gw∗∥ /ϕ1/2 and the optimal update d∗ = g0 + gw∗/λ∗ . The detailed derivation is provided in Appendix A.3 and the entire CAGrad algorithm is summarized in Alg . 1 . The dimension of w equals to the number of objectives K , which usually ranges from 2 to tens and is much smaller than the number of parameters in a neural network . Therefore , in practice , we solve the dual objective to perform the update of CAGrad . Remark In Alg . 1 , when c = 0 , CAGrad recovers the typical gradient descent with d = g0 . On the other hand , when c→∞ , then minimizing F ( w ) is equivalent to minw ∥gw∥ . This coincides with the multiple gradient descent algorithm ( MGDA ) [ 6 ] , which uses the minimum norm vector in the convex hull of the individual gradients as the update direction ( see Fig . 2 ; second column ) . MGDA is a gradient-based multi-objective optimization designed to converge to an arbitrary point on the Pareto set , that is , it leaves all the points on the Pareto set as fixed points ( and hence can not control which specific point it will converge to ) . It is different from our method which targets to minimize L0 while using gradient conflict to regularize the optimization trajectory . As we will analyze in the following section , to guarantee that CAGrad converges to an optimum of L0 ( θ ) , we have to ensure 0 ≤ c < 1 .

Summary: This work presents a first-order multi-task learning method which uses, in place of the average task gradient, a vector which maximises the minimum local single-task improvement in a neighborhood of the average task gradient.  In multi-task learning, the goal is to find the optimal point achieving low losses across all tasks. However, gradient descent methods commonly used for single tasks in computer vision and reinforcement learning achieve suboptimal performance when applied on the average task loss due to conflicting gradients (single tasks gradients pointing in opposite directions). Recently, several methods have proposed modification to the update direction that alleviate this problem. However, these methods are either heuristics or can converge to any point in the pareto optimal set of the task losses (e.g. MGDA [5]), which might not be a minimum of the average task loss.   The proposed method, called Conflict-Averse Gradient descent (CAGrad), introduces an additional hyperparameter $c \geq 0$ which allows to interpolate between gradient descent ($c = 0$) and MGDA ($c \to \infty$). Furthermore, CAGrad provably converges to a stationary point of the average task loss with the same rate as gradient descent (when $c < 1$) or to a pareto stationary point (when $c \geq 0$). At each step of the method, the update direction is obtained by solving  an optimization problem in $K$ dimensions, where $K$ is the number of tasks, but has an overhead comparable to state of the art multi-task learning methods of the same kind (see last plot of Figure 5). Experiments on several computer vision and reinforcement learning multi-task benchmarks, where the method is used to optimize deep network parameters, show that the method achieves state of the art performance in all the settings.