Efficiently Identifying Task Groupings for Multi-Task Learning

Christopher Fifty; Ehsan Amid; Zhe Zhao; Tianhe Yu; Rohan Anil; Chelsea Finn

Source-linked AI summary

Efficiently Identifying Task Groupings for Multi-Task Learning

Christopher Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, Chelsea Finn

arXiv:2109.04617v2Sep 10, 2021cs.LGcs.AIcs.CV

TL;DR

Multi-task learning needs efficient ways to identify beneficial task groupings because training all tasks together can degrade performance and exhaustive grouping searches are expensive. The paper measures inter-task affinity in one joint training run to select groupings, achieving lower loss than joint training and competitive results with greater efficiency.

Problem
Identifying which tasks should train together remains underexplored, despite the performance and computational costs of multi-task grouping choices.
Method
The method trains all tasks together once, measures how each task’s gradient affects another task’s loss, averages these affinities, and groups tasks accordingly.
Results
10.0% lower performance loss than multi-task training, while TAG outperforms several baselines and remains competitive with state-of-the-art grouping methods.
Takeaways & Limitations
Inter-task affinity provides a systematic, computationally efficient basis for selecting task groupings across multi-task learning settings.
Takeaways & Limitations
Inter-task affinity can be mistaken for task similarity or causation, especially for sensitive prediction quantities.

Abstract

from arXiv · show

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.

1 Introduction

Multi-task learning can improve performance but may degrade it when tasks compete, making task grouping both important and difficult. The paper proposes a single-run affinity-based method that efficiently selects task groupings and reports strong benchmark results.

Multi-task learning can improve modeling by encouraging shared representations, but competing tasks may severely degrade performance.
Task grouping is understudied, complex, and often left to human experts despite sensitivity to datasets, architectures, hyperparameters, capacity, and convergence.
Exhaustive search over 2^|T| − 1 networks can become prohibitively costly as task counts, model complexity, or task sets grow.
The method trains all tasks together, measures how one task’s gradient affects another’s loss, averages this affinity across training, and groups tasks to maximize affinity.
Theoretical analysis supports affinity-maximizing groupings in convex settings, while experiments report advantages over independent and joint training and more than order-of-magnitude runtime savings.

2 Related Work

Prior work addresses task grouping through similarity assumptions, costly validation, transfer-learning relationships, architecture changes, optimization methods, and lookahead techniques. This paper instead emphasizes determining which tasks should share a network using a single training run.

Task Groupings: Early task-grouping methods cluster tasks using shared latent-space basis vectors, but their assumptions do not scale naturally to deep neural networks.
Task Groupings: Traditional neural task grouping relies on costly cross-validation or high-variance human intuition.
Task Groupings: Transfer-learning methods for task similarity do not necessarily carry over to multi-task learning, motivating multi-task-specific frameworks.
Task Groupings: A prior multi-task framework approximates exhaustive search with between |T| and 2^|T| − 1 models, whereas this approach computes groupings from one training run.
Architectures and Training Dynamics: Related architecture research changes parameter sharing, while this work focuses on when tasks should share a multi-task network.
Architectures and Training Dynamics: Optimization methods adjust loss weights or gradient directions to improve multi-task dynamics, representing a different focus from task grouping.
Looking into the Future: Lookahead methods reuse information from future parameter states and have been applied in meta-learning, optimization, and auxiliary-task learning.

3 Task Grouping Problem Definition

The task-grouping problem assigns tasks to multi-task networks while maximizing aggregate performance under inference-time latency and memory constraints. A network may train auxiliary tasks that are not served at inference.

Inference-time latency measures prediction speed, whereas memory budget measures the total parameters used across deployed models.
The framework seeks task assignments maximizing average performance while respecting per-network parameter limits, network-count limits, and complete task coverage.
Each task must be served by exactly one network, and the number of networks cannot exceed the inference-time memory budget.
A network can concurrently train tasks that it does not serve during inference, allowing auxiliary tasks to support served-task performance.
The grouping algorithm selects the network predicted to provide each task’s highest performance from the constructed model set.

4 Grouping Tasks by Measuring Inter-Task Affinity

The paper measures inter-task affinity from how one task’s shared-parameter gradient update changes another task’s loss, then uses these scores to select task groupings. TAG supports efficient network selection under inference constraints, with theoretical guarantees in a convex setting and empirical effectiveness in selecting high-performance groupings.

