Source-linked AI summary

Gradient Surgery for Multi-Task Learning

Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, Chelsea Finn

arXiv:2001.06782v4cs.LGcs.CVcs.ROstat.ML

TL;DR

Multi-task learning promises more efficient learning through shared structure, but its optimization challenges are not fully understood. The paper identifies a three-condition source of detrimental gradient interference and proposes PCGrad, which projects conflicting gradients to reduce interference. Across supervised and reinforcement-learning problems, PCGrad substantially improves optimization efficiency and performance.

  • Problem

    Multi-task optimization can be less efficient and perform worse than independent learning, while the optimization challenges behind this difficulty remain insufficiently understood.

  • Method

    PCGrad projects each task gradient onto the normal plane of any conflicting task gradient, modifying shared-parameter gradients without depending on model architecture.

  • Results

    PCGrad substantially improves optimization across a variety of multi-task supervised-learning and reinforcement-learning problems, with gains in efficiency and asymptotic performance.

  • Takeaways & Limitations

    Mitigating conflicting gradients can enable multi-task settings to realize hypothesized efficiency and asymptotic-performance benefits.

  • Takeaways & Limitations

    In theory, PCGrad can converge slowly when the cosine similarity between task gradients remains near −1, and exactly −1 can yield a sub-optimal solution.

Abstract

from arXiv · show

While deep learning and deep reinforcement learning (RL) systems have demonstrated impressive results in domains such as image classification, game playing, and robotic control, data efficiency remains a major challenge. Multi-task learning has emerged as a promising approach for sharing structure across multiple tasks to enable more efficient learning. However, the multi-task setting presents a number of optimization challenges, making it difficult to realize large efficiency gains compared to learning tasks independently. The reasons why multi-task learning is so challenging compared to single-task learning are not fully understood. In this work, we identify a set of three conditions of the multi-task optimization landscape that cause detrimental gradient interference, and develop a simple yet general approach for avoiding such interference between task gradients. We propose a form of gradient surgery that projects a task's gradient onto the normal plane of the gradient of any other task that has a conflicting gradient. On a series of challenging multi-task supervised and multi-task RL problems, this approach leads to substantial gains in efficiency and performance. Further, it is model-agnostic and can be combined with previously-proposed multi-task architectures for enhanced performance.

1 Introduction

Multi-task learning aims to share structure across tasks for greater efficiency, but conflicting optimization dynamics can make joint training less effective than independent learning. This work identifies conditions behind gradient interference and introduces PCGrad to mitigate it.

  • Multi-task learning seeks greater efficiency and performance by training tasks jointly and discovering shared structure.Joint training is motivated by avoiding the data demands of learning every task independently.
  • Joint optimization can instead produce worse overall performance and data efficiency than solving tasks independently.Some multi-task RL methods therefore train independent models before distilling them into a multi-task model, sacrificing efficiency gains.
  • The paper hypothesizes that detrimental interference occurs when conflicting gradients coincide with high positive curvature and large gradient-magnitude differences.Gradients conflict when they point away from one another, measured by negative cosine similarity.
  • Deep valleys with high curvature and unequal gradient magnitudes can let one task dominate the multi-task gradient, degrading another task's performance.The illustrative 2D landscape uses task objectives with deep valleys and highlights this domination effect.
  • PCGrad mitigates gradient interference by projecting each conflicting task gradient onto the normal plane of the other gradient.The method is model-agnostic, requires one modification to gradient application, and can combine with other multi-task architectures.

2 Multi-Task Learning with PCGrad

