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.
Multi-task learning shows strikingly inconsistent results -- sometimes joint training helps substantially, sometimes it actively harms performance -- yet the field lacks a principled framework for predicting these outcomes. We identify a fundamental but unstated assumption underlying gradient-based task analysis: tasks must share training instances for gradient conflicts to reveal genuine relationships. When tasks are measured on the same inputs, gradient alignment reflects shared mechanistic structure; when measured on disjoint inputs, any apparent signal conflates task relationships with distributional shift. We discover this sample overlap requirement exhibits a sharp phase transition: below 30% overlap, gradient-task correlations are statistically indistinguishable from noise; above 40%, they reliably recover known biological structure. Comprehensive validation across multiple datasets achieves strong correlations and recovers biological pathway organization. Standard benchmarks systematically violate this requirement -- MoleculeNet operates at <5% overlap, TDC at 8-14% -- far below the threshold where gradient analysis becomes meaningful. This provides the first principled explanation for seven years of inconsistent MTL results.
Multi-task learning (MTL) has emerged as a pivotal paradigm in machine learning by leveraging shared structures across multiple related tasks. Despite its empirical success, the development of likelihood-based efficiently solvable algorithms--even for shared linear representations--remains largely underdeveloped, primarily due to the non-convex structure intrinsic to matrix factorization. This paper introduces a first-order algorithm that jointly learns a shared representation and task-specific parameters, with guaranteed efficiency. Notably, it converges in O(1) iterations and attains a near-optimal estimation error of O(dk/(TN)), improving over existing likelihood-based methods by a factor of k, where d, k, T, N denote input dimension, representation dimension, task count, and samples per task, respectively. Our results justify that likelihood-based first-order methods can efficiently solve the MTL problem.
Shihong Ding, Fangyu Du, Cong Fang
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
Multi-Task Learning (MTL) is a foundational machine learning problem that has seen extensive development over the past decade. Recently, various optimization-based MTL approaches have been proposed to learn multiple tasks simultaneously by altering the optimization trajectory. Although these methods strive to de-conflict and re-balance tasks, we empirically identify that their effectiveness is often undermined by an overlooked factor when employing advanced optimizers: the instant-derived gradients play only a marginal role in the actual parameter updates. This discrepancy prevents MTL frameworks from fully releasing its power on learning dynamics. Furthermore, we observe that Muon-a recently emerged advanced optimizer-inherently functions as a multi-task learner, which underscores the critical importance of the gradients used for its orthogonalization. To address these issues, we propose APT (Applicability of advanced oPTimizers), a framework featuring a simple adaptive momentum mechanism designed to balance the strengths between advanced optimizers and MTL. Additionally, we introduce a light direction preservation method to facilitate Muon's orthogonalization. Extensive experiments across four mainstream MTL datasets demonstrate that APT consistently augments existing MTL approaches, yielding substantial performance improvements.
Zhipeng Zhou, Linxiao Cao, Pengcheng Wu +2
Nanyang Technological University · Hong Kong University of Science and Technology (Guangzhou) · School of Artificial Intelligence, Shanghai Jiao Tong University