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

TL;DR

This paper introduces MARIGOLD, a bi-level optimization framework that efficiently balances gradients in multi-task learning, improving performance and reducing computational costs compared to existing methods.

Contribution

The paper proposes a novel bi-level optimization approach for multi-task gradient balancing, addressing efficiency issues of prior gradient-based methods like MGDA.

Findings

MARIGOLD outperforms existing methods on public datasets.

It reduces computational costs in multi-task learning.

Demonstrates effectiveness on industrial-scale datasets.

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.