This repository contains the code and data for the paper “Ensemble Prediction of Task Affinity for Efficient Multi-Task Learning”, published at the 14th International Conference on Learning Representations (ICLR 2026).
arXiv version: https://arxiv.org/abs/2602.18591
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.
Task affinity computed from a baseline MTL model and ground-truth MTL gains are fed into an ensemble framework. Non-linear transformations yield initial predictions, which are later refined by residual correction through regularized regression.
Prediction performance (
| Method | Computational Cost (# of Training Groups) | CelebA | ETTm1 | Chemical | Ridership |
|---|---|---|---|---|---|
| TAG (Fifty et al. 2021) | n.a. | 0.10 ± 0.0 | 0.47 ± 0.0 | 0.05 ± 0.1 | 0.15 ± 0.1 |
| MTGNet (Song et al. 2022) | 5 | 0.10 ± 0.2 | 0.43 ± 0.1 | 0.22 ± 0.1 | 0.43 ± 0.1 |
| 10 | 0.22 ± 0.1 | 0.54 ± 0.2 | 0.34 ± 0.2 | 0.61 ± 0.0 | |
| ETAP | 5 | 0.41 ± 0.2 | 0.77 ± 0.1 | 0.40 ± 0.1 | 0.68 ± 0.1 |
| 10 | 0.45 ± 0.1 | 0.84 ± 0.0 | 0.50 ± 0.1 | 0.74 ± 0.0 |
Correlation between ground-truth and predicted MTL gains for groups (higher values are better).