summaries: Authors propose a multitask learning approach to reduce gradient conflict by minimizing harm to the worst-performing task given any gradient update. Proofs are provided for convergence and pareto optimality of the resultant weight configuration. Experiments are run on a variety of settings, including computer vision and multitask reinforcement learning, with good results throughout.  Multi-task learning (MTL) is an efficient way of learning multiple tasks efficiently by training all tasks together, amortizing their cost by sharing some parameters across tasks. In order to perform this, a weighted sum of the task losses is usually optimized, which can lead to poor final results due to sub-optimal optimization, an effect produced by the competition between tasks for the shared resources. Specifically, when computing the gradient updates for the shared resources, gradients of different tasks can conflict with each other, cancelling each other out (an effect known as *conflicting gradients*).  This work introduces CAGrad, an algorithm which minimizes the average loss function (the network's loss), while maximizing the worst local improvement across all tasks, provably converging to a local minimum of the average loss. This algorithm results in a generalization of the usual gradient descent, as well as the MGDA algorithm. Empirical results show the effectiveness of the method. 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.  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.  The paper proposes a new solution to multi-task learning  that should improve in cases with gradients conflict of a task loss and the average loss gradient. The solution optimizes for an update vector that maximizes the worst task loss relative improvement. Experiments show the superiority of the proposed solution over similar prior methods. 
abstractText: The goal of multi-task learning is to enable more efficient learning than single task learning by sharing model structures for a diverse set of tasks. A standard multi-task learning objective is to minimize the average loss across all tasks. While straightforward, using this objective often results in much worse final performance for each task than learning them independently. A major challenge in optimizing a multi-task model is the conflicting gradients, where gradients of different task objectives are not well aligned so that following the average gradient direction can be detrimental to specific tasks’ performance. Previous work has proposed several heuristics to manipulate the task gradients for mitigating this problem. But most of them lack convergence guarantee and/or could converge to any Pareto-stationary point. In this paper, we introduce Conflict-Averse Gradient descent (CAGrad) which minimizes the average loss function, while leveraging the worst local improvement of individual tasks to regularize the algorithm trajectory. CAGrad balances the objectives automatically and still provably converges to a minimum over the average loss. It includes the regular gradient descent (GD) and the multiple gradient descent algorithm (MGDA) in the multi-objective optimization (MOO) literature as special cases. On a series of challenging multi-task supervised learning and reinforcement learning tasks, CAGrad achieves improved performance over prior state-of-the-art multi-objective gradient manipulation methods. Code is available at https://github.com/Cranial-XIX/CAGrad.
references.authors: Vijay Badrinarayanan Alex Kendall Roberto Cipolla
references.title: Segnet: A deep convolutional encoder-decoder architecture for image segmentation
references.venue: IEEE transactions on pattern analysis and machine intelligence,
references.authors: Rich Caruana
references.title: Multitask learning
references.venue: Machine learning,
references.authors: Zhao Chen Vijay Badrinarayanan Chen-Yu Lee Andrew Rabinovich
references.title: Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks
references.venue: In International Conference on Machine Learning,
references.authors: Zhao Chen Jiquan Ngiam Yanping Huang Thang Luong Henrik Kretzschmar Yuning Chai Dragomir Anguelov
references.title: Just pick a sign: Optimizing deep multitask models with gradient sign dropout
references.authors: Li Deng
references.title: The mnist database of handwritten digit images for machine learning research
references.venue: IEEE Signal Processing Magazine,
references.authors: Jean-Antoine Désidéri
references.title: Multiple-gradient descent algorithm (mgda) for multiobjective optimization
references.venue: Comptes Rendus Mathematique,
references.authors: Steven Diamond Stephen Boyd
references.title: CVXPY: A Python-embedded modeling language for convex optimization
references.venue: Journal of Machine Learning Research,
references.authors: Dibya Ghosh Avi Singh Aravind Rajeswaran Vikash Kumar Sergey Levine
references.title: Divide-andconquer reinforcement learning
references.venue: arXiv preprint arXiv:1711.09874,
references.authors: Michelle Guo Albert Haque De-An Huang Serena Yeung Li Fei-Fei
references.title: Dynamic task prioritization for multitask learning
references.venue: In Proceedings of the European Conference on Computer Vision (ECCV),
references.authors: Tuomas Haarnoja Aurick Zhou Pieter Abbeel Sergey Levine
references.title: Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor
references.venue: In International Conference on Machine Learning,
references.authors: Kazuma Hashimoto Caiming Xiong Yoshimasa Tsuruoka Richard Socher
references.title: A joint manytask model: Growing a neural network for multiple nlp tasks
references.venue: arXiv preprint arXiv:1611.01587,
references.authors: Geoffrey Hinton Oriol Vinyals Jeff Dean
references.title: Distilling the knowledge in a neural network
references.venue: arXiv preprint arXiv:1503.02531,
references.authors: Harold M Hochman James D Rodgers
references.title: Pareto optimal redistribution
references.venue: The American economic review,
references.authors: Adrián Javaloy Isabel Valera
references.title: Rotograd: Dynamic gradient homogenization for multi-task learning
references.venue: arXiv preprint arXiv:2103.02631,
references.authors: Alex Kendall Yarin Gal Roberto Cipolla
references.title: Multi-task learning using uncertainty to weigh losses for scene geometry and semantics
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Diederik P Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.venue: arXiv preprint arXiv:1412.6980,
references.authors: Alex Krizhevsky Geoffrey Hinton
references.title: Learning multiple layers of features from tiny images
references.authors: Yann LeCun Léon Bottou Yoshua Bengio Patrick Haffner
references.title: Gradient-based learning applied to document recognition
references.venue: Proceedings of the IEEE,
references.authors: Liyang Liu Yi Li Zhanghui Kuang Jing-Hao Xue Yimin Chen Wenming Yang Qingmin Liao Wayne Zhang
references.title: Towards impartial multi-task learning
references.venue: In International Conference on Learning Representations,
references.authors: Shikun Liu Edward Johns Andrew J Davison
references.title: End-to-end multi-task learning with attention
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
references.authors: Xingchao Liu Xing Han Na Zhang Qiang Liu
references.title: Certified monotonic neural networks
references.venue: arXiv preprint arXiv:2011.10219,
references.authors: Debabrata Mahapatra Vaibhav Rajan
references.title: Multi-task learning with user preferences: Gradient descent with controlled ascent in pareto optimization
references.venue: In International Conference on Machine Learning,
references.authors: Kevis-Kokitsi Maninis Ilija Radosavovic Iasonas Kokkinos
references.title: Attentive single-tasking of multiple tasks
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
references.authors: Ishan Misra Abhinav Shrivastava Abhinav Gupta Martial Hebert
references.title: Cross-stitch networks for multi-task learning
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Emilio Parisotto Jimmy Lei Ba Ruslan Salakhutdinov
references.title: Actor-mimic: Deep multitask and transfer reinforcement learning
references.venue: arXiv preprint arXiv:1511.06342,
references.authors: Clemens Rosenbaum Tim Klinger Matthew Riemer
references.title: Routing networks: Adaptive selection of non-linear functions for multi-task learning
references.venue: arXiv preprint arXiv:1711.01239,
references.authors: Sebastian Ruder
references.title: An overview of multi-task learning in deep neural networks
references.venue: arXiv preprint arXiv:1706.05098,
references.authors: Andrei A Rusu Sergio Gomez Colmenarejo Caglar Gulcehre Guillaume Desjardins James Kirkpatrick Razvan Pascanu Volodymyr Mnih Koray Kavukcuoglu Raia Hadsell
references.title: Policy distillation
references.venue: arXiv preprint arXiv:1511.06295,
references.authors: Ozan Sener Vladlen Koltun
references.title: Multi-task learning as multi-objective optimization
references.venue: arXiv preprint arXiv:1810.04650,
references.authors: Baifeng Shi Judy Hoffman Kate Saenko Trevor Darrell Huijuan Xu
references.title: Auxiliary task reweighting for minimum-data learning
references.venue: Advances in Neural Information Processing Systems,
references.authors: Shagun Sodhani Amy Zhang
references.title: Mtrl - multi task rl algorithms
references.venue: Github,
references.authors: Shagun Sodhani Amy Zhang Joelle Pineau
references.title: Multi-task reinforcement learning with context-based representations
references.venue: arXiv preprint arXiv:2102.06177,
references.authors: Charles Stein
references.title: Inadmissibility of the usual estimator for the mean of a multivariate normal distribution
references.authors: Kevin Swersky Jasper Snoek Ryan Prescott Adams
references.title: Multi-task bayesian optimization
references.authors: Antti Tarvainen Harri Valpola
references.title: Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results
references.venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,
references.authors: Yee Whye Teh Victor Bapst Wojciech Marian Czarnecki John Quan James Kirkpatrick Raia Hadsell Nicolas Heess Razvan Pascanu
references.title: Distral: Robust multitask reinforcement learning
references.venue: arXiv preprint arXiv:1707.04175,
references.authors: Simon Vandenhende Stamatios Georgoulis Wouter Van Gansbeke Marc Proesmans Dengxin Dai Luc Van Gool
references.title: Multi-task learning for dense prediction tasks: A survey
references.venue: IEEE Transactions on Pattern Analysis and Machine Intelligence,
references.authors: Han Xiao Kashif Rasul Roland Vollgraf
references.title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms
references.venue: arXiv preprint arXiv:1708.07747,
references.authors: Ruihan Yang Huazhe Xu Yi Wu Xiaolong Wang
references.title: Multi-task reinforcement learning with soft modularization
references.venue: arXiv preprint arXiv:2003.13661,
references.authors: Tianhe Yu Saurabh Kumar Abhishek Gupta Sergey Levine Karol Hausman Chelsea Finn
references.title: Gradient surgery for multi-task learning
references.venue: arXiv preprint arXiv:2001.06782,
references.authors: Tianhe Yu Deirdre Quillen Zhanpeng He Ryan Julian Karol Hausman Chelsea Finn Sergey Levine
references.title: Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning
references.venue: In Conference on Robot Learning,
references.authors: Amir R Zamir Alexander Sax William Shen Leonidas J Guibas Jitendra Malik Silvio Savarese
references.title: Taskonomy: Disentangling task transfer learning
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
sections.heading: 1 Introduction
sections.text: 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.
sections.heading: 2 Background
sections.text: 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.
sections.heading: 2.1 Multi-task Learning and its Challenge
sections.text: 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.
sections.heading: 2.2 Prior Attempts and Convergence Issues
sections.text: 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.
sections.heading: 3 Method
sections.text: We introduce our main algorithm, Conflict-Averse Gradient descent in Sec. 3.1, and then show theoretical analysis in Sec. 3.2.
sections.heading: 3.1 Conflict-Averse Gradient Descent
sections.text: 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.
sections.heading: 3.2 Convergence Analysis
sections.text: In this section we first formally introduce the related Pareto concepts and then analyze CAGrad’s convergence property. Particularly, in Alg. 1, when c < 1, CAGrad is guaranteed to converge to a minimum point of the average loss L0.
  Pareto Concepts Unlike single task learning where any two parameter vectors θ1 and θ2 can be ordered in the sense that either L(θ1) ≤ L(θ2) or L(θ1) ≥ L(θ2) holds, MTL could have two parameter vectors where one performs better for task i and the other performs better for task j ̸= i. To this end, we need the notion of Pareto-optimality [13].
  Definition 3.1 (Pareto optimal and stationary points). Let L(θ) = {Li(θ) : i ∈ [K]} be a set of differentiable loss functions from Rm to R. For two points θ, θ′ ∈ Rm, we say that θ is Pareto dominated by θ′, denoted by L(θ′) ≺ L(θ), if Li(θ′) ≤ Li(θ) for all i ∈ [K] and L(θ′) ̸= L(θ). A point θ ∈ Rm is said to be Pareto-optimal if there exists no θ′ ∈ Rm such that L(θ′) ≺ L(θ). The set of all Pareto-optimal points is called the Pareto set. A point θ is called Pareto-stationary if we have minw∈W ∥gw(θ)∥ = 0, where gw(θ) = ∑K i=1 wi∇Li(θ), andW is the probability simplex on [K].
  Similar to the case of single-objective differentiable optimization, a local Pareto optimal point θ must be Pareto stationary (see e.g., [6]).
  Theorem 3.2 (Convergence of CAGrad). Assume the individual loss functions L0, L1, . . . , LK are differentiable on Rm and their gradients ∇Li(θ) are all H-Lipschitz, i.e. ∥∇Li(x)−∇Li(y)∥ ≤ H ∥x− y∥ for i = 0, 1, . . . ,K where 0 ≤ H ≤ ∞. Assume L∗0 = infθ∈Rm L0(θ) > −∞. With a fixed step size α satisfying 0 < α ≤ 1/H , we have for the CAGrad in Alg. 1:
  1) For any c ≥ 1, all the fixed points of CAGrad are Pareto-stationary points of (L0, L1, . . . , LK). 2) In particular, if we take 0 ≤ c < 1, then CAGrad satisfies
  T∑ t=0 ∥∇L0(θt)∥2 ≤ 2(L0(θ0)− L∗0) α(1− c2) .
  This means that the algorithm converges to a stationary point of∇L0 if we take 0 ≤ c < 1. The proof is in Appendix A.3. As we discuss earlier, unlike our method, MGDA is designed to converge to an arbitrary point on the Pareto set, without explicit control of which point it will converges to. Another algorithm with similar property is PCGrad [41], which is a gradient-based algorithm that mitigates the conflicting gradients problem by removing the conflicting components of each gradient with respect to the other gradients before averaging them to form the final update; see Fig. 2, third column for the illustration. Similar to MGDA, as shown in [41], PCGrad also converges to an arbitrary Pareto point without explicit control of which point it will arrive at.
