Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xiaohan Qin, Xiaoxing Wang, Ning Liao, Junchi Yan

TL;DR

This work uses Neural Tangent Kernel theory to address task imbalance in multi-task learning by analyzing an extended NTK matrix that encodes inter-task training dynamics. It introduces NTKMTL, which balances convergence speeds across tasks by weighting task losses according to per-task NTK eigenvalues, and NTKMTL-SR, an efficient shared-representation variant that reduces computational cost while maintaining performance. The methods are validated across diverse MT domains, including dense prediction, multi-task regression, and multi-task reinforcement learning, achieving state-of-the-art or competitive results and demonstrating improved task balance. The approach provides a principled, theoretically grounded mechanism to mitigate task imbalance and offers practical variants for scalable training with many tasks.

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.

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

TL;DR

This work uses Neural Tangent Kernel theory to address task imbalance in multi-task learning by analyzing an extended NTK matrix that encodes inter-task training dynamics. It introduces NTKMTL, which balances convergence speeds across tasks by weighting task losses according to per-task NTK eigenvalues, and NTKMTL-SR, an efficient shared-representation variant that reduces computational cost while maintaining performance. The methods are validated across diverse MT domains, including dense prediction, multi-task regression, and multi-task reinforcement learning, achieving state-of-the-art or competitive results and demonstrating improved task balance. The approach provides a principled, theoretically grounded mechanism to mitigate task imbalance and offers practical variants for scalable training with many tasks.

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.
Paper Structure (22 sections, 3 theorems, 37 equations, 5 figures, 9 tables, 2 algorithms)

This paper contains 22 sections, 3 theorems, 37 equations, 5 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Preliminaries
  5. Extended NTK in Multi-Task Learning
  6. The Proposed NTKMTL
  7. Approximation via Shared Representation
  8. Experiments
  9. Protocols
  10. Multi-Task Supervised Learning
  11. Multi-Task Reinforcement Learning
  12. Conclusion, Limitations and Future Work
  13. Definitions of Notations
  14. Theoretical Analysis
  15. Proof of Theorem \ref{['th:ntk_ode']}
  16. ...and 7 more sections

Key Result

Theorem 3.1

Let $\mathcal{O}(t)=\{f(\theta, x_i)\}_{i=1}^n$ be the outputs of the neural network at time $t$. $\mathbf{x}=\{x_i\}_{i=1}^n$ is the input data, and $\mathbf{y}= \{y_i\}_{i=1}^n$ is the corresponding label, Then $\mathcal{O}(t)$ follows this evolution:

Figures (5)

  • Figure 1: Training time per epoch for various methods on CelebA (40-task) dataset.
  • Figure 2: Visualization for the NTK eigenvalues during training.
  • Figure 3: Ablation study on hyperparameter $n$ on QM9. Each experiment is repeated over 3 random seeds, and the mean and stderr are reported.
  • Figure 4: Training time per epoch for LS and NTKMTL-SR (with different hyperparameter $n$) on QM9.
  • Figure : NTKMTL

Theorems & Definitions (6)

  • Theorem 3.1
  • Theorem 3.2
  • Proposition 3.3
  • proof
  • proof
  • proof