Exploring Task Affinities through NTK Alignment and Early Training Dynamics in Multi-Task Learning

Yoann Morello, Emilie Gregoire, Sam Verboven

Keywords: MTL, Multi-Task Learning, NTK, Neural tangent Kernel, Alignment

TL;DR: After expressing existing gradient-based measures in terms of the Neural Tangent Kernel (NTK), we explore methods to average out the initial stochasticity, allowing us to compute task similarities without requiring full training.

Abstract: Multi-task learning (MTL) aims to leverage shared representations among tasks to improve generalization and training efficiency. However, the potential for negative transfer poses a significant challenge. This work explores modifications to gradient-based measures for task similarity to identify effective task groupings early in training. We highlight key connections between existing measures through the Neural Tangent Kernel (NTK). Our analysis reveals that computing these measures during the initial training stages, averaged over multiple runs, provides a robust estimation of task affinities. We demonstrate the method's effectiveness on synthetic data, capturing both linear and non-linear relationships, and suggest its potential applicability to more complex datasets.

Is Neurips Submission: No

Submission Number: 45