TLDR: NTKMTL is a new method for Multi-Task Learning (MTL) that tackles the problem of task imbalance, where some tasks learn faster than others. It uses Neural Tangent Kernel (NTK) theory to analyze and balance the convergence speeds of different tasks by adjusting their weights based on NTK eigenvalues. An efficient version, NTKMTL-SR, also provides computational benefits. Experiments show NTKMTL achieves state-of-the-art performance in various multi-task scenarios, including supervised and reinforcement learning.
Multi-Task Learning (MTL) is a powerful approach in artificial intelligence where a single model learns to perform several tasks simultaneously. This method is highly beneficial as it allows tasks to share information and representations, often leading to improved performance on individual tasks and more efficient use of computational resources. MTL has found widespread applications in diverse fields such as computer vision, natural language processing, and robotics.
However, despite its advantages, MTL faces a significant hurdle: task imbalance. This occurs when some tasks dominate the training process, receiving more optimization, while others are neglected, leading to suboptimal performance across the board. Previous research has shown that achieving a more balanced optimization across all tasks is crucial for overall success.
One common strategy to address task imbalance is to balance the convergence speeds of different tasks. Yet, accurately understanding and characterizing how multiple tasks train and converge within a complex MTL system is incredibly challenging. Many existing methods approximate convergence speeds based on simple loss value comparisons, which often fall short because different tasks have vastly different loss scales and ultimate performance goals.
Introducing NTKMTL: A Kernel-Based Solution
To overcome these limitations, a new method called NTKMTL (Neural Tangent Kernel Multi-Task Learning) has been proposed. This approach leverages Neural Tangent Kernel (NTK) theory, a framework that provides deep insights into how deep neural networks learn. In single-task learning, NTK theory explains a phenomenon known as “spectral bias,” where networks tend to learn simpler, low-frequency components of a task faster than complex, high-frequency ones. This concept bears a strong resemblance to task imbalance in MTL.
NTKMTL extends this theory to the multi-task setting by introducing an “extended NTK matrix.” This matrix helps to jointly characterize the training dynamics of all tasks. The core idea is that tasks with larger NTK eigenvalues (which indicate faster learning) can dominate the training, causing imbalance. NTKMTL addresses this by assigning appropriate weights to each task during training. These weights are derived from a spectral analysis of each task’s NTK matrix, effectively balancing their convergence speeds and mitigating task imbalance.
Efficiency with NTKMTL-SR
Recognizing that computing the full NTK matrix can be computationally intensive, the researchers also developed an efficient approximation called NTKMTL-SR (NTKMTL-Shared Representation). This variant takes advantage of the shared parameters common in MTL models. By analyzing the NTK of the shared representation, NTKMTL-SR significantly reduces computational cost, requiring only a single gradient backpropagation per iteration for shared parameters, while still maintaining competitive performance.
Also Read:
- A New Approach to Bayesian Meta-Learning: Neural Variational Dropout Processes
- A Consolidated LoRA Approach for Adaptive Domain Incremental Learning
Demonstrated Performance
Extensive experiments have validated the effectiveness of both NTKMTL and NTKMTL-SR. They have achieved state-of-the-art performance across a wide array of benchmarks, including multi-task supervised learning (on datasets like NYUv2, CityScapes, and CelebA with up to 40 tasks) and multi-task reinforcement learning (on the MT10 environment). For instance, on the challenging NYUv2 dataset, NTKMTL was one of only two methods to consistently outperform single-task learning across all three tasks, demonstrating a truly balanced optimization.
The research highlights that NTKMTL-SR, in particular, offers training speeds comparable to traditional, less sophisticated methods, making it highly practical for real-world applications. This work provides a robust theoretical foundation for understanding and solving task imbalance in multi-task learning. For more technical details, you can refer to the full research paper available here.
