Independent Component Alignment for Multi-Task Learning
Dmitry Senushkin Nikolay Patakin Arseny Kuznetsov Anton Konushin
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 .