NTKMTL: Mitigating Task Imbalance in Multi-Task Learning from Neural Tangent Kernel Perspective

TL;DR

This paper introduces NTKMTL, a novel multi-task learning method leveraging Neural Tangent Kernel theory to analyze and balance task convergence speeds, effectively mitigating task imbalance and improving performance across various benchmarks.

Contribution

The paper develops an NTK-based analysis framework for MTL and proposes NTKMTL and NTKMTL-SR to balance task convergence, addressing task imbalance effectively.

Findings

Achieves state-of-the-art results on multiple benchmarks.

Effectively mitigates task imbalance in MTL.

Maintains competitive performance with improved training efficiency.

Abstract

Multi-Task Learning (MTL) enables a single model to learn multiple tasks simultaneously, leveraging knowledge transfer among tasks for enhanced generalization, and has been widely applied across various domains. However, task imbalance remains a major challenge in MTL. Although balancing the convergence speeds of different tasks is an effective approach to address this issue, it is highly challenging to accurately characterize the training dynamics and convergence speeds of multiple tasks within the complex MTL system. To this end, we attempt to analyze the training dynamics in MTL by leveraging Neural Tangent Kernel (NTK) theory and propose a new MTL method, NTKMTL. Specifically, we introduce an extended NTK matrix for MTL and adopt spectral analysis to balance the convergence speeds of multiple tasks, thereby mitigating task imbalance. Based on the approximation via shared representation, we further propose NTKMTL-SR, achieving training efficiency while maintaining competitive performance. Extensive experiments demonstrate that our methods achieve state-of-the-art performance across a wide range of benchmarks, including both multi-task supervised learning and multi-task reinforcement learning. Source code is available at https://github.com/jianke0604/NTKMTL.