Improvable Gap Balancing for Multi-Task Learning

TL;DR

This paper introduces two novel improvable gap balancing algorithms for multi-task learning that dynamically assign task weights to improve performance, demonstrating superior results on benchmark datasets and complementing gradient balancing methods.

Contribution

The paper proposes the first deep reinforcement learning-based improvable gap balancing algorithm for MTL, enhancing loss balancing with dynamic task weight assignment.

Findings

IGB algorithms outperform existing loss balancing methods.

Combining IGB with gradient balancing yields further performance improvements.

Extensive experiments validate the effectiveness of the proposed methods.

Abstract

In multi-task learning (MTL), gradient balancing has recently attracted more research interest than loss balancing since it often leads to better performance. However, loss balancing is much more efficient than gradient balancing, and thus it is still worth further exploration in MTL. Note that prior studies typically ignore that there exist varying improvable gaps across multiple tasks, where the improvable gap per task is defined as the distance between the current training progress and desired final training progress. Therefore, after loss balancing, the performance imbalance still arises in many cases. In this paper, following the loss balancing framework, we propose two novel improvable gap balancing (IGB) algorithms for MTL: one takes a simple heuristic, and the other (for the first time) deploys deep reinforcement learning for MTL. Particularly, instead of directly balancing the losses in MTL, both algorithms choose to dynamically assign task weights for improvable gap balancing. Moreover, we combine IGB and gradient balancing to show the complementarity between the two types of algorithms. Extensive experiments on two benchmark datasets demonstrate that our IGB algorithms lead to the best results in MTL via loss balancing and achieve further improvements when combined with gradient balancing. Code is available at https://github.com/YanqiDai/IGB4MTL.