4.1 Inter-Task Affinity: Inter-task affinity measures how one task’s successive shared-parameter gradient updates affect another task’s objective.The measure serves as a proxy for information transfer between tasks in hard-parameter sharing.
4.1 Inter-Task Affinity: A lookahead loss compares task j’s loss before and after task i’s shared-parameter update using a scale-invariant loss ratio.Task-specific parameters and the input batch remain unchanged during this comparison.
4.1 Inter-Task Affinity: Positive affinity indicates that a shared-parameter update lowers the receiving task’s loss, whereas negative affinity indicates an antagonistic effect.Per-step affinities can be averaged across training, every n steps, or a contiguous subset of steps.
4.2 Network Selection Algorithm: TAG trains all tasks together once, measures pairwise affinities, selects groupings maximizing total affinity, and trains the resulting networks for evaluation.The selected networks must satisfy an inference-time memory budget and serve every task from exactly one network.
4.2 Network Selection Algorithm: For groups of three or more tasks, affinity onto each task is approximated by averaging the pairwise affinities from the other tasks.The selection problem is NP-hard but is solved efficiently with a branch-and-bound-like algorithm.
4.3 Theoretical Analysis: In the convex setting under mild assumptions, affinity-based grouping induces lower loss for a task than grouping it with a lower-affinity alternative.The guarantee is analyzed for strongly convex, strongly smooth losses and includes conditions on learning rate, gradient norms, and cosine similarity.

5 Experiments

Experiments evaluate TAG for selecting task groupings on CelebA and Taskonomy, alongside ablations of affinity computation, training dynamics, and hyperparameter sensitivity. TAG generally improves performance and efficiency over comparison methods, while its grouping quality depends on training and task-specific affinity structure.

5.1 Supervised Task Grouping Evaluation: TAG surpasses HOA, RG, and CS task grouping methods while operating 22 times faster than HOA.Optimization augmentations performed worse than TAG- or HOA-selected groups, although PCGrad sometimes improved TAG groups.
5.1 Supervised Task Grouping Evaluation: 10.0% lower total test loss than MTL was achieved by TAG on Taskonomy, while TAG required 140 fewer TeslaV100 GPU hours than CS.TAG’s 2-split grouping surpassed HOA’s by 2.5%, but HOA’s 3-split grouping performed better than TAG’s.
Multi-Task Ablation Studies: TAG struggled to identify the best auxiliary task for attribute a8 because no task showed especially high or low affinity onto it.The best-versus-worst predicted affinity difference for a8 was 0.04, compared with 0.16 for the next-smallest difference.
Should inter-task affinity be computed at every step?: Computing inter-task affinities every 10 steps eliminated redundancy without degrading performance and increased training-time efficiency by 319%.Computing less frequently after that threshold weakened the signal and increased task-grouping error.
Should inter-task affinity be computed at every step?: Affinity scores changed throughout training, although some tasks maintained higher-than-average affinity differences across time.Early training showed positive, similar affinities across tasks, followed by rapid specialization; affinity relationships later shifted between task pairs.
Hyperparameter Sensitivity: Changing batch size or learning rate could change which tasks should train together, with groupings from lr=2x generalizing worse than those from b=0.5x.These results indicate grouping depends on detailed model and training aspects as well as relationships among tasks.

6 Conclusion

The paper presents inter-task affinity as an efficient basis for determining multi-task groupings and reports competitive performance with substantially lower computational cost. It also highlights that small hyperparameter changes can alter suitable groupings, motivating automated identification methods.

The approach quantifies inter-task affinity in a single training run to determine which tasks should train together.
It outperforms Uncertainty Weights, GradNorm, and PCGrad while remaining competitive with HOA.
Its computational efficiency improves by over an order of magnitude relative to a state-of-the-art task grouping method.
Task groupings are sensitive to even small changes in hyperparameters, encouraging efficient and automatic grouping methods.

7 Broader Impact

Efficient task grouping can save substantial computational resources, but inter-task affinity must not be interpreted as task similarity or as evidence of association or causation.

Efficient task grouping may save significant time and computational resources in academic and industry environments.
Inter-task affinities can be mistakenly interpreted as task similarity, producing unfounded associations or causal claims.
These risks are especially consequential for datasets containing sensitive predictions related to race, gender, religion, age, status, or physical traits.

A.2 Proof of Proposition 1

The proposition formalizes when one auxiliary task's gradient yields a larger reduction in another task's loss than a competing auxiliary gradient. Under smoothness, convexity, and equal-gradient-norm conditions, the comparison can be expressed through cosine similarity.

