[Paper Review] Efficiently Identifying Task Groupings for Multi-Task Learning

Christopher Fifty, Ehsan Amid+4 more|arXiv (Cornell University)|Sep 10, 2021

Domain Adaptation and Few-Shot Learning52 references29 citations

TL;DR

TAG trains all tasks in one run to measure inter-task affinity via gradient impact, then constructs task groups that maximize affinity, yielding faster and better multi-task performance.

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.

Motivation & Objective

• Motivate the need for efficient task grouping in multi-task learning to avoid negative transfer and excessive search costs.
• Propose a single-run method to quantify inter-task affinity via gradient-based lookahead losses.
• Develop a network selection algorithm that forms multi-task groups maximizing per-task affinity under inference constraints.
• Demonstrate scalability and efficiency on large vision datasets (CelebA, Taskonomy).
• Provide theoretical insights showing why affinity-based grouping can outperform alternatives under mild convex assumptions.

Proposed method

• Train all tasks jointly to obtain shared parameter updates.
• Define inter-task affinity Z-i->j as the relative progress of task j when task i updates shared parameters (lookahead loss ratio).
• Average the per-step affinity across training to obtain a training-level score hat{Z}_{i→j}.
• Group tasks into k multi-task networks (k ≤ memory budget b) by maximizing the summed affinity onto each served task.
• Provide a theoretical analysis showing that, under alpha-strong convexity and beta-strong smoothness with certain conditions, higher affinity gives lower loss for the primary task when grouped with the higher-affinity task.
• Compare TAG against baseline MTL, STL, random grouping, cosine similarity-based grouping, and HOA, with and without additional training augmentations like PCGrad.

Experimental results

Research questions

• RQ1 Does inter-task affinity measured during a single training run align with identifying beneficial task groupings?
• RQ2 Is per-step affinity measurement more effective than per-epoch or aggregate measures for grouping decisions?
• RQ3 How does TAG perform relative to training all tasks jointly, single-task models, and existing grouping methods on large-scale vision benchmarks?
• RQ4 Do the identified groups remain effective under varying inference-time memory/latency budgets?
• RQ5 Can the affinity-based approach be complemented with training-time augmentations to further improve performance?

Key findings

• TAG reduces test loss by up to 10.0% on Taskonomy compared to training all tasks together.
• TAG operates about 11.6x faster than a state-of-the-art task grouping method (HOA) in their benchmarks.
• On CelebA, TAG outperforms MTL, UW, GN, and RG, with 2-, 3-, and 4-splits showing consistent gains, and can be further boosted by PCGrad.
• On Taskonomy, TAG consistently improves over MTL (10.0%), GN (7.7%), STL (1.5%), and RG (9.5%).
• TAG is significantly more computationally efficient than HOA (e.g., 22x faster in reported experiments).