Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Feed m Birds with One Scone: Accelerating Multi-task Gradient Balancing via Bi-level Optimization

Xuxing Chen, Yun He, Jiayi Xu, Minhui Huang, Xiaoyi Liu, Boyang Liu, Fei Tian, Xiaohan Wei, Rong Jin, Sem Park, Bo Long, Xue Feng

TL;DR

This paper introduces MARIGOLD, a unified algorithmic framework for efficiently solving MTL problems and reveals that multi-task gradient balancing methods have a hierarchical structure, in which the model training and the gradient balancing are coupled during the whole optimization process.

Abstract

In machine learning, the goal of multi-task learning (MTL) is to optimize multiple objectives together. Recent works, for example, Multiple Gradient Descent Algorithm (MGDA) and its variants, show promising results with dynamically adjusted weights for different tasks to mitigate conflicts that may potentially degrade the performance on certain tasks. Despite the empirical success of MGDA-type methods, one major limitation of such methods is their computational inefficiency, as they require access to all task gradients. In this paper we introduce MARIGOLD, a unified algorithmic framework for efficiently solving MTL problems. Our method reveals that multi-task gradient balancing methods have a hierarchical structure, in which the model training and the gradient balancing are coupled during the whole optimization process and can be viewed as a bi-level optimization problem. Moreover, we showcase that the bi-level problem can be solved efficiently by leveraging zeroth-order method. Extensive experiments on both public datasets and industrial-scale datasets demonstrate the efficiency and superiority of our method.

Feed m Birds with One Scone: Accelerating Multi-task Gradient Balancing via Bi-level Optimization

TL;DR

This paper introduces MARIGOLD, a unified algorithmic framework for efficiently solving MTL problems and reveals that multi-task gradient balancing methods have a hierarchical structure, in which the model training and the gradient balancing are coupled during the whole optimization process.

Abstract

In machine learning, the goal of multi-task learning (MTL) is to optimize multiple objectives together. Recent works, for example, Multiple Gradient Descent Algorithm (MGDA) and its variants, show promising results with dynamically adjusted weights for different tasks to mitigate conflicts that may potentially degrade the performance on certain tasks. Despite the empirical success of MGDA-type methods, one major limitation of such methods is their computational inefficiency, as they require access to all task gradients. In this paper we introduce MARIGOLD, a unified algorithmic framework for efficiently solving MTL problems. Our method reveals that multi-task gradient balancing methods have a hierarchical structure, in which the model training and the gradient balancing are coupled during the whole optimization process and can be viewed as a bi-level optimization problem. Moreover, we showcase that the bi-level problem can be solved efficiently by leveraging zeroth-order method. Extensive experiments on both public datasets and industrial-scale datasets demonstrate the efficiency and superiority of our method.
Paper Structure (17 sections, 1 theorem, 19 equations, 3 figures, 8 tables, 2 algorithms)

This paper contains 17 sections, 1 theorem, 19 equations, 3 figures, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Multi-task learning and multi-objective optimization
  4. Bi-level optimization
  5. Methodology
  6. Rethinking gradient balancing criterion
  7. Multi-task gradient balancing as bi-level optimization
  8. Hypergradient estimation via zeroth-order method
  9. Algorithms
  10. Beyond worst-case decrement -- a general framework
  11. Experiments
  12. Public data
  13. Industrial-scale data
  14. Conclusion
  15. Related Work
  16. ...and 2 more sections

Key Result

Lemma 3.1

Suppose $f(\theta):\operatorname{\mathbb R}^d \rightarrow \operatorname{\mathbb R}$ is differentiable and $\ell$-smooth, i.e., $\left\|\nabla f(\theta_1) - \nabla f(\theta_2)\right\|\leq \ell \left\|\theta_1 - \theta_2\right\|$ for any $\theta_1, \theta_2\in\operatorname{\mathbb R}^d$. Define the cl

Figures (3)

  • Figure 1: In Figure \ref{['fig: arch']} we have a sketch of a model architecture for multi-task learning. In the forward pass, a batch of data is fed into the bottom layers shared by all tasks to provide a representation for task-specific modules (also known as "heads"). Then task-specific losses are constructed based on the prediction output of heads. Finally a linear combination of all losses gives the main objective function to optimize. In the backward pass, model parameters receive gradients and are updated via certain optimization algorithms as the model training step, and the task weights are balanced based on pre-specified criterion as the tasks balancing step. In Figure \ref{['fig: mgda']} we showcase how MGDA-type algorithms sener2018multi work -- the gradient balancing step tries to find a convex combination of gradients with minimum norm, and updates the model parameters accordingly.
  • Figure 2: MTL as a bi-level optimization problem. The Upper-Level (UL) function $\Phi(\lambda, \rho)$ and Lower-Level (LL) function $g(\lambda,\rho)$ are defined in \ref{['opt: bo_ul']} and \ref{['opt: bo_ll']}.
  • Figure 3: Comparison between FAMO and Algorithm \ref{['alg: marigold']} in terms of $\Delta k\%$ over time.

Theorems & Definitions (2)

  • Definition 2.1
  • Lemma 3.1