Multi-task optimization can suffer when conflicting, unevenly scaled gradients interact with high curvature. PCGrad addresses this by removing conflicting gradient components while leaving non-conflicting gradients unchanged, with theoretical and empirical support.

  • 2.1 Preliminaries: Problem and Notation: The multi-task objective minimizes the sum of task losses over model parameters, with task-conditioned predictions using task encodings.Task gradients are defined separately for each task and combined through the multi-task loss.
  • 2.2 The Tragic Triad: Conflicting Gradients, Dominating Gradients, High Curvature: The tragic triad comprises conflicting gradients, large gradient-magnitude differences, and high curvature, which together can degrade multi-task optimization.Conflicting gradients alone are not necessarily harmful; the adverse conditions arise from their co-occurrence with domination and curvature.
  • 2.3 PCGrad: Project Conflicting Gradients: Non-conflicting gradients remain unchanged, preserving the possibility of constructive interaction between tasks.The method alters gradients only when their cosine similarity indicates conflict.
  • 2.3 PCGrad: Project Conflicting Gradients: PCGrad detects negative cosine similarity and projects each conflicting task gradient onto the normal plane of the other task’s gradient.The procedure repeats across other tasks in random order and returns the modified gradients for updating the model.
  • 2.3 PCGrad: Project Conflicting Gradients: PCGrad is model-agnostic and can be applied by passing its modified update to standard gradient-based optimizers such as SGD with momentum or Adam.The reported procedure reduces gradient interference and substantially improves learning progress in experiments.
  • 2.4 Theoretical Analysis of PCGrad: Under the stated two-task assumptions, PCGrad can achieve lower loss than standard multi-task gradient descent when conflict, gradient imbalance, and curvature satisfy sufficient conditions.Theoretical convergence reaches the multitask minimizer or a potentially sub-optimal point when gradients directly oppose one another; minibatch noise makes exact opposition unlikely in practice.

3 PCGrad in Practice

PCGrad is applied to supervised and reinforcement-learning settings by computing task-specific gradients and modifying them before optimization. The same gradient-surgery procedure extends to actor-critic methods by replacing both actor and critic task gradients.

  • Supervised Learning: In supervised learning, each task uses a labeled dataset and a task-conditioned negative-log-likelihood objective.Task identity is represented with a one-hot encoding, and examples are grouped by task within each sampled batch.
  • Supervised Learning: Task-specific gradients and pairwise cosine similarities are precomputed for the tasks represented in each batch before applying PCGrad.These quantities provide the inputs needed to deconflict gradients within the batch.
  • Reinforcement Learning: In multi-task and goal-conditioned RL, PCGrad directly modifies each task’s policy gradient, while actor-critic methods apply it to both actor and critic gradients.The procedure follows the same update rule used in supervised learning.

4 Related Work

Related multi-task learning work spans supervised learning, reinforcement learning, and domains including vision, language, and robotics. Prior approaches address optimization through architectures, decomposition, or other strategies, while this work targets gradient conflict directly.

  • Prior Multi-Task Learning: Multi-task learning trains a single model across varied tasks and has been applied in supervised learning, reinforcement learning, vision, language, and robotics.Its promise of accelerating acquisition of task repertoires is accompanied by a challenging optimization problem.
  • Architectural Approaches: Architectural approaches use multiple modules, paths, or attention mechanisms to address multi-task learning difficulties.The paper describes its own method as architecture-agnostic and complementary to these approaches.
  • Gradient-Based Approaches: This work attributes multi-task optimization difficulty to the tragic triad of conflicting gradients, high curvature, and large gradient differences.It addresses these conditions with a simple algorithm that deconflicts gradients from different tasks.
  • Continual Learning: Continual-learning methods also use gradient projection, but they target sequential learning and employ projection schemes involving past-task gradients.The cited approaches include quadratic programming, projection onto average past gradients, and projection onto orthonormal sets.

5 Experiments

Experiments evaluate PCGrad across supervised and reinforcement-learning benchmarks, including combinations with existing multi-task architectures and ablations of its gradient updates. PCGrad generally improves efficiency and performance, with results indicating benefits from modifying both gradient directions and magnitudes.

  • Experimental goals: The experiments test PCGrad across supervised, reinforcement-learning, and goal-conditioned reinforcement-learning settings, including whether it improves existing multi-task methods.The evaluation also examines whether the proposed optimization conditions explain multi-task difficulty.
  • Multi-Task Supervised Learning: 71% classification accuracy on CIFAR-100 is achieved by a single network using PCGrad, outperforming most prior methods.The dataset treats 20 coarse labels as distinct tasks, and PCGrad is also combined with routing networks.
  • Multi-Task Supervised Learning: PCGrad outperforms Sener and Koltun on CelebA average classification error across 40 tasks.This result suggests effectiveness for multi-label classification with many tasks.
  • Multi-Task Supervised Learning: MTAN with PCGrad achieves the best scores in 8 of 9 categories on the three-task NYUv2 benchmark.NYUv2 covers semantic segmentation, depth estimation, and surface normal prediction.
  • Multi-Task Reinforcement Learning: PCGrad with SAC solves all 10 MT10 tasks and about 70% of MT50 tasks with the best data efficiency.Independent SAC agents require about 2 million and 15 million more samples than PCGrad with SAC on MT10 and MT50, respectively.
  • Multi-Task Reinforcement Learning: PCGrad outperforms direction-only, magnitude-only, and GradNorm variants, indicating that both gradient directions and magnitudes matter.The ablation appears in the rightmost plot of Figure 3, alongside MT10 and MT50 learning curves.

