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.
Figures & tables
Figure 1: Visualizing ETAP: From White-box Task Affinity Scores to Data-driven Ensembled MTL Gain Predictions. 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.
Dataset
Domain, #Tasks ( ∣T∣ )
MTL Architecture
Notes
CelebA ( Liu et al., 2015b )
Facial attributes, 9
ResNet-18
T from TAG Fifty et al. (2021)
ETTm1 ( Wu et al., 2021 )
Temperature time-series, 7
Autoformer
T from MTGNet Song et al. (2022)
Chemical ( Jacob et al., 2008 )
Molecule classification, 10
FFNN encoder
T⊂35 with diverse sizes
Ridership ( Zulqarnain et al., 2023 )
Transit demand, 10
LSTM + FFNN
Real-world ridership data
Table 1: Multi-Task Learning Datasets and MTL Model Architectures
Affinity Scoring Method
CelebA
ETTm1
Chemical
Ridership
TAG ( Fifty et al., 2021 )
0.16 ± 0.00
0.43 ± 0.02
0.33 ± 0.19
0.10 ± 0.07
GRAD-TAE ( Li et al., 2024 )
0.09 ± 0.06
-0.27 ± 0.02
0.28 ± 0.06
-0.22 ± 0.19
ETAP
0.32 ± 0.00
0.47 ± 0.00
0.40 ± 0.03
0.36 ± 0.03
Table 2: Correlation between pairwise affinity scores and ground-truth MTL gains (higher values are better).
Method
Computational
Correlation with Ground Truth
Cost
CelebA
ETTm1
Chemical
Ridership
TAG ( Fifty et al., 2021 )
0.10±0.0
0.47±0.0
0.05±0.1
0.15±0.1
MTGNet ( Song et al., 2022 )
∣Gtrain∣ = 5
0.10±0.2
0.43±0.1
0.22±0.1
0.43±0.1
∣Gtrain∣ = 10
0.22±0.1
0.54±0.2
0.34±0.2
0.61±0.0
ETAP
∣Gtrain∣ = 5
0.41±0.2
0.77±0.1
0.40±0.1
0.68±0.1
∣Gtrain∣ = 10
0.45±0.1
0.84±0.0
0.50±0.1
0.74±0.0
Table 3: Correlation between ground-truth and predicted MTL gains for groups (higher values are better).
Figure 2: Prediction performance ( R2 ) vs. computational cost ( ∣Gtrain∣ ) for data-driven predictors, MTGNet and ETAP . MTGNet suffers from instability, with means outside the IQR due to outliers, whereas ETAP maintains consistency.
Dataset
#Groups
TAG
GRAD-TAE
MTGNet ( ∣Gtrain∣ )
ETAP ( ∣Gtrain∣ )
Optimal
CelebA ( ∣Gtrain∣ = 10)
2
49.67 ± 0.00
50.78 ± 0.59
50.62 ± 1.20
49.92 ± 0.23
49.27
3
50.22 ± 0.00
50.78 ± 1.93
50.31 ± 0.63
49.61 ± 0.30
48.63
4
49.94 ± 0.00
52.83 ± 1.70
50.27 ± 0.60
49.41 ± 0.25
48.38
ETTm1 ( ∣Gtrain∣ = 10)
2
4.08 ± 0.09
4.08 ± 0.06
4.04 ± 0.06
4.02 ± 0.05
3.98
3
3.96 ± 0.03
4.04 ± 0.06
3.96 ± 0.03
3.93 ± 0.07
3.83
4
3.90 ± 0.02
4.01 ± 0.04
3.92 ± 0.06
3.89 ± 0.09
3.82
Table 4: MTL performance (measured as total loss ( ± std); lower is better) with tasks grouped based on various prediction methods. Optimal values are computed using an exhaustive search for CelebA and ETTm1, and using a search over a representative 50% of all task combinations for Chemical and Ridership.
Dataset
∣Gtrain∣
Correlation with Ground Truth
R2
Ayman et al. (2023)
ETAP
Ayman et al. (2023)
ETAP
Chemical
5
0.33 ± 0.69
0.74 ± 0.04
–0.84 ± 2.12
0.15 ± 0.13
10
0.76 ± 0.05
0.75 ± 0.05
0.48 ± 0.12
0.36 ± 0.05
20
0.78 ± 0.06
0.77 ± 0.04
0.56 ± 0.12
0.39 ± 0.19
Ridership
5
–0.01 ± 0.13
0.44 ± 0.06
–0.18 ± 0.67
0.16 ± 0.08
10
0.05 ± 0.14
0.42 ± 0.07
–0.81 ± 0.57
0.12 ± 0.09
Table 5: Correlation and R2 between ground-truth and predicted total MTL gains for groups (higher values are better).
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Symbol
Description
T
Set of all tasks
n
Number of tasks ( n=∣T∣ )
Lt
Loss of task t∈T
LT
Total loss across all tasks ( LT=∑t∈TLt )
G
A task-group (i.e., a set of tasks G⊆T )
G
A collection of task groups ( G={G1,G2,⋯} )
Appendix
Table 6: List of Notations
Figure 3: Comparison of prediction performance ( R2 ) between the initial predictions ( without residual adjustment ) and the final predictions ( with residual adjustment ) across four datasets. The residual adjustment step improves accuracy, particularly on CelebA, Chemical, and Ridership, with a slight improvement observed on ETTm1.
Dataset
∣Gtrain∣
Correlation (Predicted vs. Actual)
Initial
Final
CelebA
5
0.27 ± 0.15
0.41 ± 0.16
10
0.31 ± 0.13
0.45 ± 0.13
ETTm1
5
0.75 ± 0.15
0.77 ± 0.12
10
0.84 ± 0.02
0.85 ± 0.02
Chemical
5
0.32 ± 0.06
0.40 ± 0.07
Appendix
Table 7: Correlation with ground-truth MTL gains across different computational costs ( ∣Gtrain∣ ) for the predictor on all benchmarks (higher values are better).
Dataset
R2 (Affine)
R2 (B-spline)
Corr. (Affine)
Corr. (B-spline)
CelebA
0.13 ± 0.00
0.13 ± 0.01
0.37 ± 0.00
0.37 ± 0.01
ETTm1
0.18 ± 0.02
0.68 ± 0.05
0.45 ± 0.00
0.84 ± 0.02
Chemical
-0.06 ± 0.11
-0.04 ± 0.11
0.35 ± 0.00
0.37 ± 0.04
Ridership
0.35 ± 0.03
0.42 ± 0.04
0.62 ± 0.00
0.67 ± 0.02
Appendix
Table 8: Effect of non-linear transformation on task-affinity scores: comparison of affine vs. B-spline mapping in terms of R2 and correlation with ground-truth MTL gains.
Figure 4: Runtime for CelebA and ETTm1 datasets (measured in hours) and the Chemical dataset (measured in minutes). Running times include task-affinity scores calculations and complete multi-task learning training for all tasks together in the benchmark using baseline MTL models.
Dataset
Training
Surrogate Model
ETAP
Cost ( ∣Gtrain∣ )
Corr.
F1
Corr.
F1
CelebA
32
0.49
0.61
0.57
0.70
ETTm1
21
0.35
0.18
0.64
0.31
Chemical
37
0.57
0.92
0.62
0.94
Ridership
35
0.67
0.81
0.77
0.85
Appendix
Table 9: ETAP vs. Linear Surrogate Model Li et al. (2023) for identifying positive vs. negative transfer.
Dataset
Naive MTL
PCGrad
ETAP (SPLITS)
2
3
4
CelebA
50.02 ± 0.4
49.99 ± 0.5
49.92 ± 0.2
49.61 ± 0.3
49.41 ± 0.3
ETTm1
4.20 ± 0.0
4.19 ± 0.0
4.02 ± 0.1
3.93 ± 0.1
3.89 ± 0.0
Chemical
4.83 ± 0.2
6.20 ± 0.2
4.67 ± 0.0
4.74 ± 0.2
4.83 ± 0.2
Ridership
18.95 ± 0.7
18.84 ± 0.3
17.77 ± 0.3
17.83 ± 0.3
17.59 ± 0.2
Appendix
Table 10: Comparison of multi-task learning performance (measured in loss; lower is better) across different methods (Naive MTL, PCGrad, and our proposed ETAP method with 2, 3, and 4 task-group splits). Results are reported as mean ± standard deviation across multiple runs.
Dataset
Computational Cost
Sample Size
Wilcoxon p -value
CelebA
5
10
0.0020
10
10
0.0039
20
10
0.0039
30
10
0.0020
40
10
0.0039
50
10
0.0020
Appendix
Table 11: Wilcoxon signed-rank test results comparing MTGNet vs. ETAP across datasets and training sample sizes. Bold p -values indicate statistical significance at p<0.05 .
Dataset
Sample Size
Cost ( ETAP )
Wilcoxon p -values
CelebA
7
5
0.0313
ETTm1
7
5
0.0313
Chemical
7
5
0.0156
Ridership
7
10
0.0078
Appendix
Table 12: Wilcoxon signed-rank test p -values comparing TAG vs. ETAP across benchmarks.
Figure 5: Comparison of prediction performance ( R2 ) for different approaches applied in the first step of the ensemble model ( fnon-linear ) on the CelebA, ETTm1, Chemical, and Ridership datasets. These performances are reported after applying a residual correction step ( fresidual ), where regularized ridge regression is used for residual prediction.
Figure 6: Prediction performance ( R2 ) for different approaches applied in the residual learning step of the ensemble model ( fresidual ) on the CelebA, ETTm1, Chemical, and Ridership datasets. Results are reported for both Ridge and XGB-based residual predictors applied to the initial prediction methods: BSpline, KNN, and Random Forest (RF).
Figure 7: Evolution of per-step affinity score zti→tjk during training between a source task i ( 5_o_Clock_Shadow ) and each target task j ( Black_Hair , Blond_Hair , Brown_Hair , Goatee , Mustache , No_Beard , Rosy_Cheeks , Wearing_Hat ) in dataset CelebA.
State Key Lab of General AI, School of Intelligence Science and Technology, Peking University, China · School of Mathematical Sciences, Peking University, China · Institute for Artificial Intelligence, Peking University, China
Nanyang Technological University · Hong Kong University of Science and Technology (Guangzhou) · School of Artificial Intelligence, Shanghai Jiao Tong University