The proposition assumes a strongly convex and strongly smooth loss and a learning rate satisfying η ≤ 1/β.
If task b has higher inter-task affinity than task c for task a, the combined gradient of tasks a and b reduces La more than the combined gradient of a and c under the stated condition.
The proof applies the strong-smoothness upper bound to the updated loss obtained with the combined gradient ga + gb.
The comparison eliminates common terms and reduces the claim to an inequality involving gradient norms, dot products, α, β, and η.
With equal gradient norms, the sufficient comparison condition is expressed using cosine similarity between gradients.

A.3 Quadratic Counterexample

A quadratic-loss counterexample shows that grouping tasks by inter-task similarity can produce inferior performance. The example compares auxiliary gradients and records their effects on the target loss after gradient steps.

The quadratic task-a loss has α = 1, β = 10, global minimum (0, 0), and initial point (−2, −1) with loss 7.
For an alternative task-c gradient, the individual losses are 6.96 versus 7.96, while combined gradients produce 6.09 versus 6.98.
Figure 6 compares classification error or total test loss across task-grouping splits as a function of inference-time memory budget.

B Additional Experimental Results

The supplementary experiments remove the fixed inference-time latency constraint by scaling baseline capacity to match each task grouping. Task grouping methods still outperform training all tasks together and optimization augmentations, while extra channels reduce CelebA performance.

Additional Experimental Results: The experiments remove the fixed inference-time latency constraint and scale baseline capacity to match each task grouping.A two-split grouping uses twice as many channels per convolutional layer as the baseline.
Additional Experimental Results: Task grouping approaches continue to outperform training all tasks together, Uncertainty Weights, and GradNorm.
Additional Experimental Results: For CelebA, increasing convolutional channels reduces performance, indicating the ResNet model is already near-optimal capacity for this dataset.

B.2 Additional CelebA Task Grouping Results

The supplementary results report raw benchmark scores and task-grouping results for CelebA and Taskonomy, including experiments with two, three, and four splits and higher-capacity models.

B.2 Additional CelebA Task Grouping Results: CelebA supplementary tables report means and standard errors for benchmark methods and task groupings.Separate tables cover benchmark methods, two-split, three-split, four-split, and high-capacity experiments.
B.2 Additional CelebA Task Grouping Results: CelebA quantities are averaged across three independent runs, whereas Taskonomy is evaluated using a single run because of its dataset size and computational cost.The reported Taskonomy baseline requires approximately 146 Tesla V100 GPU hours for one model.
B.2 Additional CelebA Task Grouping Results: Taskonomy supplementary tables report baseline methods, two-split and three-split groupings, and high-capacity model performance.

B.4 Supplementary Information on Ablation Studies

The ablation studies examine whether inter-task affinity tracks useful groupings, remains stable across task sets, and has identifiable failure cases and training-versus-validation limitations. They report natural affinity-based groups, cross-set stability in CelebA, and boundaries on when the method can select groupings.

Ablation Scope: The supplementary analysis evaluates affinity-grouping correspondence, cross-task-set robustness, a failure case, and the training-versus-validation limitation.The experiments use CelebA and Taskonomy analyses alongside the reported ablation studies and released implementation details.
Affinity Robustness: Inter-task affinity scores between two tasks are largely maintained when the surrounding task set changes.In CelebA, affinity from a7 onto a1 and from a1 onto a7 remains consistent across two different task sets.
Affinity Robustness: Affinity stability across task sets is not guaranteed when the second set contains a task that sharply reduces performance or causes divergence.Under those conditions, affinity between two tasks may differ substantially across task sets.
Affinity and Task Groupings: Affinity matrices reveal natural CelebA groups {a1, a7}, {a2,a3,a4}, and {a5, a6, a7}, while Taskonomy forms segmentation-depth-normals and edges-keypoints groups.The Network Selection Algorithm uses this information to assign tasks to networks.
Failure Case: TAG may fail to select a positive partner for a8 because no other CelebA task substantially changes a8’s loss during training.Affinity onto a8 is significantly lower than affinity onto every other task and four times smaller than the next-smallest difference.
Training-versus-Validation Ablation: The cross-task approximation between training-loss and validation-loss changes does not hold for self-affinity when j = i.Using training data therefore prevents TAG from grouping a task by itself, although loading validation data could trade efficiency for single-task selection.