6 Conclusion

The paper identifies conflicting gradients, high positive curvature, and large gradient differences as major multi-task optimization challenges, and proposes PCGrad to mitigate them. Across supervised and reinforcement-learning problems, alleviating these challenges improves optimization efficiency and asymptotic performance, while broader applications remain future directions.

  • Conclusion: The paper identifies conflicting gradients, high positive curvature, and large gradient differences as conditions underlying major multi-task optimization challenges.These conditions are described as the tragic triad of multi-task learning.
  • Conclusion: PCGrad uses gradient surgery to mitigate these challenges and substantially improve optimization across multi-task supervised and reinforcement-learning problems.The method is presented as a simple algorithm for deconflicting gradients from different tasks.
  • Future directions: The authors suspect conflicting gradients may also occur in meta-learning, continual learning, multi-goal imitation learning, and multi-task natural language processing.Applying PCGrad in these settings is proposed as a promising direction for future investigation.

Broader Impact

The paper frames PCGrad as a way to improve data-efficient multi-task learning, while noting broader machine-learning and reinforcement-learning risks. These include safety, reward specification, dataset bias, and compute- or data-intensive training.

  • Applications and Benefits: PCGrad is intended to mitigate interference between tasks and enable data-efficient multi-task learning.The broader-impact discussion contrasts this with prior multi-task methods that can have comparable or higher data complexity than independent learning.
  • Risks: The paper notes risks from unsafe out-of-distribution behavior, difficult real-world reward specification, training-data bias, and compute- or data-intensive procedures.The authors state that their method is not immune to these risks.

A.1 Proof of Theorem 1

Under convex, differentiable objectives and a Lipschitz multi-task gradient, the analysis characterizes when PCGrad decreases the objective and where repeated updates can converge. The results extend to non-convex and momentum-based settings under additional assumptions.

  • Convex convergence: Theorem 1 shows that, with sufficiently small step size, PCGrad converges either to the multi-task minimizer or to a point where the task gradients directly oppose one another.The latter case occurs when cos(φ12) = −1, making the PCGrad-modified gradient zero.
  • Convex convergence: When task gradients are non-conflicting, the standard gradient update strictly decreases the convex objective unless its gradient is zero.For conflicting gradients, the proof uses a quadratic expansion under Lipschitz-gradient assumptions.
  • Multiple objectives: For multiple convex objectives, convergence likewise reaches the optimum or a location where every pair of task gradients has cosine similarity −1.This conclusion assumes the stated Lipschitz and step-size conditions.
  • Non-convex convergence: In the non-convex two-task setting, PCGrad converges either to a location with cos(φ12) = −1 or to an almost stationary point.The convergence rate depends on α, which lower-bounds the gradient cosine similarity away from −1.
  • One-step loss comparison: Theorem 2 gives sufficient conditions for one PCGrad update to achieve lower loss than the standard multi-task update, involving gradient conflict, curvature, and step size.The comparison requires the stated bounds on cos φ12, multi-task curvature, and t.
  • Momentum-based convergence: With strongly convex and smooth objectives, momentum-based PCGrad converges linearly either to a fully opposing-gradient location or the optimum.The heavy-ball analysis defines curvature and smoothness bounds from task-gradient angles and magnitudes.

B Empirical Objective-Wise Evaluations of PCGrad

