Efficiently Identifying Task Groupings for Multi-Task Learning
Authors: Christopher Fifty
Ehsan Amid
Zhe Zhao
Tianhe Yu
Rohan Anil
Chelsea Finn
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: multi-task learning
task groupings
which tasks should train together

Abstract: Multi-task learning can leverage information learned by one task to benefit the training of other tasks. Despite this capacity, naively training all tasks together in one model often degrades performance, and exhaustively searching through combinations of task groupings can be prohibitively expensive. As a result, efficiently identifying the tasks that would benefit from training together remains a challenging design question without a clear solution. In this paper, we suggest an approach to select which tasks should train together in multi-task learning models. Our method determines task groupings in a single run by training all tasks together and quantifying the effect to which one task's gradient would affect another task's loss. On the large-scale Taskonomy computer vision dataset, we find this method can decrease test loss by 10.0% compared to simply training all tasks together while operating 11.6 times faster than a state-of-the-art task grouping method. 

Efficiently Identifying Task Groupings for Multi-Task Learning
Christopher Fifty! Ehsan Amid Zhe Zhao', Tianhe Yu!
Rohan Anil', Chelsea Finn Google Research, Brain Team", Stanford University?
of ofiiftygooglee com
Abstract
Multi-task learning can leverage information learned by one task to benefit the training other tasks. Despite this capacity, naively training all tasks together in one model often degrades performance, and exhaustively searching through combinations of task groupings can be prohibitively expensive. As a result, effi- ciently identifying the tasks that would benefit from training together remains a challenging design question without a clear solution. this paper, we suggest an approach to select which tasks should train together in multi-task learning models. Our method determines task groupings in a single run by training all tasks together and quantifying the effect to which one task's gradient would affect another task's s loss. On the large-scale Taskonomy computer vision dataset, we find this method decrease test loss by 10.0% compared to simply training all tasks together while operating 1 1.6 times faster than a state-of-the-ar task grouping method.
Introduction
Many of the forefront challenges in applied machine learning demand that a single model performs well on multiple tasks, optimizes multiple objectives while simultaneously adhering unmovable inference-time constraints. For instance, autonomous vehicles necessitate low inference time latency to make multiple predictions real-time video feed to precipitate driving action [27]. Robotic arms are asked to concurrently learn how to pick, place, cover, align, and rearrange various objects improve learning efficiency [25], and online movie recommendation systems model multiple engagement metrics to facilitate low-latency personalized recommendations [13]. Each of the above applications depends on multi-task learning, and advances which improve multi-task learning performance have the potential to make an outsized impact on these and many other domains.
Multi-task learning can improve modeling performance by introducing an inductive bias prefer hypothesis classes that explain multiple objectives and by focusing attention on relevant features [43]. However, it may also lead severely degraded performance when tasks compete for model capacity or are unable to build a shared representation that can generalize to all objectives. Accordingly, finding groups of tasks that derive benefit from the positives of training together while mitigating the negatives often improves the modeling performance multi-task learning systems.
While recent work has developed new multi-task learning optimization schemes [28, 10, 45, 53, 11, 50], the problem deciding which tasks should be trained together in the first place is an understudied and complex issue that is often left to human experts [56]. However, human's understanding of similarity is motivated by their intuition and experience rather than a prescient knowledge of the underlying structures learned by neural network. To further complicate matters, the benefit or detriment induced from multi-task learning relies on many non-trivial decisions including, but not limited to, dataset characteristics, model architecture, hyperparameters, capacity, and convergence [51, 49, 46, 47]. As result, a systematic technique to determine which tasks
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

