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.
Figures & tables
Figure 1: In Figure 1(a) 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 1(b) we showcase how MGDA -type algorithms ( 59 ) 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 Φ(λ,ρ) and Lower-Level (LL) function g(λ,ρ) are defined in ( 7 ) and ( 8 ).
Nanyang Technological University · Hong Kong University of Science and Technology (Guangzhou) · School of Artificial Intelligence, Shanghai Jiao Tong University
1Cheriton School of Computer Science, University of Waterloo, Waterloo, ON, Canada · 3Vector Institute · 2Computer Science, The Chinese University of Hong Kong