Independent Component Alignment for Multi-Task Learning

TL;DR

This paper introduces Aligned-MTL, a novel multi-task learning optimization method that uses a condition number criterion to improve stability and convergence, allowing explicit control over task trade-offs.

Contribution

The paper proposes a new stability criterion based on the condition number of gradient systems and an optimization approach that aligns gradient components to enhance multi-task learning.

Findings

Consistently improves performance across various MTL benchmarks.

Ensures convergence to a pre-defined task trade-off point.

Provides a theoretically grounded method for stable and controllable MTL optimization.

Abstract

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 .