Ensemble Prediction of Task Affinity for Efficient Multi-Task Learning
Authors: Afiya Ayman
Ayan Mukhopadhyay
Aron Laszka
Venue: ICLR 2026 Poster (2026)
Keywords: Machine Learning
Multi-Task Learning
Task Affinity
AutoML
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: We introduce an ensemble approach for predicting performance gains from multi-task learning, bridging the gap between white-box and data-driven approaches for task grouping.

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.

Published as a conference paper at ICLR 2026
ENSEMBLE PREDICTION OF TASK AFFINITY FOR EFFICIENT MULTI-TAST LEARNING
Afiya Ayman
Ayan Mukhopadhyay College of William & Mary
Aron Laszka
Pennsylvania State University
Pennsylvania State University
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 possi- ble combinations of tasks is prohibitively expensive for large task sets, crucial component of efficient and effective task grouping predicting whether group of tasks would benefit from learning together, measured per-task performance gain over single-task learning. In this paper, we propose ETAP (Ensemble Task Affin- ity Predictor), scalable framework that integrates principled and data-driven timators to predict MTL performance gains. First, we consider the gradient- traseent-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 estimat- ing affinity within group of tasks. Second, to refine these estimates, we train predictors that apply non-linear transformations and correct residual errors, cap- turing 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--at baselines across diverse application domains.
INTRODUCTION
Multi-task learning (MTL. is approach to learn multiple machine-learning tasks together to im- prove generalization performance and reduce inference time Caruana) (1997): Standley et al. (2020). improves the performance of individual tasks by leveraging related information in other tasks through inductive transfer, i.c., the learning architecture is designed to introduce inductive bias fa- voring hypothesis classes that perform well across tasks Caruana ( (1997). MTL has proven to be effective in multiple domains, including computer vision Zhang et al. (2014); Liu et al. (2015a), health informatics Puniyani et al.) 2010);(Zhou et al. (2011), and natural language processing Luong et al. (0015). However, simply learning set of tasks together does not ensure better performance; fact, performance degrade compared performing single-task learning (STL) for each task, a phenomenon known as negative transfer Rossenstei et al. (205).. This observation highlights critical challenge: finding groups of tasks that maximize the performance gain from MTL.
A naive strategy to find well-performing task groups is to run multi-task learning for every possible group; however, this is computationally infeasible in most cases. So, recent research focused on predicting MTL gains, that is, predicting the improvement in performance when a group of tasks are trained together versus independently. Accurate prediction of MTL gains enables finding optimal (or near optimal) task groups, leading to more efficient and effective MTI.
common approach for predicting MTL gains-whit we refer to as the data-driven or black-box approach-is to first obtain ground-truth MTL gain values by performing multi-task learning for several task groups and measuring performance gains, and to then predict gains for other groups based on these ground-truth values. For example, Ffigher-Order Affinities (HOA) relies ground- truth MTL gains for all task pairs (Standley et al. 2020). Other methods such as MTGNet (Song et al. 2022) and Linear Surrogate (Li al. (023) learn meta-models from large numbers of sam- pled training groups, relying on ground- -ruth MTL gains from extensive multi-task learning. Despite methodological differences, these approaches share a key limitation: they computationally

---

Published as a conference paper at ICLR 2026
pensive and often require large numbers of training groups to provide accurate, robust predictions. This robsstobsss tradeoff becomes prohibitive as the number of tasks grows.
In contrast, white-box approaches aim to predict MTL performance by analyzing training synamics--peciffcallyy how gradient updates from one task affect the losses of other tasks. A rep- resentative example is Task-Affinity Grouping (TAG), which estimates inter-task affinity by training a single MTL model while performing additional hypothetical training steps to evaluate the im- pacts of task-specific parameter updates (Fifty et al. (021). While theoretically grounded, TAG is limited to only pairwise estimates, which are highly correlated to MTL gains but are not a direct predictors. TAG also struggles model higher-order interactions, limiting its accuracy in realistic multi-task settings. Further, it incurs significant computational overhead due to the additional for- ward/backward passes, making its overall cost comparable to training many separate MTL. models.
Contributions:
This paper introduces ETAP (Ensemble Task Affintt Prediction), a novel frame- work that combines the theoretical strengths of white-box affinity estimation with the generalization ability of data-driven models to predict MTL gains, which is formally defined as the reduction in loss due to learning a group of tasks jointly versus independently. ETAP addresses the limitations of prior approaches for prediction through a two-part integration. First, we propose an improved white-box technique to compute task-affinity scores from gradient interactions during single MTL training with minimal computational overhead. Compared to prior methods such as TAG, our approach is more efficient and avoids auxiliary training steps, while producing reliable affinity estimates. How- ever, leveraging these affinity scores to accurately predict higher-order MTL gains poses a challenge: task dependencies
often non-linear and cannot be fully captured by pairwise relationships alone. To address this, ETAP employs a two-stage ensemble prediction framework. The first stage learns a non-linear mapping from affinity scores predicted group-level MTL gains. The second stage uses a residual model to correct systematic prediction errors by learning from mismatches between pre- dictions and ground-truth gains. This design improves robustness and enables accurate predictions with far fewer training groups than prior dataen methods.
Our MTL-gain predictor may be integrated as an objective into a wide range of search algorithms to find optimal task groups for MTL (c.g., it can be integrated into branch-and -ntth---n--boudd algorithms from Standley et al. (2020) and Fifty et al. (202T) or beam-search algorithm from Soog et al. (2022). In our experimental evaluation, we use the efficient aranch-and- aandbound search algorithm from Standley jet al. (2020) to find task groups with our proposed predictor and with baseline predictors. Our ex- perimental results demonstrate that ETAP outperforms existing data-driven and white-box methods in terms of both accuracy and computational efficiency across diverse MTL benchmarks. By strate- gically integrating white-box affinity measurements with data-driven refinement, ETAP provides better tradeoff between efficiency and accuracy, addressing the limitations of existing approaches.
2 PROBLEM DEFINITION
We consider a set of n tasks, denoted T= (11, 1..... ,I,) n 177. A task group, denoted G, is a subset of T (i.c.. G T) and can have any number of tasks from 1 to n. An arbitrary set of such groups is denoted as G (i.e., G C 27).
=
A multi-task learning (MTL) algorithm jointly trains on the data of all tasks in a group G to build shared feature representation. In many architectures, task-specific decoders transform this shared representation to produce task-specific predictions [Ruder (2017): [Crawshaw (2020). The MTL algorithm trains multi-task model for tasks in group G with parameters B The loss for a specific task I, denoted by I,,, measures the model's prediction performance on that task. During training, each task updates the model parameters &G. based on its own loss function.
MTL gain is our formal metric for evaluating task relationships in the context of MTL performance. MTL gain measures the performance improvement for a task (i.e., reduction in loss L,) when trained using MTL compared to training the task independently using a single-task learning (STL) model. Specifically, for task t G. the relative MTL gain, denoted by VG-+t the reduction loss for task t when tasks in group G' are trained together. Formally, the relative MTL gain for task t E G is
Le(X, (..)) -Le(X, (
VG
Li(X' 0%
(1)
2

---

Published as a conference paper at ICLR 2026
Afiiniyy Scores (White
MTL Gain Prediction
MTL step,
Affintty Scores
w. ground (Shum)
Predicted
all T
perage
Non-lincar transformation (B-Splines)
MIL
L. O% setieke,
Gains
gasd
For taskes f,.6y
Regularized group caasedd
set of training
groups Sux
G:
MTL
Moi
regression
ground-truth gains (Sa))
Figure 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 trunh MTL gains are fed into an ensemble framework. Non-linear transformations yield initial predictions, which are later refined by residual correction through regularized regression.
where X' is the test dataset, (X(, (%n) is the loss of task t when trained independently to obtain 0 MIL) the loss of task & when trained using MTL with group to obtain parameters km. A positive value of yG-+ indicates performance gain (i.e., loss reduction), while a negative value indicates performance degradation.
parameters On and Le(X,
Given task-set T, the goal is to identify optimal task groups that maximize MTL performance. This involves minimizing the total loss L. across tasks in or, equivalently, maximizing the overall MTL gains for each task. Formally, we aim to select collection of task groups G2... J. under a constraint on the number groups [ s B, to maximize the sum of MTL gains for each task across all groups:
Maximize: Σ LT
max max VG-+ GCT
subject to
G S B.
This optimization problem is NP-hard due the exponential number of possible task groups. Since an exhaustive search is infeasible, we focus on efficiently predicting MTL gains for task groups. The prediction of MTL gains involves learning function f that estimates UG-,r for task G by minimizing prediction error (#G-. ) Our goal is to minimize both the error in MTL gain predictions as well as the computational cost of learning f. enabling efficient and effective task grouping for improved MTL performance.
ETAP: ENSEMBLE TASK-AFFINT PREDICTION
Our proposed approach, ETAP, is unified framework for predicting MTL gains. ETAP predicts which groups tasks will benefit from joint training (benefit measured as performance gain over STL) and use these predictions form effective task groups, avoiding exhaustive search.
ETAP combines principled white-box analysis with data-driven modeling. It first computes gradient- based iask affinity scores by analyzing how updates from task affect another's loss, providing model-aware estimates of linear relationships. These scores form the foundation of gain prediction. Next, ETAP employs data-driven models refine these estimates, learning the complex, non-linear relationships between predicted and actual MTL gains. By integrating these complementary meth- ods, ETAP achieves accurate gain prediction with limited supervision. Finally, using the predictions made by ETAP, we select a near-optimal set of task groups that maximizes total MTL gain while balancing task coverage and performance under budget constraint on the number of groups.
3.1 TASK-AFFINI SCORE VIA WHITE-BOX ANALYSIS
Our goal is to efficiently estimate the affinity between pairs of tasks, capturing how their perfor- mance improves when trained together using multi-task learning. During training, tasks implicitly exchange information through successive gradient updates to the shared parameters. We collect task-specific losses, gradients, and parameter updates each training step to compute pairwise task- affinity scores. These affinity scores correlate well with ground truth MTL gains (as demonstrated in Section 4.2 and serve robust and computationally inexpensive baseline predictor.
3

---

Published as a conference paper at ICLR 2026
3.1.1 TASK LOSS AND GRADIENTS
We begin by training a single MTL model for group with all tasks in T. The MTL model has a shared encoder with parameters 0, (a for shared) and task-specific decoders for each task where parameters for task & are represented as .t.. In gradient step k. for a batch of examples X the loss for task 4 is T, Lr, (x** 0% 0% In each step, we compute the gradient of each task's loss L, with respect to the shared parameters, 0,, which measures how a small change 0g affects task loss Lt, using the following:
(3) These gradients encode the relationship between shared parameters and task losses, capturing task- specific sensitivities changes in O,. Note that these gradients are already computed as part of the standard backpropagation process during MTL training, so no additional computation is required.
V 0e Vo - Voy (2**, O², 0%,)
3.1.2 SHARED PARAMETER UPDATES
Our approach assumes a standard MTL setup with shared encoder parameters updated via gradient descent, optionally with momentum. While not specific to our method, we leverage this common MTL optimization behavior to quantify task interactions. Understanding how tasks interact through updates shared parameters is critical, these updates reveal how tasks impact each other's per- formance, providing key insights for estimating task affinities. Let gk++1 denote the updated shared elt, parameters 0g after gradient-step & taken using only task I,'s loss. If momentum is used during train- ing, the update incorporates previous step's velocity H= and momentum coefficient / Sutskever et al. (203)). The corresponding parameter update Agk+1 with learning rate r) is:
elt;
alts A0++1 = 30 V
(4)
The velocity term is updated as v - I Without momentum, Equation
becomes
A s/b
= V og L.
3.1.3 AFFINITY SCORE
We calculate a task-affinity score for each training step, which we later aggregate to estimate the overall relationship between pairs of tasks. To calculate the step-level task-affinity score between tasks 4 and ly at training step k, we measure how the gradient update for task 4 affects the loss for task ly. Specifically, we take the alignment between the gradient for task 1y and the parameter updates A0, based on task f,, normalized by task tj loss:
V or 1
[oVor 1R. - B
I,
Here, 2f,-,,, is the affinity score from task t; to task ty at step positive value means that task i, 's update to the shared parameters 0, reduces the loss of lj. suggesting a beneficial transfer, while negative value suggests interference. We aggregate these scores over all K training steps obtain a stable estimate of pairwise affinity:
K 2 T4 K
(6)
Time-averaging the affinity score =f, over K training steps is crucial to providing robust esti- mate of pairwise affinity since the per-step affinity score - varies significantly during training (see Figure in Section[F an illustration). The time-averaged pairwise task-affinity scores ह, +t, provides a very stable estimate, which exhibits low variance between training runs (see Table 2
After computing pairwise task-affinity scores ttg we derive group-level scores for task groups with three or more tasks. For task t; E G, the affinity score EG is the average of pairwise affinities from all other tasks in G to
Ec
(7)
G
tt.))(tt)
4

---

Published as a conference paper at ICLR 2026
Here, 2t, is the pairwise affinity score from task t, to task t, We collect these scores for all tasks within a group G into a set &G = { : e G}. In the absence of momentum, due to the linearity of gradient-based measurements, averaging pairwise affinity scores across group is equiv- alent to calculating the score based on the group's ' averaged gradients. This enables simplifying the calculation without sacrificing accuracy and serves as efficient way handle large task groups.
While these groupwise affinity scores ZG-+ offer stable, low-variance estimates due to the averaging over task pairs and training steps, they have two fundamental limitations. First, affinity scores reflect gradient correlations, which may differ in scale from MTL gains. Second, affinity-based group predictions are effectively linear approximations, which tend to have low variance due to averaging but high bias due to oversimplification, overlooking complex, non-linear task interactions.
3.2 DATA-DRIVER ENSEMBLE PREDICTION
To address the limitations of affinity scores, we introduce a two-stage ensemble framework for predicting MTL gains using a combination of white-box affinity scores and data-driven methods. Building on groupwise affinity scores &G-+. which provide first-order approximations of group's s potential for positive transfers, we apply (1) an affinity-to-gain mapping that trained translate affinity scores into predicted MTL gains, and (2) a residual correction that refines these predictions by modeling residual errors between these initial predictions and observed ground-truth truth MTL gains, further reducing bias and capturing higher-order task interactions. For both stages, we use the same training set of task groups Grain with measured ground-truth MTL gains. By integrating white-box and data-driven methods, our framework balances the tradeoff between bias and variance, providing accurate predictions of MTL gain with less training data than existing data-driven methods.
3.2.1 AFFINITY SCORES TO MTL GAINS WITH NON-LINEAR MAPPING
Since affinity scores differ in scale from MTL gain values, we introduce a non-linear transformation that maps affinity scores zg predicted MTL gains gaft Besides bridging the gap between affinity scores and MTL gains-twe two correlated but distinct metrics- the non-linearity of this transformation also enables capturing higher-order task dependencies, which groupwise affinity scores ignore due to their linearity. We learn this transformation from ground-truth MTL gains in training set G using B-spline basis expansion. B-splines are piecewise polynomial functions with local support Boor (1978); Hastie et al. (2009), which we use to map affinity scores to digeer-dimensional feature space. Specifically, for groupwise affinity score 2G-st for task E G, the basis expansion is:
& (=C--<) IN; (=G-x))
(8)
Ni(-) represents the i-th B-spline basis function, and M denotes the total number of basis functions, depending on the chosen spline degree and the number of knots. After applying the B-spline ex- pansion, we fit regularized linear regression model to these higher-dimennsonall features, balancing model complexity and generalization to control for potential overfitting. Note that while the regres- sion itself is linear, the complete transformation from groupwise affinity scores to predicted MTI. gains is non-lincar. Formally, the regression model fnon-lincar maps affinity scores #G. expanded by B-splines o, using regression coefficients waff to predicted MTL gains UG for task group G::
ye fneo-lincar (ø (zc) (Watt)
Both the hyperparameters of the B-spline expansion & (i.e., spline degree and knot parameters) and the regression parameters War are optimized using cross-validation a small training set of task groups Guain- ensuring that the model generalizes well for inference on unseen task groups. We also explore alternative approaches for this non-linear transformation, which we discuss in Section B
This first stage of our ensemble prediction framework bridges the gap between affinity scores and MTL gain values, which correlated but may be on significantly different scales, and due to the non-linearity of the learned transformation, captures some of the higher-order task interactions.
3.2.2 DATA-DRIVEN RESIDUAL PREDICTION FOR IMPROVED ACCURACY
The second stage of our prediction framework improves upon the accuracy of the initial, first-stage predictions jolt by learning to predict and reduce residual errors. The motivation for this second
5

---

Published as a conference paper at ICLR 2026
stage is the need to capture higher-order interactions that pertain particular tasks, which our first- stage predictor faen-linar cannot model. Therefore, we use our training set Guain to learn to predict residual enrors based on which particular tasks included in group. Compared to purely black- box approaches like MTGNet, our residual predictor requires significantly less data since it learns to refine strong initial predictions instead of learning predict MTL gains from scratch.
The residual prediction error for group G is eaff YG jfft where YG denotes the ground-truth MTL gains, i.e., the residual error the discrepancy between the ground-truth MTL gains and the first-stage predictions. Our goal is to learn a model that predicts eG based on which particular tasks are included in group G. To this end, we train supervised model fresidual on training set Grain with labels eg, representing a group G S 'T as a multi-hot encoded vector uG. That is, we represent group G' as a zero-one vector ug E {0, 13.77, where the element of ug corresponding to task E 'T if G, and i & G. Our model fresidual maps UG to a predicted residual, minimizing the loss between predicted and actual residuals for training groups Giain
G
Whes
argmin
Σ
(Jresidual (UGG Whes)
all 2 eG
(10)
GESmmi
where whes denotes the parameters of our model. We implement fresidual as ridge regression model:
Wres
argmin
Σ
(ee - =dWwes) + Al|Wnes!|
(11)
Was
GENuin
a regularization hyperparameter that prevents overfitting by controlling the magnitude of the model coefficients wres, which is particularly important when dealing with high-dimensional input data such as multi-hot encoding of task groups. By penalizing large coefficients, we ensure robust predictions. We can tune the value of hyperparameter À using cross-validation on training Gunin.
We obtain the final MTL gain predictions for any group G by combining the initial, first-stage predictions galf with the residual corrections predicted by fresiduat:
DG
UG + fresidual (uG; Whes).
(12)
This refinement step enables capture the impact of particular tasks MTL gains in group, improving upon the first-stage predictions ب is important to note that learning to predict these residual errors is easier than learning to predict MTL gains from scratch.
3.2.3 GROUP SELECTION BASED ON MTL GAIN PREDICTIONS
The selection of optimal task groups is constrained search problem, which involves choosing set of task groups that cover all n tasks to maximize overall MTL gain (measured by, e.g., average MTL gain over all tasks) while limiting the number of selected groups to budget B. While exhaustive or recursive search over the 2 n increases. Prior work has shown that despite this problem being generally NP-hard (through reduction from Set-Cover), it can be tackled using branch-and-ae -nouudnnd- strategies or binary integer programming solvers (Standley al.) al.| 2020 Fifty al. 2021 Zamir al.) 2018). So, following prior work, we also employ branch-and sranch-and-bound bound algorithm to optimize the selection of B task groups, maximizing that the total predicted MTL gain across all tasks during inference.
combinations is feasible for small n, it quickly becomes infeasible
EXPERIMENTS AND RESULTS
To evaluate ETAP experimentally, we apply to four diverse multi-task datasets (e.g., computer vision and time series) and compare it to various bascline approaches for predicting MTL gains. Our software implementation and the novel ridership dataset are available as supplementary material under open-source licenses.
4.1 EXPERIMENTAL SETUP
Dataset and MTL Architecture We evaluate ETAP on three diverse, widely used multi-task learning datasets, spanning vision (CelebA), time series (ETTm1), and molecular classification

---

Published as a conference paper at ICLR 2026
Table 1: Multi Task Learning Datasets and MTL Model Architectures
Dataset
Domain, #Tasks ([T)
MTL Architecture
Notes
CelebA ct al
2015b
Facial attributes. 9
ResNet-18 Autoformer
T from TAG Filly et al. 2021
ETTml al
2021
Temperature time-scries, 7
T from MTGNet Song et al, 2022
Chemical el al
Molecule classification,
FFNN encoder
T with diverse
Ridership
Zulqaraiin
2023
Transit demand, 10
LSTM FENN
Real-world ridership data
(Chemical), and on a novel, real-world transit Ridership dataset- each with distinct task structures and MTL architectures that selected to align with prior task-grouping studies. Table [ summa- rizes the datasets and MTL model architectures; we provide additional details Section B.1
Hyperparameter Tuning and Evaluation Setting All ETAP hyperparameters tuned via cross-validation on a small set of training groups to balance accuracy and efficiency. The initial prediction stage uses B-spline expansion affinity scores followed by regularized linear regression, with spline degree, knot count, and regularization strength optimized (details are in Section B.2) Specifically, we tune the spline degree (E 22, 6)), the number of knots (up to VIguain ), and the ridge- regularization strength (a E [0.001, 1.0).. The residual correction stage applies ridge regression with cross-vall pross-validated regularization. Alternative models (e.g., KNN, Random Forest) are discussed in Section E We compare ETAP with baselines in terms of MTL gain prediction and task group se- lection. All gain predictors trained using randomized task subsets with different task groups, and evaluation is performed on a fixed hold-out test set, shared across all methods to ensure consistency (130 groups for CelebA, 40 for ETTm1, 230 for Chemical, and 200 for Ridership). For group selec- tion evaluation, all possible combinations are considered for CelebA and ETTml, while 50% of all combinations are considered for Chemical and Ridership datasets due to scalability constraints. To ensure robustness and reduce randomness, all training and evaluation are repeated (10x for gain pre- diction, 6x for group selection) with different random seeds. We report the average outcomes along with their respective confidence intervals. We evaluate prediction performance across all methods using the coefficient of determination (R²) and Pearson correlation coefficients. For group selection, we assess total task performance across each caaissts-classfiicatiii error (CelebA), cross-entropy (Chemical), and mean-squared error (others).
4.2 EVALUATION OF THE PROPOSED FRAMEWORK
We evaluate ETAP with respect to three criteria: (I) how well proposed task-affinity scores capture inter-task relationships; (2) our MTL gain predictions vs. established baselines; and (3) our selection of task groups vs. baseline approaches. We choose several representative baselines across affinity scoring, white-box and data-driven gain prediction, and MTL. performance optimization. (i) TAG (Fifty al., 2222)) estimates MTL gains for higher-order task groups by averaging pairwise affinity scores obtained through a single MTL model with n extra forward/backward passes per up- date; (ii GRAD-TAE et al. 2024) derives pairwise task affinities by estimating the performance of random task groups and averaging their resulting performance: (iii) MTGNet (Song et al. 2022) uses self-attention transformer to predict MTL gains; (iv) Linear Surrogate (Li et al. (023) pre-
dicts binary transfer outcomes in MTL. We also include (v) PCGrad (Yu et 2020), a gradient projection method orthogonal to our goal of pre-training group selection. Additionally, we include
an ablation study in Section C evaluating the contribution of each component of ETAP
Table 2: Correlation between pairwise affinity scores and ground-truth MTL gains (higher values are better).
Affinity Scoring Method
CelebA
ETTml
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 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 1 0.03
0.36 0.03
Prediction of Pairwise MTL Affinity A key contribution of our approach is the proposed task- affinity score (Section 3.1.3.. We compare our pairwise affinity scoring method to two state-of-ther art methods, TAG and GRAD-TAR Table @ compares their correlation with ground-trutt -tuth pairwise MTL gains- higher correlation indicates better reflection of task dependencies, which is crucial for predicting group-level MTL gains. Since the affinity scores output by TAG
different
7

---

Published as a conference paper at ICLR 2026
Table Correlation between ground truth and predicted MTL gains for groups (higher values are better).
Computational Cost
Correlation with Ground Truth
Method
CelebA
ETTml
Chemical
Ridership
TAG (Fifty et al. 2021)
0.10 + 0.0
0.47 0.0
0.05 0.1
0.15 1.15+00 0.1
Suraia =
0.10 + 0.2
0.43 0.1
0.22 :.1
0.43 1 0.1
MTGNet (Song al.
2022
Gmai
= 10
0.22 + 0.1
1.540022
0.34 + 0.2
0.61 + 0.0
Suxin =5
0.41 0.2
0.77
0.40 1 0.1
0.68 =: 0.1
ETAP
Mmain
0.45 1 0.1
0.84 0.0
0.50 1 0.1
0.74 4: 0.0
scale than ground-truth MTL gains, their R² may be very low. To ensure that our comparison is fair, we compare the methods based on their correlation, which is independent of scale. While the MTL-performance estimation method of the GRAD-TAE approach can estimate MTL model performance reasonably well for groups, their pairwise task task-affinity scoring approach struggles to aligm with ground-trule truth pairwise MTL gains. TAG achieves higher correlation, but incurs significant computational overhead. In contrast, ETAP achieves comparable correlation with actual pairwise affinities while substantially reducing runtime (46% on CelebA, 71 1 on ETTml, 9% on Chemical, and 63% on Ridership; see Section D.I
Accuracy of MTL Gain Prediction for Groups Since ETAP integrates both white-box and data- driven components to predict MTL gains, we compare with three baselines, TAG, MTGNet, and I Linear Surrogate, in terms of prediction accuracy and data efficiency. TAG provides affinity scores to estimate MTL gains, but since these differ in scale from actual gains, a direct comparison based on R² would require rescaling. Instead, we report the correlation between predicted and ground-truth MTL gains. which is scale-invariant and ensures a fair evaluation. Table 3 shows correlations for TAG, MTGNet, and ETAP. TAG achieves moderate correlations but requires full MTL training with n additional forward and backward passes for every task pair, making computationally expensive. MTGNet is unstable with smaller training sets Guain- performs better with larger ones, but fails to achieve consistently high correlations. In contrast, ETAP demonstrates strong and stable correlations even with limited training sets (19mainl 5), achieving correlations of 0.41 on CelebA, 0.77 on ETTml, 0.40 on Chemical, and 0.68 on Ridership. Correlation improves further with more training data. These findings underscore the efficiency and robustness of our approach, which outperforms both TAG and MTGNet while significantly reducing computational cost. ETAP also achieves higher correlation and F1 scores than Linear Surrogate (Section D.2) at the same computational cost on all datasets, e.g., improving F1 (0.18 + 0.31) on ETTml and correlation (0.49 + 0.57) on CelebA
0.6
0
R² -5
0.4
MIGNNet
MIGNNe
-2
MIGNet
0.2
MIGNer
ETAP
STAP
ETAP
ETAP
-10
10
30
50
75
10
30 40 50
10
75
10
20
40
50
Cost ((Grrii)) )
Cost ( ( ))
Cost (1rrai
Cost (19urai))
(a) CelebA (n
9)
(b) ETTml (n (n==7)
(c) Chemical (n1 10)
(d) Ridership (n 10)
Figure 2: Prediction performance (AP) vs. computational cost ( (Gunia) for data-driven predictors. MTGNet and ETAP. MTGNet suffers from instability, with means outside the IQR due outliers, whereas ETAP maintains consistency.
Data Efficiency MTL Gain Prediction We compare the computational cost of ETAP with the cost of MTGNet, a data-driven baseline, measuring computational cost as the number of MTL models trained to build training set Grain Figure reports test-set I22 values with interquartile range (IQR) for various training-set sizes GGrain| We finetuned the hyperparameters of MTGNet for each training set to optimize its performance. However, MTGNet exhibits highly random and unstable performance with small training sets, starting with large negative R² values and wide ranges. The mean R² value sometimes falls outside the upper or lower bounds due to extreme outliers, which highlights the need for data-driven predictor that is reliable and stable. In contrast, ETAP consis-
8

---

Published as a conference paper at ICLR 2026
tently yields positive R values with narrow IQR across all datasets and training set sizes, indicating greater robustness.
Table 4: MTL performance (measured as total loss (# std); lower better) with tasks grouped based on various prediction methods. Optimal values are computed using an exhaustive search for ColebA and ETTml, and using a scarch over a representative 50% of all task combinations for Chemical and Ridership.
Dataset
#Groups
TAG
GRAD-TAE
MTGNet ( [9rin))
ETAP ((rrai)
Optimal
2
49.67 1 0.00
50.78 1 0.59
50.62 1 1.20
49.92 - 0.23
49.27
CelebA
(uninl
10)
3
50.22 1 0.00
50.78 1 1.93
50.31 1: 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
4.08 0.09
= 4.04 = 0.06
4
4.02 +
3.98
ETTml
3
0.03
3.96 0.03
3.93 + 0.07
3.83
[Sain| 10)
4
0.02
4.01 = 0.04
0.06
3.89 = 0.09
3.82
2
4.69
4,80 0.14
4.79 0.31
4.67 + 0.12
4.56
Chemical
0.06
4.83 = 0.05
4.89
4.74 0.07
4.52
( (19min] = 10)
4
4.95
4,96 = 0.10
4.94 0.09
4.83 + 0.10
4.67
Ridership
2
17.50 0.00
18.48 + 0.80
17.86 0.43
17.77 0.26
17.03
18.31 + 0.00
17.94 0.47
18.12 0.30
17.83 0.26
16.90
( (19min)
10)
18.25 + 0.00
18.45 + 0.36
18.06 0.42
17.59 = 0.19
16.79
Prediction-based Group Selection We evaluate ETAP against the baseline methods, TAG, GRAD-TAF, and MTGNet,, in terms of the MTL performance of task groups selected based on their predictions. Using each method's MTL gain predictions, we apply franch-and-bound algo- rithm to select task-groups from candidate set of groups For GRAD-TAE, we used the pairwise affinity values estimated by this method the branch-and-bound algorithm. Table reports MTL. performances under various selection budgets, along with the MTL performance of optimal task grouping as reference points. The color gradients indicate relative performance, darker shades in- dicating better results. ETAP consistently achieves near-optimal performance, outperforming TAG and MTGNet, with lower variance. Even with small training set, ETAP enables more effective grouping decisions and further improves with more data. Compared PCGrad (Section D.3. ETAP consistently achieves lower loss reduuingg error by up to 7.4% on ETTml (4.20 - 3.89) and 6.6% Ridership (18.84 17.59) highlighiin the benefit of explicit grouping over implicit gradient-based conflict resolutions.
Table 5: Correlation and R² between ground-truth and predicted total MTL gains for groups (higher values are better).
Correlation with Ground Truth
R²
Dataset
Guiin
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
1
0.77 0.04
1 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
20
0.13 + 0.08
0.46 1 0.03
-0.25 + 0.13
0.18 + 0.05
Accuracy of Total MTL Gain Prediction for Groups We also compare ETAP to the data-driven prediction method proposed by [Ayman et al. 2202). Note that the method of Ayman al. (2023) predicts the total MTL gain for all tasks in group, instead of predicting MTL gain for each task in a group as ETAP and most other methods do (e.g., TAG and MTGNet). For a fair comparison, we adapted ETAP to the evaluation setting of Ayman et al. (2023):
used ETAP predict MTL gain for each task in a group, and then we summed these gains up to obtain a predicted total MTI gain for the group. We present the results of these experiments in Table 5 which shows the correlation and R² between ground-truth and predicted total MTL gains (higher values are better). The results demonstrate that ETAP significantly outperforms this bascline many cases.
Note that the computational cost of the method of Ayman et al. (2023) includes the substantial cost of training MTL models for (*) pairs of tasks. Further, the method also requires measuring ground-truth MTI gains for set of training groups Guain which incurs the cost of training Gurai

---

Published as a conference paper at ICLR 2026
MTL models for these groups. In contrast, ETAP achieves better accuracy at lower cost: method requires training only one MTL model to compute affinity scores (with an additional (
dot-product operations over already-computed gradient vectors, which incurs very low overhead)
and measuring ground-truth MTL. gains for a set of training groups Guain- Since the number of training groups in most practical applications can be at least order of magnitude lower than (?) ETAP incurs a significantly lower overall computational cost.
5 RELATED WORK
key challenge of multi-task learning is modeling task relationships and minimizing negative trans- fers between tasks. Early work (Argyriou et al. 20088 Caruana 1997) showed that tasks with simi- larities or shared underlying structures benefit from joint training. More recent research has explored transferability from the perspectives of zero-shot learning (Pal & Balasubramanian 2019), represen- tation learning (Dwivedi Roig 2019), and information theory (Achille et al,[(222)) Surveys by Zhang & Yang (222)) and Ruder (2117) review techniques for modeling task relatedness via shared features, low-rank structures, or clustering. As MTL applications scale, modeling task relatedness becomes critical not just for interpretability, but for guiding optimization either during training or through informed task grouping beforehand.
MTL Optimization via Training Dynamics Gradient-based MTL strategies (e.g.. PCGrad Yu et al. 2020), GradNorm (Chen et al.) 2018). Uncertainty Weighting (Kendall et al. (0088) aim to mitigate task interference during training by optimizing shared training dynamies, rather than pre- selecting task groups. Our work is orthogonal to these MTL optimization methods since we focus on predicting group-level MTL. gains in advance to guide task-grouping decisions prior training.
MTL Optimization via Pre-selecting Task Groups Recent research efforts have explored pre- dicting MTL gains to optimize task grouping. Data-driven methods, such as HOA (Standley al.) 2020) and MTGNet (Song et al. 2022), require performing MTL training for many task groups to obtain ground-truth MTL gains, making them computationally expensive. [i et al. (2023) assume linear task interactions and introduce linear regression based prediction models, which also require ground-truth gains for many task groups. [i al. (2024) propose an alternative, GRAD-TAE: in- stead of directly measuring ground-truth gains, GRAD-TAE fine-tunes MTL hyperparameters on random task groups to estimate MTL performance. Although these estimates correlate well with the task performance of MTL, they remain resource intensive. Other approaches include analyzing task characteristics in NLP to understand MTL performance (Bingel & Søgaard 2017) and meta- learning from task-specific features (Ayman et al.) 2023); though these often underperform with limited data. Fifty et al.] (202I) propose TAG, a white-box method that analyzes training dynamics through gradient-basce affinity measurements, but requires multiple forward and backward passes per step, leading to high computational overhead. These methods, while offering low-variance esti- mates of pairwise affinities, may not fully capture complex task dependencies in larger groups.
6 CONCLUSION
Our proposed Ensemble Task Affinity Predictor (ETAP) framework addresses the limitations of existing approaches-innttaiilie in data-driven models and high computational cost in white-box methods-by refining low-cost affinity scores with data-driven prediction and residual correction, balancing computational efficiency and prediction accuracy. We demonstrated that by leveraging non-linear transformations and residual correction, ETAP achieves consistent, robust performance across multiple benchmark datasets, spanning diverse domains. Our framework lays foundation for scalable MTL solutions that can optimize task grouping decisions across diverse domains, enabling more efficient multi-task learning applications.
This material is based upon work supported by the National Science Founda- tion (NSF) under Award No. CNS-1952011 and by the U.S. Department of Energy (DOE) under Award No. DE-EE001 188. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the NSF and the U.S. DOE. We would like to thank the anonymous reviewers of ICRL 2026 for their feedback on work and their suggestions to improve our manuscript.
Acknowledgmentts
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Alessandro Achille, Giovanni Paolini, Glen Mbeng, and Stefano Soatto. The information complexity of learning tasks, their structure and their distance. Information and Inference: A Journal of the IMA, 10(1): 2021.
Andreas Argyriou, Theodoros Exgeniou, and Massimiliano Pontil. Convex multi-task feature learn- ing. Machine Learning, 73:243-277 2008.
Afiya Ayman, Ayan Mukhopadhyay, and Aron Laszka. Task grouping for automated multi-task machine learning via task affinity prediction. arXiv preprint arXiv:2310.1624/ 2023.
Joachim Bingel and Anders Søgaard. Identifying beneficial task relations for multi- learning in deep neural networks. In 15th Conference of the European Chapter of the Association for Com- putational Linguistics (EACL), volume 2, pp. 164 169, Valencia, Spain, April 2017. Association for Computational Linguisties.
Carl De Boor. practical guide to splines. Applied Mathematical Sciences. Springer, 1978. ISBN
Leo Breiman. Random forests. Machine learning, 45:5-32, 2001.
Rich Caruana. Multitask learning, Machine Learning, 28(1):41-75, 1997.
Tianqi Chen and Carlos Guestrin. XGBoost: A scalable tree boosting system. In 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), pp. 785- 794, San Francisco, CA, USA, 2016. ACM.
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. GradNorm: Gradient normalization for adaptive loss balancing in deep multitask networks. 35th International Con- ference on Machine Learning (ICML), pp. 794-803. PMLR, 2018
Michael Crawshaw. Multi-task learning with deep neural networks: A survey. arXiv preprint
N0t1:20009009966.. 2020.
Kshitij Dwivedi and Gemma Roig. Representation similarity analysis for efficient task laxonomy & transfer learning. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp.
Long Beach, CA, USA, 2019. IEEE.
Max Ehrlich, Timothy Shields, Timur Almaev, and Mohamed Amer. Facial attributes classi- fication using multi-task representation learning. In IEEE Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pp. 47-55, Las Vegas, NV, USA, 2016. IEEE.
Chris Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, and Chelsea Finn. Efficiently iden- tifying task groupings for multi-task learning. 35th Annual Conference Neural Information Processing Systems (NeurIPS), 34:27503-275 2021.
Pengsheng Guo, Chen-Yu Lee, and Daniel Ulbricht. Learning to branch for multi-task learning- In 37th International Conference on Machine Learning (ICML), pp. 3854-3863, PMLR, JMLR, 2020.
Hu Han, Anil K Jain, Fang Wang, Shiguang Shan, and Xilin Chen. Heterogencous face attribute estimation: A deep multi-task learning approach. IEEE Transactions Pattern Analysis and Machine Intelligence, 40(11):2597-266 2017.
Trevor Hastie, Robert Tibshirani, and Jerome Friedman. The elements of statistical learning: data mining, inference, and prediction. Springer, 2009. ISBN 978-0-387-88 0-3877-4855700
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, Las Vegas, NV, USA, 2016. IEEE.
Laurent Jacob, Jean-philippe Vert, and Francis Bach. Clustered multi-task learning: A convex formulation. Advances in Neural Information Processing Systems (NeurIPS 2008), 21:745-752, 2008.

---

Published as a conference paper at ICLR 2026
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 7482-7491, 2018.
Dongyue Li, Huy L Nguyen, and Hongyang R Zhang. Identification of negative transfers in multi- task learning using surrogate models. arXiv preprint wrXiv:2303.14585 2023.
Dongyue Li, Aneesh Sharma, and Hongyang Zhang. Scalable multitask learning using gradient- based estimation of task affinity. In 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD), pp. 1542-1553, Barcelona, Spain, 2024. Association for Computing Machinery.
Wu Liu, Tao Mei, Yongdong Zhang, Cherry Che, and Jiebo Luo. Multi-task deep visual-semantic embedding for video thumbnail selection. In 28th IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 3707-3715, Boston, MA, USA, 2015a. IEEE.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In 2015 IEEE International Conference Computer Vision (ICCV), pp. 3730-3733 Los Alamitos, CA, USA, 2015b. IEEE Computer Society.
Minh-Thang Luong, Quoe V Le, llya Sutskever, Oriol Vinyals, and Lukasz Kaiser. Multi-task sequence sequence learning, arXiv preprint arXiv:1511.0611 14, 2015.
Arghya Pal and Vineeth N Balasubramanian. Zero-shot task transfer. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2189-2198, Long Beach, CA, USA, 2019. IEEE.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research,
12:2825-28322 2011.
Kriti Puniyani, Seyoung Kim, and Eric P Xing. Multi-population GWA mapping via multi-task regularized regression. Bioinformatics, 26(12):i208-22 -20..,221,,2220 2010.
Michael T. Rosenstein, Zvika Marx, Leslie Pack Kaelbling, and Thomas G. Dietterich. To transfer or not to transfer. In NIPS 2005 Workshop Inductive Transfer: 10 Years Later, 2005.
Sebastian Ruder. An overview of multi-task learning in deep neural networks. abs/1706.05098, 2017.
CoRR,
Xiaozhuang Song, Shun Zheng, Wei Cao, James Yu, and Jiang Bian. Efficient and effective multi- task grouping via meta learning on task combinations. 36th Annual Conference on Neural Infor- mation Processing Systems (NeurIPS), 35:37647-37659, 2022.
Trevor Standley, Amir Zamir, Dawn Chen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? In 37th International Conference on Machine Learning (ICML), volume 119, pp. 9120-9132, virtual, 2020. JMLR.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum deep learning. In 30th International Conference on Machine Learning (ICML), pp. 1139-1147, 2013.
Haixu Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. Autoformer: Decomposition trans- formers with auto-correlation for long-term series forecasting. 35th Annual Conference on Neural Information Processing Systems (NeurIPS), 34:22419-224444444
2021.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. 34th Annual Conference Neural Information Pro- cessing Systems (NeurIPS), 33:5824-5836, 2020.
Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. In 2018 IEEE Conference Com- puter Vision and Pattern Recognition (CVPR 2018), pp. 3712-3722, Salt Lake City, UT, USA, 2018. IEFE.
12

---

Published as a conference paper at ICLR 2026
Yu Zhang and Qiang Yang. A survey multi learning. IEEE Transactions on Knowledge and Data Engineering, 34(11))5586- 5609, 2021.
Zhanpeng Zhang, Ping Luo, Chen Change Loy, and Xiaoou Tang. Facial landmark detection by deep multi-task learning. In 13th European Conference on Computer Vision (ECCV 2014), pp- 94-108, Zurich, Switzerland, 2014. Springer, Springer.
Jiayu Zhou, Lei Yuan, Jun Liu, and Jieping Ye. A multi-task learning formulation for predicting disease progression. In 17th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD 201 1), pp. 814-822, San Diego, CA, USA, 2011. ACM.
Ammar Bin Zulqarnain, Samir Gupta, Jose Paolo Talusan, Dan Freudberg, Philip Pugliese, Ayan Mukhopadhyay, and Abhishek Dubey. Addressing APC data sparsity in predicting occupancy and delay of transit buses: A multitask learning approach. In 9th IEEE International Conference on Smart Computing (SMARTCOMP), pp. 17-24. IEEE, 2023.
13

---

Published as a conference paper at ICLR 2026
A
NOTATION
To help readers navigate the theoretical framework and empirical findings effectively, Table presents comprehensive guide the notation used throughout the paper.
Table 6: List of Notations
Symbol
Description
T
Set of all tasks
Number of tasks (n [T]
L
Loss of task &
LT
Total loss across all tasks (LT = EiST Lj)
G
A lask-group (i.e., set of tasks G -
A collection of task groups (g [G1, ...
X
Input to any model (i.e., training/ test data for any task) Notations for STL and MTL Models Parameters for an STL model trained on a single lask t 7
Smn.
Parameters for an MIL model trained on group of tasks G CT Bsm
yG-+8
MTL gain for task in group G (#G-+ [L, (X, 0sn) L((X (xn))//L((X(,,
yG
Vector of MTL gains for all tasks in group G
Notations for Affinity Score Calculation
zk
Task-affinity score task fy on Ł at gradient step
Aggregated task-affinity
for task f; iú across entire training
Task-affinity score of a
G on task & (EG-1. Ø-I Ee,ec,jas Z1y+
2G
Affinity scores for all tasks within group G ((zg-e: E G)).
Notations for MTL Gain Predictor
f foon-linear Φ
Initial predictor function ETAP (non-linear mapping)
Non-lincar transformation function Regression coefficients for fuon-lincur
Waif ge
Prediction from fnon-tincar for a task group C
Error between actual and predicted MTL gains for task group G' (yG ( Residual predictor function in ETAP (residual model)
fresidual
Wres
Model coefficients for fresidual
Multi-hot encoding the membership of task group (ug E {0. 1)1 Gall
UG
iual
Final MTL. gain predictions for tasks in group, G (gilina Ve fresidual UGg Whes)))))))
B EXPERIMENTAL SETUP
B. DATASETS AND MTL ARCHITECTURES
We evaluate our proposed approach on four multi-task learning benchmarks from diverse domains, which had been used assess task grouping strategies in prior MTL studies.
CelebA Liu et al. (015b) is a large-scale dataset of face images with multiple attributes, commonly used in MTL research Ehrlich al. (2016); Guo
(2020); ) Han al. (2017). The Task-Affinity Grouping (TAG) approach Fifty et al. al. (2021) evaluated task combinations on subset of nine tasks from CelebA. We adopt the same dataset and use the optimized ResNet- 18 architecture He et al. (016)) as the baseline MTL model, following the experimental setup in prior work Fifty et al. (2021).
ETTm1 (Electricity Transformer Temperature) Wu al. (2021) an electric load dataset compris- ing seven time-series forecasting tasks which are used to evaluate MTGNet approach [Song et al.
(2022). We follow prior work Wu al.] (202T), framing it as a multivariate time-series forecasting problem, where each series forecast is treated as separate task. The Autoformer Wu et al. (2021)
model serves as the MTL architecture with an SGD optimizer.
Chemical (MHC-I) Jacob al. (2008) is dataset designed to classify peptide binding affinities for different molecules. The dataset contains 35 tasks, each representing a molecule's 's binding clas- sification. We select 10 tasks with various numbers of training samples for experiments. Each task has 180 binary input features. We implement feed-forward neural networks for both multi-task learning (MTL) and single-task learning (STI.), using shared encoder followed by task-specific
14

---

Published as a conference paper at ICLR 2026
output layers. The shared encoder consists of two fully connected layers with 32 and 16 ReLU units, respectively, capturing a common representation across tasks. Each task-specific decoder contains single sigmoid-an activated output neuron. Models are trained using binary cross-entropy loss and optimized with stochastic gradient descent (learning rate = 0.001, momentum = 0.9). This single architecture used consistently for all STL and MTL. runs.
Alongside these widely used benchmarks, we also employ a novel, real-world transit ridership dataset from U.S. city, curated by Zulqarnain et al. (2023), containing stop-level passen- ger counts for buses, enriched with predictive features. Each task is to predict bus ridership for a specific route and direction. We select 10 such tasks with various data availability, from data-rich to highly sparse. We use a hybrid multi-input MTL architecture to model ridership prediction across routes. The predictive features are of three types: (i) sequential statisties from the past 10 trips (e.g., median/max load, time gaps), (ii) numerical features such as weather, delays, and service headway, and (iii) categorical features like time-of-day and calendar indicators, embedded via learnable layers (embedding dim 4). The shared encoder combines stacked Bidirectional LSTMs (64 units each) for temporal features, feed-forward layers for numerical features, and embeddings for categorical 16-dimensional latent space. Each task-specific decoder contains two dense layers (8 and 1 units) for final regression. A shared encoder captures common structure across tasks, while lightweight task-specific decoders generate the final predictions. Models trained with Adam (learning rate = 0.001), using MSE loss, dropout (0.2 - 0.3), and batch normalization for regularization.
Ridership:
inputs, projecting
Please note that the MTL architectures are not a critical component of our approach and serve only enable experimentation with predicting MTL gains and grouping tasks.
B.2 HYPERPARAMETER SEARCH
All ETAP hyperparameters are tuned via cross-validation on randomly selected subsets of training groups. The first stage of ETAP employs regularized linear regression model over B-spline transformation of affinity scores. We perform a grid search for spline degree (2-6), number of knots (up to /9rrain samples, where Grain indicates the number of groups for training the predictor), and Ridge regularization strength (a ranging from 0.001 1.0). The residual correction stage uses Ridge regression with a selected via cross- cross-validation over similar range. For each training instance of the gain predictor, the best hyperparameters are determined using cross- coss-validation over training groups only, and then fixed for prediction on unseen groups To ensure robustness and minimize the impact of randomness, all training and evaluation procedures are repeated with different random seeds: 10x for gain prediction and 6x for group selection. The seeds are randomly chosen, and the specific seeds used to generate the results reported in the main text are documented the code appendix.
For the MTGNet baseline Song et al. 2022), we conduct separate hyperparameter search over the number of Transformer layers and the embedding dimension, selecting from values in the range of to 64. For further architectural details, we refer readers to the code implementation provided in Song al. (2022). All models are trained and validated using only training groups, ensuring fair generalization to held-out test groups. To replicate TAG Fifty et al. (202T) and PCGrad Yu et al. the official software implementations provided by the original authors.
2020),
Hardware Configuration All model training (including both single-task and multi-task experi- ments) is performed on a Unix-based high serh-peffomance server running Ubuntu 22.04.5 LTS. The system is equipped with two NVIDIA RTX A5000 GPUs (24 GB each, CUDA 12.8). GPU acceler- ation is utilized for training where applicable.
C ABLATION
To evaluate the contribution each component ETAP, we perform an ablation study, focusing particularly on the residual prediction step. By removing this step and comparing the resulting performance to the entire method, we can evaluate its role in improving MTL gain predictions and refining trok-group estimates.
15

---

Published as a conference paper at ICLR 2026
C.1 CONTRIBUTION OF RESIDUAL PREDICTION
To investigate the significance of incorporating residual prediction step, we train a model without this step and evaluate under identical conditions as the full framework. The comparison highlights how this step influences prediction accuracy and its alignment with ground-truth MTI gains.
0.2
R²
R
0.6
30
30
Cost ( (9tzain]
Cost (9irain)
(a) CelebA (n 9).
(b) ETTml (n 7).
0.5
R?
R²
0.5
0.4
Cost (1
Cost (9uainn)
(c) Chemical (n 10).
(d) Ridership (n 10).
Figure 3: Comparison of prediction performance (H²) between the initial predictions (without residual adjust- ment -*) and the final predictions (with residual adjustment +) across four datasets. The residual adjustment step improves accuracy, particularly CelebA, Chemical, and Ridership, with slight improvement observed on ETTml.
Figure shows comparison in terms of prediction performance (R²-values) on the groups. The results show that residual prediction significantly improves accuracy for the CelebA, Chemical, and Ridership datasets, with higher alignment between predicted and ground-truth MTL gains.
Table Correlation with ground-truth MTL gains across different computational costs ((raii I predictor on all benchmarks (higher values are better).
Correlation (Predicted vs. Actual) Initial
Dataset
Grain
Final
0.27 0.15
0.41 0.16
CelebA
10
0.31 = 0.13
0,45 + 0.13
5
0.75 + 0.15
0.77 + 0.12
ETTml
10
0.84 = 0.02
0.85 + 0.02
5
0.32 = 0.06
0.40 + 0.07
Chemical
10
0.36 + 0.03
0.50 + 0.07
10
0.66 = 0.02
0.71 + 0.03
Ridership
15
0.67 + 0.03
0.73 + 0.04
When comparing initial predictions based on task affinity scores and non- non-linear mapping with those enhanced through residual correction, the differences some cases are marginal. However, this not a limitation of the approach. Even in cases where both approaches achieve similar performance, our suggested method generally maintains a stronger overall correlation with the ground-truth MTL gains and outperforms baseline approaches across all benchmarks shown in Table 7 This con- sistency demonstrates the framework's robustness, particularly its ability to handle complex and non-linear task dependencies at much lower cost. The residual correction step provides a crucial mechanism to address systematic errors in initial gain estimates, leading to more reliable predictions for MTL grouping.
16

---

Published as a conference paper at ICLR 2026
C.2 CONTRIBUTION OF NON-LINEAR MAPPING ON AFFINITY SCORES
To assess the impact of the non-linear transformation in ETAP, we compare ETAP's performance against variant that uses an affine mapping to align task affinity scores with moound-tuute MTL gains. We aim to determine whether the added expressiveness of non-linear mappings (e.g., B- spline) meaningfully improves prediction accuracy over an affine mapping.
Table 8: Effect of non-linear transformation on task -isk-fffintty scores: comparison affine vs. B-spline mapping in terms of Ra and correlation with ground-trum truth MTL gains.
Dataset
Rd (Affine)
R (B-spline)
Corr. (Affine)
Corr: (B-spline)
CelebA
0.13 + 0.00
0.13 + 0.01
0.37 + 0.00
0.37 = 0.01
ETTml
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
As shown in Table 8 the non-linear B-spline transformation consistently matches or outperforms the affine baseline in terms both R2 and correlation across all four datasets. The gains are especially pronounced on ETTml and Ridership, where complex, non-linear relationships exist between affin- ity and MTL gain. These results confirm that MTL gains are not linearly explained by affinity scores alone, and that the non- non-linear mapping plays key role in modeling this complexity. Overall, the B-spline transformation enhances ETAP's ability to predict MTL gains by capturing non-linear pat- terns that affine transformation cannot express, making it essential component of the pipeline.
D ADDITIONAL RESULTS
D.1 RUN-TIME COMPARISON FOR PREDICTION OF PAIRWISE MTL AFFINITY
To predict pairwise multi-task learning affinity, we compute task-affinity scores that estimate the compatibility between task pairs. In this section, we compare the run-time performance of our affinity scoring approach (introduced in Section 3.T with the Inter-Task Affinity (ITA) score used in Pask-Affinity Grouping (Ftt))Fift al. (2021).
CelebA
ETTml
6
8
6.88
4.69
Hours
4
Hours
6
2.51
4
2
2.02
2
0
Inter- Task Affinity (Hlly et 2021
Task Affinity Score (ETAP)
Inter- Task Affinity (TAG) (nflty et 2021
Task Affinity Score (ETAP)
Chemical
Ridership
2
6
Minutes
1.5
1.4
1.27
3.88
Hours
1
2
1.42
0.5
0
0
Inter-Task Affinity (TAG) Fifty at 2021
Affinity (ETAP)
Inter. Affinity (TAG) Fifty et all 2021
Affinity (ETAP)
Figure 4; Runtime for CelebA and ETTml datasets (measured in hours) and the Chemical dataset (measured in minutes). Running times include task-affinit scores calculations and complete multi-task leamnigg training for all tasks together in the benchmark using baseline MTL models.
17

---

Published as a conference paper at ICLR 2026
Figure presents side-by-side comparison of the computational efficiency of the two methods, evaluated under identical hardware and software conditions for faimess. Our approach consistently achieves lower run-times across all datasets-Cellbb,, ETTml, Chemical, and Ridership. For in- stance, on CelebA, TAG requires 6.88 hours to compute affinity scores, whereas our method com- pletes the process in just 2.51 hours. Similar improvements observed other datasets: 71% reduction in runtime on ETTml, 9% on Chemical, and 63% on Ridership.
These results highlight the scalability of our approach and its ability to deliver robust affinity pre- dictions with significantly reduced computational overhead.
D.2 COMPARISON WITH LINEAR SURROGATE MODEL FOR GAIN PREDICTION
We also compare our proposed method, ETAP, with the Linear Surrogate model proposed by Li al. (2025) in terms of prediction accuracy. Linear surrogate approach was originally designed for predicting binary transfer outcomes (positive vs. negative transfer) in MTL. To ensure a fair comparison, we train both ETAP and the surrogate baseline on the same set of task groups across four benchmarks.
Since the approach of Li et (2023) focuses on binary classification of transfer gains, we evaluate both methods using two metrics:
Correlation: Pearson correlation between predicted and ground-truth MTL gains.
F1 Score: Computed on the binary classification task of predicting the correct sign of the MTL gain.
Table 9: ETAP vs. Linear Surrogate Model Li et al.
2023 for identifying positive vs. negative transfer.
Training Cost ((9main) )
Surrogate Model
ETAP
Dataset
Corr.
Fl
Corr.
Fl
CelebA
32
0.49
0.61
0.57
0.70
ETTml
21
0.35
0.18
0.64
0.31
Chemical
0.57
0.92
0.62
0.94
Ridership
35
0.67
0.81
0.77
0.85
As shown in Table 9] ETAP consistently outperforms the surrogate modeling approach Li et al. (2023) across all datasets in both correlation with ground-trun MTL gains and F1 score for classi- fying positive
negative transfer. The performance gap is particularly notable on ETTml, where ETAP achieves a correlation of 0.64 and an F1 score of 0.31, significantly surpassing the surrogate model's weaker correlation (0.35) and low F1 score (0.18). On CelebA and Ridership. ETAP also improves both metrics, demonstrating better alignment with ground-truth and stronger classification of transfer outcomes. Even Chemical, where the surrogate model performs relatively well, ETAP offers a modest but consistent improvement. These results highlight ETAP's superior ability to cap- ture complex, non-linear, and higher-order task interactions caapablitis that the surrogate model, constrained by its linear assumption, fails to match.
D.3 COMPARISON WITH PCGRAD FOR OPTIMIZED MTL PERFORMANCE
To further assess the effectiveness of ETAP, we include comparison with PCGrad Yu et al. (2020). gradient-projection pwwpeeriio--asseeeeeee based method designed resolve conflicts between task gradients during train- ing. Unlike our grouping-based strategy, PCGrad assumes a single unified task set (i.e., one group that includes all tasks) and operates entirely in the gradient space, making it orthogonal to our ap- proach.
Table 10 presents MTL performance (lower is better) across four diverse datasets. ETAP outper- forms both Naive MTL (which jointly trains all tasks without grouping) and PCGrad in most cases, especially when tasks are split into 3-4 groups For example, on the ETTml dataset, ETAP with 4 groups achieves loss of 3.89, significantly lower than 4.20 (Naive MTL) and 4.19 (PCGrad). Similar trends are observed for Ridership, where ETAP reduces loss to 17.59, compared to 18.95 and 18.84.
18

---

Published as a conference paper at ICLR 2026
Table 10: Comparison of multi-task learning performance (measured in loss; lower is better) across different methods (Naive MTL, PCGrad, our proposed ETAP method with 2, 3, and task-group splits). Results are reported as mean + standard deviation across multiple runs.
ETAP (SPLITS)
Dataset
Naive MTL
PCGrad
3
4
CelebA
50.02 + 0.4
49.99 0.5
49.92 + 0.2
49.61 0.3
49.41 - 0.3
ETTmi
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
Interestingly, on the Chemical dataset, PCGrad underperforms Naive MTL, suggesting that gradi- ent projection alone may not resolve deeper incompatibilities between certain tasks. Meanwhile, ETAP maintains competitive performance by grouping synergistic tasks. On CelebA, all methods yield comparable results, though ETAP still benefits from modest improvements with more refined groupings.
These results reinforce that task grouping can capture task relationships beyond gradient conflict and may complement or even outperform purely optimization- strategies like PCGrad.
D.4 STATISTICAL SIGNIFICANCE
To assess whether the performance differences between ETAP and the baselines are statistically significant, we conduct paired Wilcoxon signed-rank test across all datasets and training sample
Table 11: Wilcoxon signed-rank test results comparing MTGNet ETAP across datasets and training sample
sizcs. Bold p-values indicate statistical significance p < 0.05.
Dataset
Computational Cost
Sample Size
Wilcoxon p-value
CelebA
10
0.0020
10
10
0,0039
20
10
0,0039
0,0020
10
0.0039
50
10
0.0020
75
10
0.0098
ETTml
5
10
0.0020
10
10
0,0020
0.0137
0,0137
40
10
0,0195
45
10
0.8457
50
10
0,0059
Chemical
10
0,0020
10
0.0039
0,0020
30
10
0.0039
40
10
0,0039
50
10
0.0020
0.0020
Ridership
10
10
0.0020
0.0020
20
10
0,0020
25
10
0.0020
30
10
0.1602
40
10
0.2304
0.0137
19

---

Published as a conference paper at ICLR 2026
sizes. This con-aaamettri test evaluates whether the distribution of prediction performance (e.g., R or correlation with ground-truth MTL gains) differs consistently between two approaches.
Table I1 shows whether the differences between ETAP and MTGNet are statistically significant in terms of prediction performance (JRº). For each configuration, we collect paired performance values from 10 randomized trials and compute p-value using the Wilcoxon signed-rank test. p- value tests the null hypothesis that the two models perform equally, considering deviations in either direction (i.e., whether ETAP better or worse). Results reported in Table where statistically significant differences (p < 0.05) are shown in bold. Across most settings, ETAP significantly outperforms MTGNet, particularly at lower computational costs. This confirms the robustness of the observed improvements and suggests that the gains from ETAP are unlikely to be due to random chance.
Table 12; Wilcoxon signed-rank test p-values comparing TAG vs. ETAP across benchmarks.
Dataset
Sample Size
Cost (ETAP)
Wilcoxon p-values
CelebA
7
5
0.0313
ETTml
7
5
0.0313
Chemical
7
5
0.0156
Ridership 7
0.0078
As shown in Table 12 we assess the statistical significance of the difference in gain prediction performance between TAG and ETAP. Since TAG is unsupervised method that does rely on training samples, we compare it against ETAP at fixed training size (lguail 5 or 10 depending on the benchmark), using correlation with ground-truth MTL gains across multiple random seeds. A Wilcoxon signed-rank test reveals that ETAP significantly outperforms TAG on all benchmarks (p 0.05).
E ALTERNATIVE COMPONENTS FOR ETAP FRAMEWORK
For our ensemble framework, ETAP, we experiment with different components for both stages of the ensemble.
E.1 APPROACHES FOR INITIAL PREDICTION
The main text presents results for B-spline transformation followed by a regularized linear regres- sion. Here, we present alternative methods tested for the initial prediction process.
K-nearest neighbor (KNN): We use KNN regressor predict MTL gains by determining the optimal number of neighbors (k). The input features, derived from pairwise task affinities, are scaled before training. We use cross-validation select & by minimizing the mean squared error (MSE).
Random Forest Regression (RF): We utilize Random Forest regressor Breiman 2001 with task- affinity-based predictions as input features. Hyperparameters such as the number of trees, tree depth, and minimum sample requirements are optimized via a grid search with cross-validation based on MSE.
For all approaches, we implement the models using the standard functions from Sci kit Learn [Pedregosa al, (2011). We determine the optimal grid search and cross- cross-validand validation on the training groups, where the respective models are evaluated based on the mean squared error (MSE) the evaluation metric. The best models are selected based on cross-validation performances. Once the optimal hyper-parameters are identified, the models are trained on the entire training data and used to make predictions both the training and test sets, which are later utilized in the residual prediction step.
for the models using
Results Figure 5 illustrates the performance of various approaches applied to the first step of an ensemble model all datasets. These approaches include our suggested method, BSplines Basis
20

---

Published as a conference paper at ICLR 2026
BSpline
JOKNN In Random Forest
BSpline IDKNN I Random Forest
0.28
0.26
0.7
R
0.24
R
0.22
0.65
0.2
10
20
30
40
50
10
20
30
40
Training Groups ((uain))
Training Groups ((uain))
(a) CelebA (n = 9).
(b) ETTml (n = 7).
BSpline
KNN Random Forest
BSpline
KNN In Random Forest
0.4
0.6
0.4
0.2
R2
n22
0.2
0C
0
10
20
30
40
50
10
20
30
40
50
Training Groups (9nmin)
Training Groups (9
(c) Chemical (n 10).
(d) Ridership (n = 10).
Figure 5: Comparison prediction performance (R²) different approaches applied the first step of ensemble model (Soonlinear) on the CelebA, ETTml, Chemical, and Ridership datasets. These performances are reported after applying residual correction step (Festaa)), where regularized ridge regression is used for residual prediction.
Expansion followed by Regularized Linear regression, chosen for its simplicity, efficiency, and in- terpretability. Alternate methods, such as the KNN, and Random Forest with Ridge regression as the residual predictors also evaluated. Performance is assessed using R²2 scores across various training sample sizes, shedding light on how each method generalizes in different scenarios.
For CelebA and Ridership dataset, BSplines shows relatively stronger performance with all training sample sizes, indicating its ability to model non-lineat task relationships effectively with sufficient training groups. For ETTml, both B-splines with Ridge and KNN maintain robust and stable R² scores across sample sizes, with KNN slightly outperforming B-splines at larger training sizes, high- lighting their capacity to balance complexity and regularization in smaller datasets. For the Chem- ical dataset, all approaches achieve comparable performance, with consistently higher R² scores at larger training sizes.
While Random Forest delivers strong performance, its lack of interpretability and poor handling of sparse data make it less suitable for our needs. Although KNN also performs competitively, espe- cially in certain datasets, its high memory requirements and limited scalability with increasing data size make it less practical choice real-world applications. These factors collectively support our rationale for suggesting BSplines with regularized linear regression as the initial affinity predictor, as it balances model simplicity, computational efficiency, and generalization performance.
E.2 APPROACHES FOR RESIDUAL PREDICTION
For the second stage of our framework, which focuses on residual error correction, we use ridge regression with the regularization parameter optimized via validaaaii...on. We also explore another boosted regression method.
Extreme Gradient Boosting (XGBoost) For residual correction, employ XGBoost regres- sor Chen & Guestrin (2016) to predict the residual emrors between the initial task-affinity- -bsed predictions and the actual MTL. gains. The input features zero-one vectors representing task
21

---

Published as a conference paper at ICLR 2026
BSpline (Ridge)
(Ridge)
RE (Ridge)
RSpline (XGB)
KNN (XGB)
RF (XGR)
0.25
Ra
0.2
10
20
30
40
50
(a) CelebA (71 9).
BSpline (Ridge)
KNN (Ridge)
RF (Ridge)
BSpline (XGB)
KNN (XGB)
RF (XGB)
0.7
0.65
10
20
30
40
(b) ETTml (nc 7).
BSpline (Ridge)
KNN (Ridge)
RF (Ridge)
BSpline (XGB)
KNN (XGB)
(XGB)
0.4
Ra
0.2
0
10
20
30
40
50
(c) Chemical (n 10).
BSpline (Ridge)
KNN (Ridge)
RF (Ridge)
BSpline (XGB)
KNN (XGB)
RF (XGBB
0.6
R²
0.5
0.4
10
20
30
40
Training Groups ((Gaain)
(d) Ridership (n 10).
Figure 6: Prediction performance (R²) for different approaches applied in the residual learning step of the ensemble model (fresidual) on the CelebA, ETTml, 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).
groups, while the labels are the corresponding residual values. We perform hyperparameter tuning using a grid search with cross-validation. optimizing parameters such as the number of estimators, tree depth, learning rate, and subsampling ratios. The model is trained on scaled input data, and predictions are made on training and test sets. This method refines initial predictions by capturing complex patterns in residuals, thereby enhancing MTL gain prediction accuracy.
Results Figure/sssoww the performance of residual models across all benchmarks, evaluated with varying numbers of training groups. After generating initial predictions using task-affinity scores
22

---

Published as a conference paper at ICLR 2026
and non-linear mapping, compare two residual predictors: Ridge regression and XGB, each represented by distinct color gradients in the figure.
Ridge regression consistently achieves high R scores across different datasets and training sizes, demonstrating strong generalization and stability. contrast, XGB exhibits greater variability, par- ticularly with smaller training groups. For instance, while XGB slightly outperforms Ridge the Chemical dataset with larger training sizes, the difference negligible when fewer groups are available. Across CelebA, ETTml, and Ridership, Ridge regression delivers better average perfor- mance, whereas XGB's sensitivity to dataset size and stability makes it less reliable under limited supervision.
Given its robustness, simplicity, and resistance to overfitting, we recommend Ridge regression for residual correction. Its regularization ensures stable improvements even with scarce training data, in contrast to XGB, which can be more sensitive to hyperparameters and noise.
F VARIABILITY OF PER-STEP AFFINITY SCORE DURING TRAINING
2lack Hair
Blond Hair
Brown Bair
Sostee
Per-Step Affinity Score
0.4
1.5
0.2
0.2
0.2
0.1
0.5
Epochs
Epochs
Epochs
Epochs
Rosy Jhheeks
Neeaa ing_Bat
Per-Step Affinity Scote
0.8
0.4
0.2
0.4
0.2
10
20
Epochs
Epochs
Epochs
Epochs
Evolution of per-step affinity score =tz- during training between a source task (5.6 Clock Shadow) and each target task j (Black Hair, Blond Hair, Brown Hair, Goatee, Mustache No Beare Rosy-Checks, Wear ) in dataset CelebA
Figure
23


Reviews:
[{"review_id": "l5IbmixKZv", "reviewer": "Reviewer_Y2PA", "summary": "Proposed ETAP builds a scalable predictor by computing a gradient-alignment affinity score for pairs and groups in shared parameters, then refining it with learned nonlinear transformations and residual corrections. Across benchmarks, ETAP improves MTL gain prediction and enables more effective task grouping, outperforming used baselines.", "questions": "It would help to clarify the concrete differences from prior work, especially how your methods distinctively improve on gradient-based affinity and group-wise tracking approaches. Practical guidance on when to prefer your method over more recent baselines would make the contribution clearer.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "It combines gradient-based affinity with learned non-linear relationship modeling to efficiently and accurately capture task relationships. And it includes thorough component-wise ablations that clarify each contribution and improve interpretability.", "weaknesses": "Dividing learning tasks into groups based on similarity is a long-standing area [1]. The paper introduces new measures of task affinity for MTL, but I am not fully convinced that the proposed methods are superior to prior work. The baselines used are relatively dated, and a comparison of computational cost and predictive performance with stronger recent baselines, such as [2], would strengthen the claims. Efficient group-wise tracking of task affinity is also not new, as [3] tracks inter-task affinity in a group-wise manner during multi-task optimization. Finally, I am not convinced that the proposed methods clearly improve over other inter-task affinity tracking approaches, including [4].\n\n[1] Which tasks should be learned together in multi-task learning?\n\n[2] Task Grouping for Automated Multi-Task Machine Learning via Task Affinity Prediction\n\n[3] Selective Task Group Updates for Multi-Task Optimization\n\n[4] Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2DBbcx4Naz", "reviewer": "Reviewer_LwS8", "summary": "Summary\n\nThis paper addresses a central challenge in multi-task learning (MTL): identifying groups of tasks that mutually improve each other’s performance when trained together. Existing methods typically fall into two categories — white-box and black-box approaches.\n\nThe paper introduces a hybrid method that integrates both. The proposed two-step framework first estimates pairwise task affinities by training a single MTL model on all tasks within a group, inferring affinities for all possible combinations. In the second step, a non-linear mapping from task affinity to performance gain is constructed and refined using residual predictions.\n\nThe method is evaluated against several state-of-the-art task grouping algorithms, with ablation studies highlighting the contribution of each component. Results demonstrate the benefits of applying non-linear mappings and residual predictions. Overall, the paper presents a robust and well-conceived methodology that successfully combines the strengths of existing approaches.", "questions": "Questions for the Authors\n\nWhat is the rationale for time-averaging the cosine similarity over the \nK\nK training steps in Equation (6)? How stable is this measure across training, and how much variability is typically observed?\n\nSince TAG’s affinity scores are on a different scale than observed gains, how does your method correct for this discrepancy?\n\nYou mention that GRAD-TAE performs well for groups — could you include these results in Table 3 for completeness?\n\nCould you provide additional details on the computational requirements of the different methods, particularly in terms of the number of backward passes and the extent of hyperparameter tuning involved?\n\nRecommendation\n\nAlthough some design choices could be better motivated and the performance gains over naive MTL are modest, the paper provides valuable insights into the problem of task grouping. The authors conduct detailed and well-structured experiments, including informative ablation studies that clarify the role of different components within the proposed framework.\n\nAdditional Feedback\n\nFigure 1 is somewhat difficult to interpret, as it presents too many elements at once. Consider splitting it into two complementary figures: one conceptual illustration to convey the overall idea, and another outlining the algorithmic steps in more detail.\n\nSection 4 could be improved for clarity and consistency. Ensure that all discussed results appear in the corresponding tables or figures, and avoid repetition of statements like “TAG achieves higher correlation but incurs significant computational overhead.”\n\nThe term “training cost” could be replaced with “computational cost” for improved clarity.\n\nIn Section 3.1.3, explicitly explain the difference with TAG and refer to the appendix where both affinity measures are compared.\n\nThe discussion of limitations could be expanded. As shown in the appendix, the method performs comparably to naive MTL in some cases, indicating room for improvement and further exploration in future work.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths\n\nTackles an important challenge in the MTL literature — the high computational overhead of existing task-grouping algorithms. Through comprehensive ablation studies, the paper provides valuable insights into the role of gradient similarities (e.g., via comparison of affine vs. non-linear mappings).\n\nDesign choices (e.g., use of B-splines and regression techniques) are justified through ablation analyses and comparative experiments.\n\nThe paper is clearly structured and effectively relates its contributions to prior work, ensuring coherence and contextual grounding.\n\nThe experimental evaluation goes beyond measuring performance gains, also analyzing the correlation between predicted affinities and ground-truth transfer gains.", "weaknesses": "Weaknesses\n\nWhile computational efficiency is claimed as a key advantage, it would be valuable to include results for the complete approach (including hyperparameter tuning in the second stage) or to explicitly state that the additional cost is negligible. Comparing against an additional data-driven baseline would strengthen the evaluation.\n\nThe performance gains over naive MTL are relatively modest; a discussion of their practical significance would help contextualize their value.\n\nSome design choices lack clear theoretical justification (e.g., why averaging gradient similarities yields effective affinity estimates, or why B-splines are particularly suitable).\n\nThe proposed method introduces several additional hyperparameters (for the B-spline expansion, regression method, and residual prediction), which may complicate tuning and reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xhfj93iMdN", "reviewer": "Reviewer_66hX", "summary": "This paper proposes ETAP (Ensemble Task Affinity Predictor), a framework for predicting multi-task learning (MTL) gains to enable efficient task grouping. The approach combines white-box gradient-based affinity scoring with data-driven ensemble prediction.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The two-stage ensemble design that uses gradient-based affinity scores as foundation and refining with data-driven models is reasonable, it combines white-box gradient-based affinity scoring with data-driven ensemble prediction.\n2. ETAP achieves impressive runtime reduction while maintaining or improving correlation with ground-truth gains. This is a meaningful practical contribution.", "weaknesses": "1. The gradient-based affinity score (Equation 5) is quite similar to existing work (TAG), just removing the auxiliary forward/backward passes. The B-spline transformation and ridge regression are standard techniques. The main novelty seems to be in combining these pieces, which feels somewhat incremental. Can you clarify what is fundamentally new here beyond engineering different existing methods together?\n2. While you claim ETAP is \"scalable,\" all experiments use relatively small task sets (n=7-10). What happens when n=50 or n=100? The affinity score computation still requires training one MTL model with all tasks, and the pairwise scores grow as O(n²). The paper doesn't really demonstrate scalability to large task pools that appear in real applications.\n3. Section 3.2.3 uses branch-and-bound from prior work. So the contribution is really just the gain prediction, not the actual grouping algorithm. This should be made more clear in the contribution claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)