LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control
Abstract
Multi-task learning (MTL) has been widely adopted for itsability to simultaneously learn multiple tasks. While existing gradientmanipulation methods often yield more balanced solutions than simplescalarization-based approaches, they typically incur a significant compu-tational overhead of O(K) in both time and memory, where K is thenumber of tasks. In this paper, we propose LDC-MTL, a simple andscalable loss discrepancy control approach for MTL, formulated froma bilevel optimization perspective. Our method incorporates two keycomponents: (i) a bilevel formulation for fine-grained loss discrepancycontrol, and (ii) a scalable first-order bilevel algorithm that requiresonly O(1) time and memory. Theoretically, we prove that LDC-MTLguarantees convergence not only to a stationary point of the bilevelproblem with loss discrepancy control but also to an ϵ-accurate Paretostationary point for all K loss functions under mild conditions. Extensiveexperiments on diverse multi-task datasets demonstrate the superior per-formance of LDC-MTL in both accuracy and efficiency. Code is availableat https://github.com/OptMN-Lab/LDC-MTL.