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

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.