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 proposes a method to reduce conflicting gradients during training of multi-task learning paradigms via gradient modification. Rather than simply take an average of the per-task gradients, the authors propose to learn a model-wide gradient direction $d$ to update the shared parameters by solving a dual objective to find per-task loss weights. Simply, the proposed method (CAGrad) finds a direction which maximizes the worst local improvement among all tasks, and then applies this gradient update to the shared parameters.