Independent Component Alignment for Multi-Task Learning
This work addresses optimization instability in multi-task learning, which is a problem for researchers and practitioners seeking efficient joint training of models, though it appears incremental as it builds on existing convergence guarantees.
The paper tackles optimization challenges in multi-task learning, such as conflicting gradients, by proposing Aligned-MTL, a method that aligns orthogonal gradient components based on a stability criterion, and demonstrates consistent performance improvements across diverse benchmarks including segmentation and reinforcement learning.
In a multi-task learning (MTL) setting, a single model is trained to tackle a diverse set of tasks jointly. Despite rapid progress in the field, MTL remains challenging due to optimization issues such as conflicting and dominating gradients. In this work, we propose using a condition number of a linear system of gradients as a stability criterion of an MTL optimization. We theoretically demonstrate that a condition number reflects the aforementioned optimization issues. Accordingly, we present Aligned-MTL, a novel MTL optimization approach based on the proposed criterion, that eliminates instability in the training process by aligning the orthogonal components of the linear system of gradients. While many recent MTL approaches guarantee convergence to a minimum, task trade-offs cannot be specified in advance. In contrast, Aligned-MTL provably converges to an optimal point with pre-defined task-specific weights, which provides more control over the optimization result. Through experiments, we show that the proposed approach consistently improves performance on a diverse set of MTL benchmarks, including semantic and instance segmentation, depth estimation, surface normal estimation, and reinforcement learning. The source code is publicly available at https://github.com/SamsungLabs/MTL .