sections.heading: 3.3 Practical Speedup
sections.text: A typical drawback of methods that manipulate gradients is the computation overhead. For computing the optimal update vector, a method usually requires K back-propagations to find all individual gradients gi, in addition to the time required for optimization. This can be prohibitive for the scenario with many tasks. To this end, we propose to only sample a subset of tasks S ⊆ [K], compute their corresponding gradients {gi | i ∈ S} and the averaged gradient g0. Then we optimize d in:
  max d∈Rm min
  ( ⟨ Kg0 − ∑ i∈S gi
  K − |S| , d⟩, min i∈S ⟨gi, d⟩
  ) s.t. ∥d− g0∥ ≤ c ∥g0∥ (4)
  Remark Note that the convergence guarantee in Thm. 3.2 still holds for Eq. 4 as the constraint does not change (See Appendix A.3). The time complexity is O((|S|N + T ), where N denotes the time for one pass of back-propagation and T denotes the optimization time. For few-task learning (K < 10), usually T ≪ N . When S = [K], we recover the full CAGrad algorithm.
sections.heading: 4 Related Work
sections.text: Multi-task Learning Due to its benefit with regards to data and computational efficiency, multi-task learning (MTL) has broad applications in vision, language, and robotics [11, 28, 22, 44, 38]. A number of MTL-friendly architectures have been proposed using task-specific modules [25, 11], attentionbased mechanisms [21] or activating different paths along the deep networks to tackle MTL [27, 40]. Apart from designing new architectures, another branch of methods focus on decomposing a large problem into smaller local problems that could be quickly learned by smaller models [29, 26, 37, 8]. Then a unified policy is learned from the smaller models using knowledge distillation [12].
  MTL Optimization In this work, we focus on the optimization challenge of MTL [38]. Gradient manipulation methods are designed specifically to balance the learning of each task. The simplest form of gradient manipulation is to re-weight the task losses based on specific criteria, e.g., uncertainty [15], gradient norm [3], or difficulty [9]. These methods are mostly heuristics and their performance can be unstable. Recently, two methods [30, 41] that manipulate the gradients to find a better local update vector have become popular. Sener et al [30] view MTL as a multi-objective optimization problem, and use multiple gradient descent algorithm for optimization. PCGrad [41] identifies a major optimization challenge for MTL, the conflicting gradients, and proposes to project each task gradient to the normal plane of other task gradients before combining them together to form the final update vector. Though yielding good empirical performance, both methods can only guarantee convergence to a Pareto-stationary point, but not knowing where it exactly converges to. More recently, GradDrop [4] randomly drops out task gradients based on how much they conflict. IMTLG [20] seeks an update vector that has equal projections on each task gradient. RotoGrad [14] separately scales and rotates task gradients to mitigate optimization conflict.
  Our method, CAGrad, also manipulates the gradient to find a better optimization trajectory. Like other MTL optimization techniques, CAGrad is model-agnostic. However, unlike prior methods, CAGrad converges to the optimal point in theory and achieves better empirical performance on both toy multi-objective optimization tasks and real-world applications.
sections.heading: 5 Experiment
sections.text: We conduct experiments to answer the following questions:
  Question (1) Do CAGrad, MGDA and PCGrad behave consistently with their theoretical properties in practice? (yes)
  Question (2) Does CAGrad recover GD and MGDA when varying the constant c? (yes)
  Question (3) How does CAGrad perform in both performance and computational efficiency compared to prior state-of-the-art methods, on challenging multi-task learning problems under the supervised, semi-supervised and reinforcement learning settings? (CAGrad improves over prior state-of-the-art methods under all settings)
sections.heading: 5.1 Convergence and Ablation over c
sections.text: To answer questions (1) and (2), we create a toy optimization example to evaluate the convergence of CAGrad compared to MGDA and PCGrad. On the same toy example, we ablate over the constant c and show that CAGrad recovers GD and MGDA with proper c values. Next, to test CAGrad on more complicated neural models, we perform the same set of experiments on the Multi-Fashion+MNIST benchmark [19] with a shrinked LeNet architecture [18] (in which each layer has a reduced number of neurons compared to the original LeNet). Please refer to Appendix B for more details.
  For the toy optimization example, we modify the toy example used by Yu et al. [41] and consider θ = (θ1, θ2) ∈ R2 with the following individual loss functions:
  L1(θ) = c1(θ)f1(θ) + c2(θ)g1(θ) and L2(θ) = c1(θ)f2(θ) + c2(θ)g2(θ), where f1(θ) = log ( max(|0.5(−θ1 − 7)− tanh (−θ2)|, 0.000005) ) + 6,
  f2(θ) = log ( max(|0.5(−θ1 + 3)− tanh (−θ2) + 2|, 0.000005) ) + 6,
  g1(θ) = ( (−θ1 + 7)2 + 0.1 ∗ (−θ2 − 8)2 ) /10− 20,
  g2(θ) = ( (−θ1 − 7)2 + 0.1 ∗ (−θ2 − 8)2) / 10− 20,
  c1(θ) = max(tanh (0.5 ∗ θ2), 0) and c2(θ) = max(tanh (−0.5 ∗ θ2), 0).
  The average loss L0 and individual losses L1 and L2 are shown in Fig. 1. We then pick 5 initial parameter vectors θinit ∈ {(−8.5, 7.5), (−8.5, 5), (0, 0), (9, 9), (10,−8)} and plot the corresponding optimization trajectories with different methods in Fig. 3. As shown in Fig. 3, GD gets stuck in 2 out of the 5 runs while other methods all converge to the Pareto set. MGDA and PCGrad converge to different Pareto-stationary points depending on θinit. CAGrad with c = 0 recovers GD and CAGrad with c = 10 approximates MGDA well (in theory it requires c→∞ to exactly recover MGDA).
  Next, we apply the same set of experiments on the multi-task classification benchmark MultiFashion+MNIST [19]. This benchmark consists of images that are generated by overlaying an image from FashionMNIST dataset [39] on top of another image from MNIST dataset [5]. The two images are positioned on the top-left and bottom-right separately. We consider a shrinked LeNet as our model, and train it with Adam [16] optimizer with a 0.001 learning rate for 50 epochs using a batch size of 256. Due to the highly non-convex nature of the neural network, we are not able to visualize the entire Pareto set. But we provide the final training losses of different methods over three independent runs in Fig. 4. As shown, CAGrad achieves the lowest average loss with c = 0.2. In addition, PCGrad and MGDA focus on optimizing task 1 and task 2 separately. Lastly, CAGrad with c = 0 and c = 10 roughly recovers the final performance of GD and MGDA. By increasing c, the model performance shifts from more GD-like to more MGDA-like, though due to the non-convex nature of neural networks, CAGrad with 0 ≤ c < 1 does not necessarily converge to the exact same point.
sections.heading: 5.2 Multi-task Supervised Learning
sections.text: To answer question (3) in the supervised learning setting, we follow the experiment setup from Yu et al. [41] and consider the NYU-v2 and CityScapes vision datasets. NYU-v2 contains 3 tasks: 13- class semantic segmentation, depth estimation, and surface normal prediction. CityScapes similarly contains 2 tasks: 7-class semantic segmentation and depth estimation. Here, we follow [41] and combine CAGrad with a state-of-the-art MTL method MTAN [21], which applies attention mechanism on top of the SegNet architecture [1]. We compare CAGrad with PCGrad, vanilla MTAN and CrossStitch [25], which is another MTL method that modifies the network architecture. MTAN originally experiments with equal loss weighting and two other dynamic loss weighting heuristics [15, 3]. For a fair comparison, all methods are applied under the equal weighting scheme and we use the same training setup from [3]. We search c ∈ {0.1, 0.2, . . . 0.9} with the best average training loss for CAGrad on both datasets (0.4 for NYU-v2 and 0.2 for Cityscapes). We perform a two-tailed, Student’s t-test under equal sample sizes, unequal variance setup and mark the results that are significant with an ∗. Following Maninis et al.[24], we also compute the average per-task performance drop of method m with respect to the single-tasking baseline b: ∆m = 1K ∑K i=1(−1)li(Mm,i −Mb,i)/Mb,i where li = 1 if a higher value is better for a criterion Mi on task i and 0 otherwise. The single-tasking baseline (independent) refers to training individual tasks with a vanilla SegNet. Results are shown in Tab. 1 and Tab. 2.
  Given the single task performance, CAGrad performs better on the task that is overlooked by other methods (Surface Normal in NYU-v2 and Depth in CityScapes) and matches other methods’
  performance on the rest of the tasks. We also provide the final test losses and the per-epoch training time of each method in Fig. 5 in Appendix B.2.
sections.heading: 5.3 Multi-task Reinforcement Learning
sections.text: To answer question (3) in the reinforcement learning (RL) setting, we apply CAGrad on the MT10 and MT50 benchmarks from the Meta-World environment [42]. In particular, MT10 and MT50 contains 10 and 50 robot manipulation tasks. Following [33], we use Soft Actor-Critic (SAC) [10] as the underlying RL training algorithm. We compare against Multi-task SAC (SAC with a shared model), Multi-headed SAC (SAC with a shared backbone and task-specific head), Multi-task SAC + Task Encoder (SAC with a shared model and the input includes a task embedding) [42] and PCGrad [41]. We also compare with Soft Modularization [40] that routes different modules in a shared model to form different policies. Lastly, we also include a recent method (CARE) that considers language metadata and uses a mixture of expert encoder for MTL. We follow the same experiment setup from [33]. The results are shown in Tab. 3. CAGrad outperforms all baselines except for CARE which benefits from extra information from the metadata. We also apply the practical speedup in Sec. 3.3 and sub-sample 4 and 8 tasks for MT10 and MT50 (CAGrad-Fast). CAGrad-fast achieves comparable performance against the state-of-the-art method while achieving a 2x (MT10) and 5x (MT50) speedup over PCGrad. We provide a visualization of tasks from MT10 and MT50, and the comparison of computational efficiency in Appendix B.3.
sections.heading: 5.4 Semi-supervised Learning with Auxiliary Tasks
sections.text: Training with auxiliary tasks to improve the performance of a main task is another popular application of MTL. Here, we take semi-supervised learning as an instance. We combine different optimization
  algorithms with Auxiliary Task Reweighting for Minimum-data Learning (ARML) [31], a state-ofthe-art semi-supervised learning algorithm. The loss function is composed of the main task and two auxiliary tasks:
  L0 = LCE(θ;Dl) + w1L 1 aux(θ;Du) + w2L 2 aux(θ;Du), (5)
  where LCE is the main cross-entropy classification loss on the labeled dataset Dl, and L1aux, L 2 aux are auxiliary unsupervised learning losses on the unlabeled dataset Du. We use the same w1 and w2 from ARML, and use the CIFAR10 dataset [17], which contains 50,000 training images and 10,000 test images. 10% of the training images is held out as the validation set. We test PCGrad, MGDA and CAGrad with 500, 1000 and 2000 labeled images. The rest of the training set is used for auxiliary tasks. For all the methods, we use the same labeled dataset, the same learning rate and train them for 200 epochs with the Adam [16] optimizer. Please refer to Appendix B.4 for more experimental details. Results are shown in Tab. 4. With all the different number of labels, CAGrad yields the best averaged test accuracy. We observed that MGDA performs much worse than the ARML baseline, because it significantly overlooks the main classification task. We also compare different gradient manipulation methods on the same task with GradNorm [3], which dynamically adjusts w1 and w2 during training. The results and conclusions are similar to those for ARML.
sections.heading: 6 Conclusion
sections.text: In this work, we introduce the Conflict-Averse Gradient descent (CAGrad) algorithm that explicitly optimizes the minimum decrease rate of any specific task’s loss while still provably converging to the optimum of the average loss. CAGrad generalizes the gradient descent and multiple gradient descent algorithm, and demonstrates improved performance across several challenging multi-task learning problems compared to the state-of-the-art methods. While we focus mainly on optimizing the average loss, an interesting future direction is to look at main objectives other than the average loss under the multi-task setting.
sections.heading: Acknowledgements
sections.text: The research was conducted in the statistical learning and AI group (SLAI) and the Learning Agents Research Group (LARG) in computer science at UT Austin. SLAI research is supported in part by CAREER-1846421, SenSE-2037267, EAGER-2041327, and Office of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML). LARG research is supported in part by NSF (CPS-1739964, IIS-1724157, FAIN-2019844), ONR (N00014-18-2243), ARO (W911NF-19-2-0333), DARPA, Lockheed Martin, GM, Bosch, and UT Austin’s Good Systems grand challenge. Peter Stone serves as the Executive Director of Sony AI America and receives financial compensation for this work. The terms of this arrangement have been reviewed and approved by the University of Texas at Austin in accordance with its policy on objectivity in research. Xingchao Liu is supported in part by a funding from BP.
title: Conflict-Averse Gradient Descent for Multi-task Learning