Ensemble Prediction of Task Affinity for Efficient Multi-Task Learning

TL;DR

This paper introduces ETAP, a scalable framework that predicts task affinity in multi-task learning, enabling efficient task grouping by estimating potential performance gains through linear and non-linear models.

Contribution

The paper presents a novel ensemble predictor combining linear affinity scores and non-linear refinements to accurately forecast task group benefits in multi-task learning.

Findings

ETAP outperforms existing baselines in predicting MTL gains.

ETAP enables more effective task grouping, improving overall multi-task learning performance.

The framework is scalable and applicable across diverse datasets.

Abstract

A fundamental problem in multi-task learning (MTL) is identifying groups of tasks that should be learned together. Since training MTL models for all possible combinations of tasks is prohibitively expensive for large task sets, a crucial component of efficient and effective task grouping is predicting whether a group of tasks would benefit from learning together, measured as per-task performance gain over single-task learning. In this paper, we propose ETAP (Ensemble Task Affinity Predictor), a scalable framework that integrates principled and data-driven estimators to predict MTL performance gains. First, we consider the gradient-based updates of shared parameters in an MTL model to measure the affinity between a pair of tasks as the similarity between the parameter updates based on these tasks. This linear estimator, which we call affinity score, naturally extends to estimating affinity within a group of tasks. Second, to refine these estimates, we train predictors that apply non-linear transformations and correct residual errors, capturing complex and non-linear task relationships. We train these predictors on a limited number of task groups for which we obtain ground-truth gain values via multi-task learning for each group. We demonstrate on benchmark datasets that ETAP improves MTL gain prediction and enables more effective task grouping, outperforming state-of-the-art baselines across diverse application domains.