LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control

TL;DR

LDC-MTL introduces a scalable bilevel optimization method for multi-task learning that effectively balances tasks with minimal computational overhead, outperforming existing gradient manipulation techniques.

Contribution

It proposes a novel bilevel optimization framework for MTL that ensures balanced task learning with only constant time and memory complexity per iteration.

Findings

LDC-MTL achieves superior accuracy on multi-task datasets.

The method converges to a Pareto stationary point under mild conditions.

It significantly reduces computational overhead compared to existing methods.

Abstract

Multi-task learning (MTL) has been widely adopted for its ability to simultaneously learn multiple tasks. While existing gradient manipulation methods often yield more balanced solutions than simple scalarization-based approaches, they typically incur a significant computational overhead of $\mathcal{O}(K)$ in both time and memory, where $K$ is the number of tasks. In this paper, we propose LDC-MTL, a simple and scalable loss discrepancy control approach for MTL, formulated from a bilevel optimization perspective. Our method incorporates two key components: (i) a bilevel formulation for fine-grained loss discrepancy control, and (ii) a scalable first-order bilevel algorithm that requires only $\mathcal{O}(1)$ time and memory. Theoretically, we prove that LDC-MTL guarantees convergence not only to a stationary point of the bilevel problem with loss discrepancy control but also to an $\epsilon$-accurate Pareto stationary point for all $K$ loss functions under mild conditions. Extensive experiments on diverse multi-task datasets demonstrate the superior performance of LDC-MTL in both accuracy and efficiency.