# Ensemble Prediction of Task Affinity for Efficient Multi-Task Learning

arXiv: https://arxiv.org/abs/2602.18591 · HTML: https://mlsanity.com/paper/2602.18591/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 40 references that are in the map.
- First posted on arXiv 20 Feb 2026.

## Builds on (most cited first, 23 of 23)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Autoformer: Decomposition Transformers with Auto-Correlation for Long-Term Series Forecasting](https://mlsanity.com/paper/2106.13008/) (2021)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [Taskonomy: Disentangling Task Transfer Learning](https://mlsanity.com/paper/1804.08328/) (2018)
- [Which Tasks Should Be Learned Together in Multi-task Learning?](https://mlsanity.com/paper/1905.07553/) (2019)
- [Multi-Task Learning with Deep Neural Networks: A Survey](https://mlsanity.com/paper/2009.09796/) (2020)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [Learning to Branch for Multi-Task Learning](https://mlsanity.com/paper/2006.01895/) (2020)
- [Multi-task Sequence to Sequence Learning](https://mlsanity.com/paper/1511.06114/) (2015)
- [Representation Similarity Analysis for Efficient Task taxonomy & Transfer Learning](https://mlsanity.com/paper/1904.11740/) (2019)
- [Identifying beneficial task relations for multi-task learning in deep neural networks](https://mlsanity.com/paper/1702.08303/) (2017)
- [Heterogeneous Face Attribute Estimation: A Deep Multi-Task Learning Approach](https://mlsanity.com/paper/1706.00906/) (2017)
- [The Information Complexity of Learning Tasks, their Structure and their Distance](https://mlsanity.com/paper/1904.03292/) (2019)
- [Zero-Shot Task Transfer](https://mlsanity.com/paper/1903.01092/) (2019)
- [Identification of Negative Transfers in Multitask Learning Using Surrogate Models](https://mlsanity.com/paper/2303.14582/) (2023)
- [Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity](https://mlsanity.com/paper/2409.06091/) (2024)
- [Task Grouping for Automated Multi-Task Machine Learning via Task Affinity Prediction](https://mlsanity.com/paper/2310.16241/) (2023)

## Abstract (arXiv)

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.