I: Train All Tasks In One Multi-Task Network
I: Compute Inter-Task Affinities
II: Network
IV: Train Networks for Each Group of Tasks
Selection
Segmeniation
Keypaints
Segmenzatinn Keypoins
Segmamiation
Kaynuints
Taska
Select Networks
Cdges
Numuak
Dapih
Edues
Depth
That Maximize
Edoas
Dopeth
Nomals
Inter-Task Affinity
Mult- Mulii-Takk Taak Modal
Topat Imaage
Figure 1: Overview of our suggested approach to efficiently determine task groupings. (I): Train all tasks
together in a multi-task learning model. (II): Compute inter-task affinity scores during training. (III): Select multi-task networks that maximize the inter-task affinity score onto each serving-time task. (IV): Train the resulting networks and deploy to inference.
should train together in a multi-task neural network would be valuable to practitioners and researchers alike [5, 6].
One approach to select task groupings is to exhaustively search over the 2/71 multi-task networks' for set of tasks T. However, the cost associated with this search can be prohibitive, especially when there is large number of tasks. It further complicated by the fact that the set tasks to which a model is applied may change throughout its lifetime. As tasks are added to dropped from the set of all tasks, this costly analysis would need to be repeated to determine new groupings. Moreover, model scale and complexity continues increase, even approximate task grouping algorithms which evaluate only subset of combinations may become prohibitively costly and time-consuming to evaluate.
In this paper, we aim to develop an efficient framework to select task groupings without sacrificing performance. We propose measure inter-task affinity by training all tasks together in single multi-task network and quantifying the effect to which one task's gradient update would affect another task's loss. This per-step quantity is averaged across training, and tasks then grouped together to maximize the affinity onto each task. visual depiction of the method shown in Figure 1. Our suggested approach makes assumptions regarding model architecture and is applicable any paradigm which shared parameters are updated with respect to multiple losses.
In summary, our primary contribution is to suggest a measure of inter-task task affinity that can be used to systematically and efficiently determine task groupings for multi-tast learning. Our theoretical analysis shows that grouping tasks by maximizing inter-task affinity will outperform any other task grouping in the convex setting under mild conditions. Further on two challenging multi-task image benchmarks, our empirical analysis finds this approach outperforms training all tasks independently, training all tasks together (with and without training augmentations), and is competitive with a oite-offthh-art task grouping method while decreasing runtime by more than an order of magnitude.
2 Related Work
Task Groupings. Prevailing wisdom suggests tasks which are similar or share a similar underlying structure may benefit from training together in multi-task system [9, 8, 4]. Early work in this domain pertaining to the convex setting assume all tasks share common latent feature representation, and find that model performance can be significantly improved by clustering tasks based on the basis vectors they share in this latent space [26, 30]. However, early convex methods to determine task groupings often make prohibitive assumptions that do not scale to deep neural networks.
Deciding which tasks should train together in multi-task neural networks has traditionally been addressed with costly cross-validation techniques or high variance human intuition. An altogether
'here
(T)) single-task networks, (!77) 2-task multi task networks. which yields Σ.7 (77 2!!1 -

---

different approach may leverage recent advances in transfer learning focused on understanding task relationships [54, 3, 15, 58, 2}: however, [46] show transfer learning algorithms which determine task similarity do not carry over to the multi-task learning domain and instead propose a multi-task specific framework which trains between (1z)) [7] and 2/7) - 1 models to approximate exhaustive search performance. Our approach differs from [46] in that it computes task groupings from only a single training run.
Architectures and Training Dynamics. A plethora of multi-task methods addressing what pa- rameters to share among tasks in a model have been developed, such as Neural Architecture Search [19, 47, 49, 35, 22, 39], Soft-Parameter Sharing [40, 14, 52], and asymmetric informa- tion transfer [31, 44, 32] to improve multi-task performance. Although this direction is promising, we direct our focus towards when to share tasks in a multi-task network rather than architecture modifications to maximize the benefits from training fixed set of tasks together. Nevertheless, both approaches are complementary, and architecture augmentations seem to perform best when trained with related tasks [43].
Significant effort has also been invested to improve the optimization dynamics of MTL systems. In particular, dynamic loss reweighing has achieved performance superior to using fixed loss weights found with extensive hyperparameter search [28, 18, 37, 10, 45, 33]. Another set of methods seek to mitigate inter-task conflict by manipulating the direction of task gradients rather than simply their magnitude [48, 57, 53, 11, 50, 36]. our experiments, we compare against Uncertainty Weights [28], GradNorm [10], and PCGrad [53] to contextualize the relative change in performance from splitting tasks into groups. We find that task grouping methods outperform all three training augmentations; nonetheless, they are naturally complementary. In Section 5, our results indicate enhancing the networks found by our method with PCGrad can lead to additional improvements in performance.
Looking into the Future. "Lookahead" methods in deep learning can often be characterized by saving the current state of the model, applying one or more gradient updates to subset of the parameters, reloading the saved state, and then leveraging the information learned from the future state to modify the current set of parameters. This approach has been used extensively in the meta- learning [16, 42, 7, 17, 291, optimization [41, 21, 55, 23, 24]. and recently auxiliary task learning domains [34]. Unlike the above mentioned methods which look into the future to modify optimization processes, our work adapts this central concept to the multi-task learning domain to characterize task interactions and assign tasks to groups of networks.
3 Task Grouping Problem Definition
We draw a distinction between inference-time latency constraints and inference-time memory budget. The former characterizes the speed at which predictions can be computed, with similarly sized models running in parallel having latency roughly equivalent to a single model running by itself. The latter relates to the number of parameters used by all models during inference, with a n-times parameter model having a similar budget to an n- n-group of normal-sized models. We configure our analysis to span both dimensions, but also provide analysis into only the latter in the Appendix.
Given a set of tasks T, a fixed inference-tine memory budget and latency constraint c, our aim is to assign tasks to networks such that average task performance is maximized. Additionally, each network must have a parameter count less than e; the networks must span our set of tasks T; and the total number of networks is less than or equal to our memory budget b. Moreover, tasks can be trained a network without being served from this network during inference. In this case, they are used assist the learning of the other tasks, and during inference, are served from one of the other multi-task networks. Formulating our task grouping framework in this manner aligns our work with industry trends where inference- time parameter budgets are limited and latency unmovable.
More formally, for a set of n tasks T = {πi, T2,... Jn we would like to construct group of multi-task neural networks M = (m1, m2,.... ..., m&} such that VT; E T, 3 exactly one multi-task network m, E M, parameter count of m, < c, such that my makes an inference-time prediction for 1, subject to S where b our memory budget. my is the jth multi-task network which takes input X, and concurrently trains a set of tasks (Ta, Te..... J. }. but only serves a subset of those tasks at inference. For given performance measure P. we can then define the aggregate performance of our task grouping En= PTT[[M) where MM) computes the performance of task T; from the
3

---

set of models M using the model my which the task grouping algorithm predicts the performance of T, will be highest.
4 Grouping Tasks by Measuring Inter-Task Affinity
We propose method to group tasks by examining the effect to which one task's gradient would increase or decrease another task's loss. We formally define the method Section 4.1, describe systematic procedure to go from inter-task affinity scores to grouping of tasks in Section 4.2, and provide theoretical analysis in Section 4.3.
4.1 Inter-Task Affinity
Within the context of a hard-parameter sharing paradigm, tasks collaborate to build a shared feature representation which is then specialized by individual task-specific heads output a prediction. Specifically, through the process of successive gradient updates to the shared parameters, tasks implicitly transfer information each other. As consequence, we propose to view the extent to which task's 's successive gradient updates on the shared parameters affect the objective of other tasks in the network as a proxy measurement of inter-task affinity,
Consider a multitask loss function parameterized by {0.} U {0.1 E T where 0g represents the shared parameters and 8, represents the task T specific parameters. Given batch of examples X, let
Luual (X, 0s (0.)) L.
iET
denote the total loss where I, represents the non-negative loss of task i. For simplicity of notation, we set the loss weight of each task to be equal to I, though our construction generalizes arbitrary weightings.
For given training batch X at time-step define the quantity gf+1 to represent the updated shared djs parameters after a gradient step with respect to the task Assuming stochastic gradient descent for simplicity, we have
g = 0 Vo L((X", 0%, 0.)
We calculate a lookahead loss for each task by using the updated shared parameters while keeping the task-specific parameters as well as the input batch unchanged. That is, in order to assess the effect of the gradient update of task i on a given task we compare the loss of task before and after applying the gradient update from task onto the shared parameters. To eliminate the seale discrepancy among different task losses, gradient step on the shared parameters as a scale invariant measure of relative progress. We can then define asymmetric measure for calculating the affinity of task
consider the ratio of a task's 's loss before and after the
given time-step / on task j as (X 1 03)
E =
(1)
T., (X's , 0G 05)
Notice that positive value of E|, indicates that the update on the shared parameters results in lower loss task j than the original parameter values, while negative value of 2 indicates that the shared parameter update is antagonistic for this task's performance. Our suggested measure of inter-task affinity is computed at per-step level of granularity, but be averaged across all steps, every n steps, or a contiguous subset of steps to derive a "training-level" score:
2 T z
In Section 5, we find 2nn is empirically effective in selecting high performance task groupings and provide ablation study along this dimension in Section 5.2.
4.2 Network Selection Algorithm
At a high level, our proposed approach trains all tasks together in one model, measures the pairwise task affinities throughout training, identifies task groupings that maximize total inter-task affinity,

---

and then trains the resulting groupings for evaluation on the test set. We denote this framework as Task Affinity Grouping (TAG) and now describe the algorithm to convert inter-task affinity scores into a set of multi-task networks. More formally, given (12) values representing the pairwise inter- task affinity scores collected during a single training run, our network selection algorithm should produce multi-task networks, k I b where b is the inference- memory budget, with the added constraint that every task must be served from exactly one network at inference.
For group composed of pair of tasks {a, b), the affinity score onto task a would simply be 2,sa and the affinity score onto task b would be 2a-b- For task groupings consisting of three more tasks, we approximate the inter-task affinity onto a given task by averaging the pairwise affinities onto this given task. Consider the group consisting of tasks {a, b, c). We can compute the total inter-task affinity onto task a by averaging the pair-wise affinities from tasks and e onto a: (Z + 20
After approximating higher-order affinity scores for each network consisting of three or more tasks, we select a set of k multi-task networks such that the total affinity score onto each task used during inference is maximized. Informally, each task that is being served at inference should train with the tasks which most decrease its loss throughout training. This problem NP-hard (reduction from -a----n---uddddiii algorithm detailed in [46] or
Set-Cover) but can be solved efficiently with with a binary integer programming solver as done by [54].
4.3 Theoretical Analysis
We now offer a theoretical analysis of our measure of inter-task affinity. Specifically, our goal is to provide an answer to the following question: given that task b induces higher inter-task affinity than task e on task a, does training {a, b} together result in a lower loss on task a than training {a, c)? Intuitively, we expect the answer to this question to always be positive. However, it is easy to construct counter examples for simple quadratic loss functions where training {a, 6} actually induces a higher loss than training {a, c} (see Appendix). Nonetheless, we show that in convex setting and under some mild assumptions, the grouping suggested by our measure of inter-task affinity must induce lower loss value on task a.
For simplicity of notation, we ignore the task specific parameters and use L, (6) to denote the loss of task evaluated at shared parameters 0. Additionally, we denote the gradient of tasks a, b, and at 0 as ga- Ib. and Gc, respectively.
Lemma 1. Let La be a-sttrongly convex and B-strongly smooth loss function. Given that task b induces higher inter-task affinity than task c on task a, the following inequality holds:
Bn |1
an |1 S ga 9b
9c
+
(2)
2
2
The proof is given in the Appendix.
Proposition 1. Let La be a-strongly convex and 8-strongly smooth loss function. Let 1) S B be the learning rate. Suppose that in given step, task b has higher inter-task affinity than task on task a. Moreover, suppose that the gradients have equal norm, i.e. [| 1196 llgell. Then, taking a gradient step on the parameters using the combined gradient of a and h reduces Taa more so than taking a gradient step on the parameters using the combined gradient of a and c, given that cos(ga.
< n 8 1 where cos(u. u, v) = u.v is the cosine similarity between u and v. 1 B-a
[||l
The proof is given in the Appendix. Proposition 1 intuitively implies the grouping chosen by maximizing per-task inter-task affinity is guaranteed to make more progress than any other group. Moreover, the assumptions for this condition to hold in the convex setting are fairly mild. Specifically, our first assumption relies on the gradients to have equal norms, but our proof can be generalized to when the gradient norms are approximately equal. This is often the case during training when tasks use similar loss functions and/or compute related quantities. The second assumption relies on the cosine similarity between the lower-affing affinity task e and the primary task a being smaller than a constant. This is mild assumption since the ratio trivially when the Hessian of the loss has sufficiently small condition number.
2/e 8/o-1 becomes sufficiently large for the assumption to hold
5

---

Train Groupings
Train Groupings Taskanomy
2048
Testa VI00 GIRUI Hours
TAG
TMG
1G
1G
104
GIPU
1850 Touro:02 hours
M04
BH Neurs
512
Teriv Hours
Teeaai00
Ted
252
Tedior100 Hoon
128
Two Groups
Thece Groupe Number Splits
Pour Groups
Twos Groupa
Tires Groups
Number Splits
Figure 2: (Left) time to determine and train task groupings in CelebA. (Right) time to determine and train task groupings in Taskonomy. Note the y-axis is in log scale, and the time to determine task groupings is incurred only once to determine groupings for all splits.
Experiments
We evaluate? the capacity of TAG to select task groupings CelebA, large-scale face attributes dataset [38] and Taskonomy, a massive computer vision dataset of indoor scenes [54]. Following this analysis, direct our focus towards answering the following questions with ablation experiments on CelebA:
Does our measure of inter-task affinity align with identifying which tasks should train together?
Should inter-task affinity be measured at every step of training determine task groupings?
Is measuring the change in train loss comparable with the change in validation loss when computing inter-task affinity?
How do inter-task affinities change over the course of training?
Do changes in model's 's hyperparameters change which tasks should be trained together?
As described in Section 3, we constrain all networks to have the same number of parameters to adhere fixed inference-tinn latency constraint. We provide experimental results removing this constraint, as well as additional experimental results and detail relating to experimental design, in the Appendix.
5.1 Supervised Task Grouping Evaluation
For our task grouping evaluation, we compare two classes of approaches: approaches that determine task groupings, and approaches that train on all tasks together but alter the optimization. In the first class, we consider simply training all tasks together in the same network (MTL), training every task by itself (STL), the expected value from randomly selecting task groupings (RG), grouping tasks by maximizing inter-task cosine similarity between pairs of gradients (CS), our method (TAG ), and HOA [46] which approximates higher-order task groupings from pair-wise task performance. For the latter class, we consider Uncertainty Weights (UW) [28], GradNorm (GN) [10], and PCGrad [53]. In principle, these two classes of approaches are complementary and can be combined.
Our empirical findings are summarized in Figure 3. For the task grouping methods TAG,, CS, and HOA, we report the time to determine task groupings, not determine task groupings and train the resultant multi-task networks. The runtime of RG fixed to the time it takes train a single multi- task network, and the runtime of STL is reported as the time it takes to train all single task networks. This is to facilitate comparison between the efficiency different task grouping methods and provide high-level overview of how long it takes to select task groupings compared to popular multi-task learning benchmarks. We also include a comparison the resultant multi-task networks among task grouping methods in Figure 2.
time to determine task groupings and train
CelebA. We select subset of attributes (al, a2, a3, a4, a5, a6, a7, a8, a9) from the 40 possible attributes in CelebA and optimize the baseline MTL model by tuning architecture, batch size, and
"Our code is available wtw.hh/wwwooooogoeeoooooooooooee

---

Task Groupings in CelebA
Task Groupings in Taskonomy
Spitts Three Splits
Spitss Three Spilts
(8.5x)
RG
50.0
11.0xt
Four Splits
G4
Total Test Error
50.6
12. 05x0
Total Loss
0.56
60.4
50.2
NTL
GN (1.9k]
PCGras
0.54
1AG
CS
13.8kl
A
50.0
UW
0.52-
(15K))
(1.04kl
51.4 Testevi00 louus
TAG
0.50
49.6
V10O
10a
16* Relative Runtime
10a
307
Relative Runtime
Figure 3: (Left) average classification error for 2, 3, and 4-split task groupings for the subset of 9 tasks in CelebA. (Right) total test loss for 2 and 3-split task groupings for the subset of 5 tasks in Taskonomy. All models were run on a TeslaV 100 instance with the time to train the full MTL model being approximately 83 minutes in CelebA and 146 hours in Taskonomy. Note the x-axis is in log scale, and the relative runtime for TAG, CS, and HOA only considers the time to find groups
learning rate to maximize the performance of training all tasks together on the validation set. We do not tune other methods with the exception of GradNorm, for which we search over {0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0} for alpha. For task grouping algorithms, we evaluate the set of {2-splits, 3-splits, 4-splits) inference-tinee memory budgets. Each of the multi-task networks has the same number of parameters as the base MTL model; however, the inference-time latency constraint is satisfied as all models may run in parallel. Our findings are summarized in Figure (left). Alternatively, increasing the number of parameters within a single model may significantly increase the time to make a forward pass during inference. Figure 6 visualizes our findings when we remove the inference-time latency constraint from comparison systems.
We find the performance of TAG surpasses that of the HOA, RG, and CS task grouping methods, while operating 22 times faster than HOA. We also find UW, GN, and PCGrad to perform worse than the groups found by either HOA or TAG, suggesting the improvement from identifying tasks which train well together cannot be replaced with current multi-task training augmentations. Nevertheless, we find multi-task training augmentations can be complementary with task grouping methods. For example, augmenting the groups found by TAG with PCGrad improves 2-splits performance by 0.85%, 3-splits performance by 0.18%, but changes 4-splits performance by -0.08%.
Similar to the results from [46], we find GradNorm [10] can sometimes perform worse than training all tasks together. We reason this difference is due to our common experimental design that loads the weights from the epoch with lowest validation loss to reduce overfitting, and differs from prior work which uses model weights after training for 100 epochs [45]. Reformulating our design to mirror [45], we find training the model to 100 epochs without early stopping results in worse performance on MTL, GradNorm, UW, and PCGrad; however, with this change, each training augmentation method now significantly outperforms the baseline MTL method.
Taskonomy. Following the experimental setup of [46], we evaluate the capacity of TAG to select task groupings on the "Semantic Segmentation", 'Deptth Estimation", "Keypoint Detection", "Edge Detection", and "Surface Normal Prediction" " objectives in Taskonomy. Unlike [46], our evaluation uses an augmented version of the medium Taskonomy split (2.4 TB) as opposed to the Full+ version (12 TB) to reduce computational overhead and increase reproducibility.
Our results are summarized in Figure 3 (right). Similar our findings on CelebA, TAG continues to outperform MTL by 10.0%, GN by 7.7%, STL by 1.5%, and RG by 9.5%. Comparing TAG to HOA, we find the 2-split task grouping found by TAG to surpass the performance of that found by HOA by 2.5%, but HOA's 3-split task grouping performance is superior to that of TAG. In terms of compute, TAG is significantly more efficient than HOA, with HOA demanding an additional 2,008 1100 GPU hours to find task groupings. To put this cost into perspective, on an 8-GPU, on-demand p3. 16laare AWS instance, the difference in monetary expenditure between TAG and HOA would be $6, 1 144.48. On a similar note, the performance of TAG and CS on Taskonomy are equivalent, but TAG is more efficient, requiring 140 fewer Tesla 1000 GPU hours to compute task groupings.
7

---

5.2 Multi-Task Task Ablation Studies
Does our measure of inter-task affinity correlate with op- timal task groupings? To further evaluate if TAG can be used to select which tasks should train together in multi-task learning models, we evaluate its capacity to select the best training partner for a given task. We compare with the optimal (best) and worst auxiliary task computed from the test set and normalize scores with respect to the expected performance of selecting an auxiliary task at random.
Improvement in Test Accuracy Relative to Random Grouping
Tasks
optimal
(ours)
werst
a1 2.60%
2.6%
+3.03%
a2 1.53%
1.29%
-1.95%
a3
2.37%
1.72%
.304%.
a4
2.67%
0.72%
44.4%
a5
2.75%%
2.291
-3.8%%
a6 2.04%
-2.08%.
a7 2.40%
0.74%
-2.43%
Our findings are summarized in Table and indicate the per- formance of auxiliary tasks found by TAG correlates with the performance of the optimal auxiliary task (Pearson's 's Corre- lation of 0.93%). However, a notable exception occurs with attribute a8 where TAG actually selects the worst partner. In this instance, and unlike every other objective in our dataset, no task manifests especially high or low inter-task affinity onto a8. The difference in normalized inter-task affinity for a8 between the best and worst partner predicted by inter-task affinity is 0.04, while the next smallest difference is 4x larger at 0. for a3. A table showing these differences is included in the Appendix. This case represents a limitation of TAG, and it will struggle to find the best auxiliary task for objectives like a8.
a8 1.61%
-1.59%
-1.59%
29 8.58%
8.38%
6.21%
Table 1: Performance of each task when trained with it's partner task relative to the expected performance from random groupings.
However, this weakness does not seem significantly affect the capacity of TAG to select strong task groupings. As no other task exhibits especially high or low inter-task affinity onto a8, it is often slotted into a larger task group rather than being one of the tasks with a large difference in inter-task affinity which precipitates the formation of a new group.
Should inter-task affinity be computed every step?
likely that the inter-task affinity of consecu- tive steps are similar, and it could be the case that inter-task affinity signals at the beginning. mid- dle, end of training are sufficient for selecting which tasks should train together. In particular, if task relationships crystallize early in training, computing inter-task affinities during the initial stages of training may be sufficient. We evaluate both hypotheses on CelebA and our results are summarized in Table 2.
Method
Relative Performance
Relative Specdup
Every I Step
5.13%
1.0x
Every Steps
5. 5.13%
2.56x
Every Steps
5.13%
3.19x
Every 25 Steps
4.84%
3.73A
Every Steps
3.73%
3.96x
Every 100 Steps First 25%
2,06%
4,08A.
3.67%
4.00x
Middle 25%
4.31%
2.95x
Final 25%
4.02%
2.34x
Table 2: Change in total test accuracy across inference-time memory budget of (2-groups. groups. 4-groups) relative to the expected perfor- mance from random groupings. Speedup is relative to computing inter-task affinity in each step.
Our findings indicate significant redundancy can be climinated without degrading performance by computing inter-task affinities every 10 steps rather than every step. This modification increases training- maining-time efficiency by 319% and is used in Section 5.1 to compute task groupings. After this threshold, signal strength decreases and leads to higher task grouping error. We also find that comput- ing inter-task affinities in the first 25%, middle 25%, or final 25% of training degrades task- aask-grouping performance. This result suggest the relationships among tasks change throughout training as mea- sured by inter-task affinity. As result, we choose to average inter-task affinity scores throughout the entirety of training determine which tasks should train together.
Is change in train loss comparable to change in validation loss? Given multi-task learning's 's capacity to improve generalization, computing the change in validation loss after gradient step may capture a more informative signal as to which tasks should train together. On the other hand, certain datasets may not contain validation split and loading batch from the validation set every 10 steps of training would decrease efficiency.
To our surprise, the taser-taskan affinity scores computed on the validation set are very similar to the inter-task affinity scores computed on the training set (Pearson's Coefficient: 0.9804). For context, this similarity with computing inter-task affinities every step is greater than any other ablation in Table 2 with the exception of "Every Steps" as measured by Pearson's Coefficient. Moreover, the performance of groupings found by both methods are similar; 49.574 average total error across our
8

---

Affinity
Inter-Task Affinity
Affinity
Per-Epoch Inter-Task Affinity
15
Number of Epochs Before Early Stopping
Figure 4: Total per-epoch inter-task affinities onto tasks [al, a5, a6) in the CelebA dataset. The y-axis signifies the total change in train loss after applying an update to the shared parameters. Note the relationship among affinities changes. For instance in the center plot representing the inter-task affinity onto a5, a6 initially manifests higher affinity onto a5 than does a7. later stages of training, this trend is reversed. In the same subplot, al manifests higher inter-task affinity over the tasks {a3, a4, a8, } at different stages of training.
inference time budget of [2- (2-groups, 3-groups, 4- groups) compared with 49.576 for groupings found on the validation set.
More formally, let X, and X be independent and identically distributed random variables for training and validation, respectively. Given the updated shared parameter 6++1 using the gradient of task i calculated on X the loss of task j (s.t. j * i) yields similar expectations with respect to X, and X That is,
E. LE 0 00)1 se E E. 1.110111.00) E,,C, 10 11,)1
The leftmost term is joint expectation w.r.t. both X,, and X in which X,, is only used to calculate the updated shared parameters gt+1 Assuming both tasks have distinct loss functions that do not directly depend on each other during training, we can move the second expectation inside the loss function to obtain the middle term. By using the fact that X,, and X are identically distributed, thus yielding the same expectation, we can move from the middle term to the rightmost term. Hence, the inter-task affinity computed on the training dataset would approximately equal the inter-task affinity computed on the validation set.
a/:
How do inter-task affinities change over the course of training? Our analysis on CelebA and Taskonomy suggests inter-task affinities change throughout training, but no obvious patterns emerge related to how inter-task affinities change a function of time across tasks. Nevertheless, our analysis indicates certain tasks exhibit higher-than higher-than-average inter-task affinity and tend to maintain this difference throughout training. We visualize this effect in Figure 4 for three tasks in the CelebA dataset. Figure 4(left) shows that task a7 exhibits significantly higher affinity onto task al than any other task in the network, and in similar comparison, Figure 4(right) indicates that tasks (a5, a7) manifest significantly higher affinity onto a6 than any other task. Shifting our focus to a5 in Figure 4(center), we find (a6, a7) display markedly higher affinity, and sometimes displays higher affinity, onto a5 than the other tasks in the network.
One additional observation not captured in Figure 4 relates the initial steps of training. At this stage of convergence, all tasks seem to manifest positive and similar inter-task affinity. We postulate the representations learned by the model at this early stage may be common among all tasks, For example, the learned representations may extract general-purpose facial characteristics, but specialization occurs quickly thereafter.
Can changes in hyperparameters affect task groupings? We define three settings: (i) typical, the base setting, (ii) b=0.5x, halving the batch size, and (iii) lr=2x, increasing the learning rate by a factor of 2. To determine the ground-truth task groupings, we train all 511 combinations of task groupings from our 9-task subset of CelebA and select the groups in each setting with the lowest total test error. Our aim is to assess the extent to which task groupings chosen in setting b=0.5x or Ir=2x generalize the typical setting.

---

Our results are summarized in Table 3. They indicate that simply changing the batch size or learning rate of model may change which tasks should be trained to- gether, with the groupings found by 1r=2x exhibiting worse generalization than those found by b=0.5x. This result suggests that how tasks should be trained together does not simply depend on the relationships among tasks, but also on detailed aspects of the model and training. is notably difficult to build intuition for the latter, illus- trating the need to develop automated methods that can take into account these nuances.
Improvement in Test Accuracy Relative to Optimal Groupings
Budget
b=0.5x
Ir=2x
2- 2-groups 3 3-groups
-0.61 %
-1.22%
-0.25%
-2.90%
4 groups
-1.87%
-3.21%
Table 3: Accuracy of task groupings found by b=0.5x and Ir=2x relative to the typical sctting.
Conclusion
In this work, we present an approach to quantify inter-task affinity in single training run and show how this quantity can be used to systematically determine which tasks should train together in multi- task networks. Our empirical findings indicate our approach is highly competitive. outperforms multi-task training augmentations like Uncertainty Weights, GradNorm, and PCGrad, and performs competitively with tate-offthe--ae task grouping methods like HOA, while improving computational efficiency by over an order of magnitude. Further, our findings are supported by extensive analysis that suggests inter-task affinity scores can find close to optimal auxiliary tasks, and in fact, implicitly measure generalization capability among tasks.
A plethora of research has been undertaken design better multi-task learning architectures, or improve the optimization dynamics within multi-task learning systems, with relatively little work addressing the question of which tasks should train together in the first place. It our hope this work renews interest in this domain, and given the sensitivity of task groupings to even small changes in hyperparameters, encourages the development of efficient and automatic methods identify which tasks should train together in multi-task learning networks.
7 Broader Impact
Efficiently identifying task groupings in multi-task learning has the potential save significant time and computational resources in both academic and industry environments. Despite this benefit, there are several risks associated with this work. In particular, inter-task task affinities can be mistakenly interpreted as "task similarity", and incorrectly create an association and/or causation relationship among tasks with high mutual inter-task affinity scores. This association would be especially problematic for datasets involving sensitive prediction quantities related to race, gender, religion, age, status, physical traits, etc., where inter-task affinities could be mistakenly used to support an unfounded conclusion that attempts to posit similarity among tasks. That said, we believe acknowledging these risks mitigates their potential for abuse, and the benefit from this work most notably decreasing computational resources by over an order of magnitude compared with a ofate-of-the-arr task grouping method while performing competitively in terms of accuracy merits its dissemination.
Acknowledgement
We thank Vince Gatto, Bing-Rong Lin, Nick Bridle, Li Wei, Shawn Andrews, Yuan Gao, and Yuyan Wang for their thoughtful discussion related to a precursor of this work. We also thank James Chen and Linda Wang for providing feedback on earlier draft of this paper. Lastly, we would like to recognize Zirui Wang for technical support with running experiments as well as our anonymous NeurIPS reviewers for their thoughtful feedback and clear desire to improve this work.
References
[1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for
10

---

large-scale machine learning. In 12th {USENIX) symposium on operating systems design and implementation ({OSDI) 16), pages 365-283,20 2016.
[2] Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji, Charless C Fowlkes, Stefano Soatto, and Pietro Perona. Task2vec: Task embedding for meta- learning. In Proceedings of the IEEE International Conference on Computer Vision, pages 1330-6439, 2019.
[3] Alessandro Achille, Giovanni Paolini, Glen Mbeng, and Stefano Soatto. The information complexity of learning tasks, their structure and their distance. arXiv preprint 2019.
[4] Andreas Argyriou, Theodoros Evgeniou, and Massimiliano Pontil. Convex multi-task feature learning. Machine learning, 73(3):243-272, 2008.
[5] Jonathan Baxter. A model of inductive bias learning. Journal of artificial intelligence research, 12:149-198, 2000.
[6] Shai Ben-David and Reba Schuller. Exploiting task relatedness for multiple task learning. In Learning Theory and Kernel Machines, pages 567-580. Springer, 2003.
[7] Lukas Brinkmeyer, Rafael Rego Drumond, Randolf Scholz, Josif Grabocka, and Lars Schmidt- Thieme. Chameleon: Learning model initializations across tasks with different schemas. arXiv preprint arXiv:1909.133: 2019.
[8] Rich Caruana. Multitask leamning Machine Learning. 28, pages 41-75, 1997.
[9] Richard Caruana. Multitask learning: A knowledge- source inductive bias. In Pro- ceedings of the Tenth International Conference Machine Learning, pages 41-48. Morgan Kaufmann, 1993.
[10] Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gra- dient normalization for adaptive loss balancing in deep multitask networks. arXiv preprint arXiv:1711.02235
2017.
[11] Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. Just pick sign: Optimizing deep multitask models with gradient sign dropout. Advances in Neural Information Processing Systems, 33, 2020.
[12] François Chollet et al. Keras: The python deep learning library. ascl, pages ascl-1806, 2018.
[13] Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommenda- tions. In Proceedings of the 10th ACM conference on recommender systems, pages 191-198, 2016.
[14] Long Duong, Trevor Cohn, Steven Bird, and Paul Cook. Low resource dependency parsing: Cross-lingual parameter sharing in neural network parser. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 2: Short Papers), pages 845- 850, 2015.
[15] Kshitij Dwivedi and Gemma Roig. Representation similarity analysis for efficient task taxonomy transfer learning. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 12387-12396. 2019.
[16] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In Proceedings of the 34th International Conference on Machine Learning- Vearning-Vollue 70, pages 1126-1135, JMLR. org, 2017.
[17] Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient-based meta-learning as hierarchical bayes, 2018
[18] Michelle Guo, Albert Haque, De-An Huang, Serena Yeung, and Li Fei-Fei. Dynamic task prioritization for multitask learning. In Proceedings of the European Conference on Computer Vision (ECCV), pages 270-287, 2018.

---

[19] Pengsheng Guo, Chen- Yu Lee, and Daniel Ulbricht. Learning to branch for multi-task learning. arXiv preprint arXiv:2006. 01895, 2020.
[20] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and paltern recognition, pages 770-778, 2016.
[21] Geoffrey Hinton and David C Plaut. Using fast weights to deblur old memories. In Proceedings of the ninth annual conference of the Cognitive Science Society, pages 177-186, 1987.
[22] Siyu Huang, Xi Li, Zhi-Qu Qi Cheng, Zhongfei Zhang, and Alexander Hauptmann. GNAS: A greedy neural architecture search method for multi-attribute learning. 2018 ACM Multimedia Conference Multimedia Conference MM '18, 2018
[23] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407. 2018.
[24] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances in neural information processing systems, pages 315-323, 2013.
[25] Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic reinforcement learning at scale. arXiv preprint arXiv:2 2021.
[26] Zhuoliang Kang, Kristen Grauman, and Fei Sha. Learning with whom to share in multi-task feature learning. In ICML, 2011.
[27] Andrej Karpathy. Multi-task learning in the wilderness. ICML, 2019. URL https: // slideslive. com/ 38917 7600 httm/muutaaattseee waaaa--leernngg Learning- ien-iinthhhh the Setheewwleee wilderness.
[28] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semanties. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7482-7491, 2018.
[29] Taesup Kim, Jaesik Yoon, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian -anel-ag meta-learning. arXiv preprint arXiv: 03836, 2018
[30] Abhishek Kumar and Hal Daume III. Learning task grouping and overlap in multi-task learning. arXiv preprint arXiv: 1206. 6417, 2012.
[31] Giwoong Lee, Eunho Yang, and Sung Hwang. Asymmetric multi-task learning based on task relatedness and loss. In International conference on machine learning, pages 230-238. PMLR, 2016.
[32] Hae Beom Lee, Funho Yang, and Sung Ju Hwang. Deep asymmetric multi-task feature learning. In International Conference on Machine Learning, pages 2956-2964 PMLR, 2018
[33] Xi Lin, Hui-Ling Zhen, Zhenhua Li, Qing-Fu Fu Zhang, and Sam Kwong. Pareto multi-task learning. In Advances in Neural Information Processing Systems, pages 12037-12047, 2019.
[34] Xingyu Lin, Harjatin Singh Baweja, George Kantor, and David Held. Adaptive auxiliary task weighting for reinforcement learning. Advances in neural information processing systems, 32, 2019.
[35] Chenxi Liu, Barret Zoph, Maxim Neumann, Jonathon Shlens, Wei Hua, Li-Jia Li, Li Fei-Fei, Alan Yuille, Jonathan Huang, and Kevin Murphy. Progressive neural architecture search. In Proceedings of the European Conference Computer Vision (ECCV), pages 19-34, 2018.
[36] Liyang Liu, Li, Zhanghui Kuang, Jing-Hao Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Towards impartial multi-task learning. In International Confer- ence on Learning Representations, 2021. URL https: / /ppenrevi ew. net forum? id-IMPP IMP NRXWpV....
12

---

[37] Shikun Liu, Edward Johns, and Andrew J Davison. End-to-end multi-task learning with attention. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1871-1880, 2019.
[38] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. 2015 IEEE International Conference on Computer Vision (ICCV), Dec 2015.
[39] Yongxi Lu, Abhishek Kumar, Shuangfei Zhai, Yu Cheng, Tara Javidi, and Rogerio Feris. Fully-adaptive feature sharing in multi-task networks with applications in person attribute classification. 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Jul 2017.
[40] Ishan Misra, Abhinav Shrivastava, Abhinay Gupta, and Martial Hebert. Cross-stitch networks for multi-task learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3994-4003, 2016.
[41] Yu Nesterov. A method of solving a convex programming problem with convergence rate ((1/k)). In Sov. Math. Dokl, volume 27, pages 372-376, 1983.
[42] Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprini arXiv: /803.0299, 2018.
[43] Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint arXiv:/706.0500 050988 2017.
[44] Andrei Rusu, Neil Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. arXiv preprint arXiv: (6066666667) 2016.
[45] Ozan Sener and Vladlen Koltun, Multi-task learning as multi-objective optimization. In Advances in Neural Information Processing Systems, pages 527-538, 2018.
[46] Trevor Standley, Amir R Zamir, Dawn Chen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? arXiv preprint arXiv:1905.07553 2019.
[47] Ximeng Sun, Rameswar Panda, and Rogerio Feris. Adashare: Learning what to share for efficient deep multi-task learning. arXiv preprint arXiv:1911.12423, 2019.
[48] Mihai Suteu and Yike Guo. Regularizing deep multi-task networks using orthogonal gradients. arXiv preprint arXiv: 1912006444 2019.
[49] Simon Vandenhende, Stamatios Georgoulis, Bert De Brabandere, and Luc Van Gool. Branched multi-task networks: deciding what layers to share. arXiv preprint wrXiv:1904.02920 2019.
[50] Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao. Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models. arXiv preprint arXiv:2010.05874. 2020.
[51] Sen Wu, Hongyang Zhang, and Christopher Ré. Understanding and improving information transfer in multi-task learning, arXiv preprint arXiv:2005.0094 2020.
[52] Yongxin Yang and Timothy M. Hospedales. Trace norm regularised deep multi-task learning. arXiv preprint arXiv: 1606.04036 2016.
[53] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. arXiv preprint arXiv:2001.06782 2020.
[54] Amir R. Zamir, Alexander Sax, William Shen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. 2018 IEEE/CVE Conference, Jun 2018.
[55] Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, step back. In Advances Neural Information Processing Systems, pages 9597-9608. 2019.
13

---

[56] Yu Zhang and Qiang Yang. A survey on multi-task learning, arXiv preprint arXiv:1707.088) 2017.
[57] Xiangyun Zhao, Haoxiang Li, Xiaohui Shen, Xiaodan Liang, and Ying Wu. A modulation module for multi-task learning with applications in image retrieval. In Proceedings of the European Conference on Computer Vision (ECCV), pages 401-416, 2018.
[58] Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE, 2020.
14


Reviews:
[{"review_id": "olqv9n_1ohH", "reviewer": "Reviewer_LXec", "summary": "This paper presents a task grouping approach for multi-task learning performed in a single training step. It measures an inter-task affinity by quantifying the effect of a task’s gradient on other tasks. This work also shows a theoretical analysis of the task grouping by the inter-task affinity under some conditions. Experimental results show that the presented approach gives competitive results compared to other multi-task learning counterparts using two benchmark datasets.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "This work focuses on suggesting an inter-class affinity measure rather than proposing other strategies for multi-task learning such as analyzing dataset characteristics, model architectures, etc. Since this work assumes learning tasks in a single architecture (hard-parameter sharing), identifying similar tasks and grouping them are important tasks to improve learning efficiency. It presents a simple approach to identify task groups based on the response of gradient updates in the shared parameters. This strategy leads to comparable performance for two multi-task learning scenarios.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The idea of measuring inter-task affinity is simple and the technical novelty is limited. The pairwise affinity is measured by comparing the losses before and after applying the gradient update in the shared body of the network. Since it computes all the pairwise task affinities, I think it takes heavy complexity.\n\nIt is unclear and not well described how to split tasks into several groups using the affinity. Based on the score between L152 and L153, how to group tasks, and in L170, how to set k in the method? \n\nThe theoretical analysis seems not practical because it assumes a convex setting with a smooth loss function. If using popular deep neural networks, it is questionable the analysis can still be valid.\n\nIn L242, how the 9 attributes were selected out of 40 attributes in CelebA?\n\nIn Figure 2, I wonder how to collect the results (error, runtime) of the compared methods and whether they are collected under the same experimental setup.\n\nIn Table 2, it is unclear what the value means. An exact measure for the relative performance should be provided to reproduce.\n\nIn Table 3, it is insufficient to show another set of hyperparameters. A control experiment with respect to multiple different values of hyperparameters (batch size, learning rate) should be provided to show the stability of the proposal.\n\n--\n\nUpdate:\nI have read the authors' responses and other reviewers' comments. \nMany of my concerns have been addressed, but reducing the expensive cost for computing pairwise affinity is rather heuristic which is performed by a simple trick, not from the algorithm itself. It would be better to address this point to make the paper stronger.\nBy considering other reviewers' comments and their views on the work, I agree with them, and my comment is not much deviated from others' opinions. The proposed method can be helpful for practitioners due to its simplicity and promising results.\nSo, I raise my rating and wouldn't mind if the paper gets accepted.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3-4"}, "extra_text": {}}, {"review_id": "kCFFjeWeUeP", "reviewer": "Reviewer_swCN", "summary": "This paper proposes a method for grouping tasks to achieve stronger multi-task learning (MTL) performance. The approach is relatively simple and quite efficient in comparison to alternatives: during a single pass of learning, measure the impact that one task's gradient updates would have on other tasks. A theoretical analysis is conducted in a simplified (strongly convex) setting, and experiments show that this approach is as effective as alternatives that are much more computationally intensive. Overall this appears to be a useful technique that could be useful for many other researchers and also potentially for product teams.\n", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths:\n1. The approach is relatively simple, easy to understand, and (relatively) easy to implement.\n2. The approach appears highly effective in the setting where the number of parameters total (across all models) is not constrained.\n3. The experiments are conducted well, and the results are convincing.\n \nWeaknesses:\n1. While the paper makes an argument on the basis of deployed systems that it is reasonable to allow the number of parameters to grow as a function of the number of task groupings, this does introduce a confounding effect in analyzing the results. It also was unclear to me whether, for the comparison systems, if they were also allowed to have increasing numbers of parameters for the number of groups.\n2. The assumption of convexity in the theoretical analysis, while probably the only way to make the result go through, is fairly unsatisfying in MTL settings, especially when the whole motivation is to learn good shared representations.\n\n\n\nI only have a few additional questions/comments:\n\n - I would avoid the term \"co-training\" -- this is its own field of ML, quite different from MTL, and overloading the term is confusing.\n\n - I was confused throughout much of the paper about why the paper claims there are 2^T-1 possible groupings; I thought it should be B_T (the T-th Bell number). It turns out the answer is because tasks can contribute to the training of networks that do not make predictions for that task. This should be made clearer much sooner (eg around line 38).\n \n - l.103 \"combined task performance\" -- I assume this means \"average task performance\" (later text confirms this) but it would be good to be specific.\n \n - I find the terminology \"inference-time budget\" confusing because it's not clear what it's a budget over. I would call it \"parameter budget\" or \"memory budget\" or something like that.\n \n - l.107 I don't understand what \"are not served\" means.\n \n - In the theoretical analysis (4.3) it wasn't clear to me how this interacts with the affinity averaging described around line 168. It would be nice if this were clarified.\n \n - I'm very confused by Figure 2. My understanding is that what is plotted here is the runtime required to find groups (this is also what the caption says). But then why is STL 8.5x and 3.8x, respectively? There is no time required to find groups in STL. Similarly, why does \"Full MTL\" have a non-zero amount of time to \"find groups\"? With those as baselines, I would expect everything else to be infinitely longer. I'm clearly not understanding what is being plotted here, but l.237 syas \"we report the time to determine task groupings, not determine task groupings and train the resultant multi-task networks.\"\n \n - I'd ideally like to see figures showing the full training time (including groupings and training) for all results. The proof of the theorem can be moved to the appendix to make space for this. I understand that the point of this paper is a method to find good groupings, but given that the paper stresses the practical relevance of this approach, what practitioners will really care about is the full time.\n \n - l.246 I don't understand the discussion of 2-splits, 3-splits, 4-splits here. I assume that this means that you're allowed to have 2* as many parameters, 3*, and 4*? Is that right? Please clarify.\n \n - Table 3 body and caption says \"b=128x\" but the text on l.346 and l.355 says \"b=0.5x\". Which is it?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are not discussed. Societal impact is reasonable.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "LRcnTGycmPI", "reviewer": "Reviewer_suuP", "summary": "Authors propose a single run training method to find out what group of tasks should be learnt together based on the network capacity and computational budget. They evaluate TAG on CelebA and Taskonomy datasets.", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "The paper is well written and easy to understand. \n\nThe proposed method solves the task grouping problem very efficiently, which is a very challenging topic in MTL. Previous methods require training many networks to determine the task group while this method only requires a single run, which is very innovative and exciting. I like the method. \n\nThe method accumulates training gradients of different tasks during the single run and computes the task affinity, which sheds some light on the challenging task affinity computation. It is also very insightful in transfer learning.\n\nIn experiment, they evaluate on CelebA and Taskonomy datasets. TAG yields the best relative runtime and total test loss in both scenarios. They also carefully analyze the drawback of their method in Sec5.2 and also carefully discuss the frequency of computing the gradient and which data split should be used. \n\nFor Rebuttal,\nIt would be very interesting that the authors could provide some grouping examples, and show the difference when it is 2 groups, 3 groups and 4 groups. Task affinity matrix will also demonstrate the effectiveness of the method.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, they have mentioned and analyzed very carefully in sec 5.2 ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4hr"}, "extra_text": {}}, {"review_id": "AeSjbCoMrJu", "reviewer": "Reviewer_1wZL", "summary": "Given a large set of tasks, the goal of this work is to identify subsets of these tasks that can be trained together to improve performance and avoid harmful interference. In contrast to multi-task research that focuses on building a model that handles _all_ tasks as well as possible, this work looks at the cues acquired from training such a model to quickly assess which tasks to group together. This is done by seeing whether a task's gradient update helps or hurts the losses of other tasks. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Main thoughts:\n\n- (+) the proposed strategy for measuring task affinity is well justified and straightforward - would be easy enough to implement and replicate. From a first impression it is unclear that it would necessarily correlate with subsequent multi-task performance, but the experiments back up its utility as a grouping signal\n- (+) the case made for efficiency (that this method only requires a single training run of a model with all tasks) is also a strong advantage relative to the cost of other ways to find task groupings\n- (+) clear and well written work with thoughtful analysis throughout paper and in ablations\n- (-) I do not find 9 attributes on CelebA the most compelling multi-task benchmark (though I understand CelebA is a go-to setting in the multi-task literature). The difference in test error is marginal across these methods and facial attribute classification is such a restricted setting. I think it would be better to see results on diverse multi-domain multi-task settings like Visual Decathlon or VTAB. There is extra overhead in using the proposed method in those settings since all task losses cannot be evaluated in a single inference pass, but seems as though it would still be useful.\n\nI lean positive for this paper, I think that the authors propose a nice, simple strategy for grouping tasks in multi-task training and provide a solid set of experiments and analysis to justify its use.\n\nAdditional comments:\n\n- There's an ablation looking at task affinity at different stages of training, the dynamics here seem like they'd be interesting to have a clearer picture of. How does affinity typically change over the course of training, what does it look like for the tasks that match up best, are there different patterns for different pairs of tasks (moving from high -> low affinity or vice versa)\n- Are task affinities robust to the presence of other sets of tasks. Say I trained a model with Tasks A, B, C, D and a model with Tasks B, C, E, F. Would the affinity between B + C be similar in both cases?\n- Is it possible to see the distribution in performance of random groupings? Were the optimal and worst values derived from an exhaustive search or were they the best/worst to appear in the random sampling?\n\nTypo L224 \"hyperparamters\", L263 \"Segmentic Segmentation\"\n\n\n-----\n**Update:** After having read the other reviews and author responses I maintain my original recommendation to accept. I really appreciate the effort in the response along with the attached figures and details that addressed my questions.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I found the \"Broader Impact\" statement solid, much better than what I'm used to coming across in submissions", "review_text": "", "extra_scores": {"ethical_concerns": "none", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for grouping tasks to achieve stronger multi-task learning performance.  All four reviewers believe this is a well-written paper and should get accepted.   Authors need to modify their final version based on reviewers' comments.  
Decision: Accept (Spotlight)