On NYUv2, PCGrad combined with MTAN matches MTAN’s training convergence rate while improving validation convergence and final loss on two of three tasks. The broader supervised-learning results suggest a regularization effect rather than faster optimization.

  • NYUv2 evaluations: PCGrad+MTAN has a similar training convergence rate to MTAN across all three NYUv2 tasks.The comparison uses objective-wise training learning curves.
  • NYUv2 evaluations: PCGrad+MTAN converges faster and achieves lower final validation loss in two out of three NYUv2 tasks.The validation comparison is reported from the objective-wise curves.
  • NYUv2 evaluations: Both methods appear to overfit on NYUv2 task 0, suggesting that this domain may require a better regularization scheme.This observation is specific to task 0.
  • Interpretation: The results suggest that PCGrad’s supervised multi-task benefit is a regularization effect rather than an improvement in optimization speed or convergence.The authors hypothesize that greater representation sharing allows supervision from one task to regularize another.

E Additional Multi-Task Supervised Learning Results

PCGrad improves multi-task supervised learning performance across MultiMNIST and CityScapes, including when combined with the MTAN architecture.

  • MultiMNIST: 0.13% and 0.55% improvements over Sener and Koltun [53] were achieved in left and right digit accuracy, respectively.These results were obtained on the two MultiMNIST digit-classification tasks.
  • CityScapes: PCGrad + MTAN outperforms MTAN in three out of four CityScapes scores.The combined method achieves the top scores in mIoU and pixel accuracy for semantic segmentation.
  • CityScapes: The CityScapes evaluation combines seven-class semantic segmentation with depth estimation.The comparison uses the MTAN architecture with equal weighting.

F Goal-Conditioned Reinforcement Learning Results

In goal-conditioned robotic pushing, PCGrad improves both goal-reaching performance and data efficiency relative to vanilla SAC, and outperforms CosReg in MT10.

  • Goal-conditioned RL: PCGrad outperforms vanilla SAC in average distance to the goal and data efficiency.The comparison is reported for goal-conditioned reinforcement learning results.
  • Comparison with CosReg: PCGrad achieves a much better average success rate than CosReg in the MT10 benchmark.The comparison suggests that reducing interference should not remove task gradients with positive cosine similarity.
  • Task-order ablation: Random task ordering performs better than fixed task ordering in the MT50 benchmark.Random shuffling makes PCGrad symmetric with respect to task order in expectation.

I Combining PCGrad with other architectures

PCGrad further improves several multi-task architectures, including Cross-Stitch, Dense, and Multi-head SAC, showing compatibility across supervised and reinforcement-learning settings.

  • Combined architectures: PCGrad improves all four metrics across the three NYUv2 tasks when combined with Cross-Stitch and Dense.The results support combining PCGrad with existing multi-task architectures.
  • Combined architectures: PCGrad + Multi-head SAC outperforms Multi-head SAC alone.This extends the reported benefits of PCGrad to a multi-headed reinforcement-learning architecture.
  • Combined architectures: PCGrad is tested with prior architectures rather than only as a standalone optimization modification.The supervised comparisons include Cross-Stitch and Dense, while the reinforcement-learning comparison uses Multi-head SAC.
  • Implementation: PCGrad converges within 12 hours and uses at most 10 GB of GPU memory in the supervised experiments.Vanilla models converge within 8 hours and use 6 GB among the reported experiments.

J.2 Multi-Task Reinforcement Learning Experiment Details

The multi-task reinforcement-learning experiments apply PCGrad to SAC on Meta-World’s MT10 and MT50 benchmarks, using shared task-conditioned actor and critic networks and task-averaged success rates.

  • Training cost: PCGrad + SAC converges in 1 day on MT10 and 5 days on MT50, compared with 12 hours and 3 days for vanilla SAC.The reported PCGrad runs use 5M and 20M simulation steps on MT10 and MT50, respectively.
  • Meta-World benchmarks: MT10 and MT50 contain 10 and 50 diverse simulated robotic manipulation tasks, respectively.Both benchmarks use a shared tabletop environment with a simulated Sawyer arm.
  • Evaluation protocol: Success rates are averaged across tasks, with 600 samples collected per task at each data-collection step.Training samples are drawn from task-specific replay buffers.
  • Algorithms and models: PCGrad is applied to vanilla SAC with task encoding provided to both the actor and critic.The comparison includes vanilla SAC without PCGrad and independent per-task actor and critic training.
  • Additional results: The full CityScapes and NYUv2 comparisons are provided in Tables 7 and 8.These tables report the broader supervised-learning results referenced by the experiment details.
Loading 2001.06782v4…