Gradient Surgery for Multi-Task Learning

1 INTRODUCTION . While deep learning and deep reinforcement learning ( RL ) have shown considerable promise in enabling systems to perform complex tasks , the data requirements of current methods make it difficult to learn a breadth of capabilities particularly when all tasks are learned individually from scratch . A natural approach to such multi-task learning problems is to train a single network on all tasks jointly , with the aim of discovering shared structure across the tasks in a way that achieves greater efficiency and performance than solving the tasks individually . However , learning multiple tasks all at once results in a difficult optimization problem , sometimes leading to worse overall performance and data efficiency compared to learning tasks individually ( Parisotto et al. , 2015 ; Rusu et al. , 2016a ) . These optimization challenges are so prevalent that multiple multi-task RL algorithms have considered using independent training as a subroutine of the algorithm before distilling the independent models into a multi-tasking model ( Levine et al. , 2016 ; Parisotto et al. , 2015 ; Rusu et al. , 2016a ; Ghosh et al. , 2017 ; Teh et al. , 2017 ) , producing a multi-task model but losing out on the efficiency gains over independent training . If we could tackle the optimization challenges of multi-task learning effectively , we may be able to actually realize the hypothesized benefits of multi-task learning without the cost in final performance . While there has been a significant amount of research in multi-task learning ( Caruana , 1997 ; Ruder , 2017 ) , the optimization challenges are not well understood . Prior work has described varying learning speeds of different tasks ( Chen et al. , 2017 ) and plateaus in the optimization landscape ( Schaul et al. , 2019 ) as potential causes , while a range of other works have focused on the model architecture ( Misra et al. , 2016b ; Liu et al. , 2018 ) . In this work , we instead hypothesize that the central optimization issue in multi-task learning arises from gradients from different tasks conflicting with one another . In particular , we define two gradients to be conflicting if they point away from one another ( i.e. , have a negative cosine similarity ) . As a concrete example , consider the 2D optimization landscapes of two task objectives shown in Figure 1 . The optimization landscape of each task consists of a deep valley , as has been characterized of neural network optimization landscapes in the past ( Goodfellow et al. , 2014 ) . When considering the combined optimization landscape for multiple tasks , SGD produces gradients that struggle to efficiently find the optimum . This occurs due to a gradient thrashing phenomenon , where the gradient of one task destabilizes optimization in the valley . We can observe this in Figure 1 ( d ) when the optimization reaches the deep valley of task 1 , but is prevented from traversing the valley to an optimum . In Section 6.2 , we find experimentally that this thrashing phenomenon also occurs in a neural network multi-task learning problem . The core contribution of this work is a method for mitigating gradient interference by altering the gradients directly , i.e . by performing “ gradient surgery ” . If two gradients are conflicting , we alter the gradients by projecting each onto the normal plane of the other , preventing the interfering components of the gradient from being applied to the network . We refer to this particular form of gradient surgery as projecting conflicting gradients ( PCGrad ) . PCGrad is model-agnostic , requiring only a single modification to the application of gradients . Hence , it is easy to apply to a range of problem settings , including multi-task supervised learning and multi-task reinforcement learning , and can also be readily combined with other multi-task learning approaches , such as those that modify the architecture . We evaluate PCGrad on multi-task CIFAR classification , multi-objective scene understanding , a challenging multi-task RL domain , and goal-conditioned RL . Across the board , we find PCGrad leads to significant improvements in terms of data efficiency , optimization speed , and final performance compared to prior approaches . Further , on multi-task supervised learning tasks , PCGrad can be successfully combined with prior state-of-the-art methods for multi-task learning for even greater performance . 2 PRELIMINARIES . The goal of multi-task learning is to find parameters θ of a model fθ that achieve high average performance across all the training tasks drawn from a distribution of tasks p ( T ) . More formally , we aim to solve the problem : min θ ETi∼p ( T ) [ Li ( fθ ) ] , where Li is a loss function for the i-th task Ti that we want to minimize . To obtain a model that solves a specific task from the task distribution p ( T ) , we define a task-conditioned model fθ ( y|x , zi ) , with input x , output y , and encoding zi for task Ti , which could be provided as a one-hot vector or in any other form . 3 MULTI-TASK LEARNING VIA GRADIENT SURGERY . While the multi-task problem can in principle be solved by simply applying a standard single-task algorithm with a suitable task identifier provided to the model or a simple multi-head or multi-output model , a number of prior works ( Parisotto et al. , 2015 ; Rusu et al. , 2016a ; Sener & Koltun , 2018 ) have found this learning problem to be difficult , especially in the reinforcement learning setting . We hypothesize that one of the main challenges of multi-task learning can be characterized by conflicting and thrashing gradients , and find that this can significantly impede learning progress , especially when combined with iterative data collection . We identify possible causes for this problem and propose a simple and general approach to mitigate it . 3.1 THRASHING GRADIENTS IN MULTI-TASK OPTIMIZATION LANDSCAPES . We hypothesize that a key optimization issue in multi-task learning arises when gradients from multiple tasks are in conflict with one another , i.e . when gradients point away from one another . More specifically , we hypothesize that such conflict may lead to gradient thrashing . Concretely , gradient thrashing refers to the phenomenon where a large gradient for one task changes the parameter vectors in a way that substantially decreases performance on another task . Since worse performance typically leads to larger gradients , this results in alternating gradient directions , where , at the next iteration , the second task will have large gradients that dominate and reduce performance on the former task . This issue can be particularly pronounced for neural network optimization , since neural network loss landscapes are known to resemble long narrow valleys ( Goodfellow et al. , 2014 ) , where the gradient perpendicular to the direction of the valley will be small . We aim to study this hypothesis through two toy examples . First , consider the two-dimensional optimization landscape illustrated in Fig . 1a , where the landscape for each task objective corresponds to a deep and curved valley ( Fig . 1b and 1c ) . The optima of this multi-task objective correspond to where the two valleys meet . More details on the optimization landscape are in Appendix B . We observe that the gradient thrashing hypothesis is consistent with what we observe when running Adam ( Kingma & Ba , 2014 ) on this landscape in Fig . 1d , where we observe that Adam does not traverse one valley towards the other , preventing it from reaching an optimum . We also aim to detect if a similar phenomenon occurs in multi-task learning with a neural network with thousands of parameters on a toy regression problem . To measure the extent of gradient thrashing , we plot the cosine similarity between the gradients of two tasks throughout the beginning of learning in Fig . 4 ( left ) . We indeed observe a significant level of gradient thrashing at every iteration , where the cosine similarity varies between −0.75 and 0.75 at a very high frequency . Motivated by these observations , we develop an algorithm that aims to alleviate the optimization challenges caused by gradient thrashing by preventing such gradient conflict between tasks . 3.2 PCGRAD : PROJECTING CONFLICTING GRADIENTS . We aim to prevent gradient thrashing by directly altering the gradients themselves , i.e . through “ gradient surgery. ” To be maximally effective and maximally applicable , we must perform surgery in a way that still allows for positive interactions between the task gradients and does not introduce any assumptions on the form of the model . We start by first detecting whether two gradients are in conflict , by measuring whether they point away from one another . More concretely , we characterize two tasks as conflicting for the current parameter setting if they yield a negative cosine similarity between their respective gradients . The goal of PCGrad is to modify the gradients for each task so as to minimize negative conflict with other task gradients , which will in turn mitigate gradient thrashing . To deconflict gradients during optimization , PCGrad adopts a simple procedure : if the gradients between two tasks are in conflict , i.e . their cosine similarity is negative , we project the gradient from one task onto the normal plane of the gradient of the other task . This amounts to removing the conflicting component of the gradient for the task , thereby reducing the amount of destructive gradient interference between tasks . A pictorial description of this idea is shown in Fig . 2 . Suppose the gradient for task Ti is gi , and the gradient for task Tj is gj . PCGrad proceeds as follows : ( 1 ) First , it determines whether gi conflicts with gj by computing the cosine similarity between vectors gi and gj , where negative values indicate conflicting gradients . ( 2 ) If the cosine similarity is negative , we replace gi by its projection onto the normal plane of gj : gi = gi − gi·gj‖gj‖2gj . If the gradients are not in conflict , i.e . cosine similarity is non-negative , the original gradient gi remains unaltered . ( 3 ) PCGrad repeats this process across all of the other tasks sampled in random order from the current Algorithm 1 PCGrad Update Rule Require : Current model parameters θ 1 : Sample mini-batch of tasks B = { Tk } ∼ p ( T ) 2 : for Ti ∼ B in sequence do 3 : Compute gradient gi of Ti as gi = ∇θLi ( fθ ) 4 : for Tj uniformly∼ B in random order do 5 : Compute gradient gj of task Tj as gj = ∇θLj ( fθ ) 6 : Compute cosine similarity between gi as gj as cos ( φij ) = gi·gj ‖gi‖‖gj‖ . 7 : if cos ( φij ) < 0 then 8 : Set gi = gi − gi·gj‖gj‖2 gj // Subtract the projection of gi onto gj 9 : end if 10 : end for 11 : Store gproji = gi 12 : end for 13 : return update ∆θ = ∑ i g proj i batch Tj ∀ j 6= i , resulting in the gradient gproji that is applied for task Ti . We perform the same procedure for all tasks in the batch to obtain their respective gradients . The full update procedure is described in Algorithm 1 and a discussion on using a random task order is included in Appendix D. This procedure , while simple to implement , ensures that the gradients that we apply for each task per batch interfere minimally with the other tasks in the batch , mitigating the thrashing gradient problem , producing a variant on standard first-order gradient descent in the multi-objective setting . In practice , the PCGrad gradient surgery method can be combined with any gradient-based optimizer , including commonly used methods such as SGD with momentum and Adam ( Kingma & Ba , 2014 ) , by simply passing the computed update to the respective optimizer instead of the original gradient . Our experimental results verify the hypothesis that this procedure reduces the problem of thrashing gradients , and find that , as a result , learning progress is substantially improved . Finally , we analyze the convergence of this procedure in Theorem 1 in the two-task setting , to ensure that the procedure is sensible under the standard assumptions in optimization . Theorem 1 . Consider two task loss functions L1 : Rn → R and L2 : Rn → R which are convex and differentiable . For all θ ∈ Rn , let L ( θ ) = L1 ( θ ) + L2 ( θ ) , i.e . L is a multi-task objective . Let φ be the angle between ∇L1 ( θ ) and ∇L2 ( θ ) . Suppose L is differentiable and that its gradient is Lipschitz continuous with constant L > 0 , i.e . we have ||∇L ( θ1 ) −∇L ( θ2 ) ||2 ≤ L||θ1 − θ2||2 for any θ1 , θ2 . Then , the PCGrad update rule with step size t ≤ 1L will converge to either ( 1 ) a location in the optimization landscape where cos ( φ ) = −1 or ( 2 ) the optimal value L ( θ∗ ) . Proof . See Appendix A. Theorem 1 states that application of the PCGrad update in the two-task setting with a convex and Lipschitz multi-task loss function L leads to convergence to either the minimizer of L or a potentially sub-optimal objective value . A sub-optimal solution occurs when the cosine similarity between the gradients of the two tasks is −1 , i.e . the gradients directly conflict , leading to zero gradient after applying PCGrad . However , in practice , since we are using SGD , which is a noisy estimate of the true batch gradients , the cosine similarity between the gradients of two tasks in a minibatch is unlikely to be −1 , thus avoiding this scenario .

Summary: The paper presents a method to boost multi-task learning performance by editing gradient to remove conflicts between tasks. The main idea is to use cosine similarity to 1) determine if two task gradients conflict and 2) to project one conflicting gradient to the normal plane of the other, thereby removing the conflict at the expense of disturbing the other gradient to some extent. Experiments are presented for classification and other computer vision tasks along with reinforcement learning problems.