Conflict-Averse Gradient Descent for Multi-task Learning

TL;DR

This paper introduces Conflict-Averse Gradient descent (CAGrad), a new optimization method for multi-task learning that balances task gradients to improve performance and guarantees convergence, outperforming previous heuristics.

Contribution

CAGrad is a novel gradient descent algorithm that minimizes average loss while regularizing based on worst task improvement, with proven convergence and superior empirical results.

Findings

CAGrad outperforms prior methods on multi-task learning benchmarks.

It guarantees convergence to a minimum of the average loss.

It effectively balances conflicting task gradients.

Abstract

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.