# Learning to Branch for Multi-Task Learning

arXiv: https://arxiv.org/abs/2006.01895 · HTML: https://mlsanity.com/paper/2006.01895/

## Key facts

- Cited by 115 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 2 Jun 2020.

## Built on by (most cited first, 30 of 115)

- [Domain Generalization: A Survey](https://mlsanity.com/paper/2103.02503/) (2021)
- [AdaMerging: Adaptive Model Merging for Multi-Task Learning](https://mlsanity.com/paper/2310.02575/) (2023)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [BEVerse: Unified Perception and Prediction in Birds-Eye-View for Vision-Centric Autonomous Driving](https://mlsanity.com/paper/2205.09743/) (2022)
- [RotoGrad: Gradient Homogenization in Multitask Learning](https://mlsanity.com/paper/2103.02631/) (2021)
- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [FAMO: Fast Adaptive Multitask Optimization](https://mlsanity.com/paper/2306.03792/) (2023)
- [In Defense of the Unitary Scalarization for Deep Multi-Task Learning](https://mlsanity.com/paper/2201.04122/) (2022)
- [Multi-Task Dense Prediction via Mixture of Low-Rank Experts](https://mlsanity.com/paper/2403.17749/) (2024)
- [MetaGPT: Merging Large Language Models Using Model Exclusive Task Arithmetic](https://mlsanity.com/paper/2406.11385/) (2024)
- [LoRA Soups: Merging LoRAs for Practical Skill Composition Tasks](https://mlsanity.com/paper/2410.13025/) (2024)
- [Recon: Reducing Conflicting Gradients from the Root for Multi-Task Learning](https://mlsanity.com/paper/2302.11289/) (2023)
- [TaskExpert: Dynamically Assembling Multi-Task Representations with Memorial Mixture-of-Experts](https://mlsanity.com/paper/2307.15324/) (2023)
- [DINO-Foresight: Looking into the Future with DINO](https://mlsanity.com/paper/2412.11673/) (2024)
- [Learning Multiple Dense Prediction Tasks from Partially Annotated Data](https://mlsanity.com/paper/2111.14893/) (2021)
- [Learning Multiple Stock Trading Patterns with Temporal Routing Adaptor and Optimal Transport](https://mlsanity.com/paper/2106.12950/) (2021)
- [Improving Multi-Task Generalization via Regularizing Spurious Correlation](https://mlsanity.com/paper/2205.09797/) (2022)
- [Parallel Vertex Diffusion for Unified Visual Grounding](https://mlsanity.com/paper/2303.07216/) (2023)
- [Universal Representations: A Unified Look at Multiple Task and Domain Learning](https://mlsanity.com/paper/2204.02744/) (2022)
- [Revisit Multimodal Meta-Learning through the Lens of Multi-Task Learning](https://mlsanity.com/paper/2110.14202/) (2021)
- [InvPT++: Inverted Pyramid Multi-Task Transformer for Visual Scene Understanding](https://mlsanity.com/paper/2306.04842/) (2023)
- [Controllable Dynamic Multi-Task Architectures](https://mlsanity.com/paper/2203.14949/) (2022)
- [Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity](https://mlsanity.com/paper/2409.06091/) (2024)
- [Egocentric Video Task Translation](https://mlsanity.com/paper/2212.06301/) (2022)
- [AutoMTL: A Programming Framework for Automating Efficient Multi-Task Learning](https://mlsanity.com/paper/2110.13076/) (2021)
- [FedHCA2: Towards Hetero-Client Federated Multi-Task Learning](https://mlsanity.com/paper/2311.13250/) (2023)
- [MLP Can Be A Good Transformer Learner](https://mlsanity.com/paper/2404.05657/) (2024)
- [Selective Task Group Updates for Multi-Task Optimization](https://mlsanity.com/paper/2502.11986/) (2025)
- [Variational Multi-Task Learning with Gumbel-Softmax Priors](https://mlsanity.com/paper/2111.05323/) (2021)
- [Prompt Guided Transformer for Multi-Task Dense Prediction](https://mlsanity.com/paper/2307.15362/) (2023)

## Abstract (arXiv)

Training multiple tasks jointly in one deep network yields reduced latency during inference and better performance over the single-task counterpart by sharing certain layers of a network. However, over-sharing a network could erroneously enforce over-generalization, causing negative knowledge transfer across tasks. Prior works rely on human intuition or pre-computed task relatedness scores for ad hoc branching structures. They provide sub-optimal end results and often require huge efforts for the trial-and-error process. In this work, we present an automated multi-task learning algorithm that learns where to share or branch within a network, designing an effective network topology that is directly optimized for multiple objectives across tasks. Specifically, we propose a novel tree-structured design space that casts a tree branching operation as a gumbel-softmax sampling procedure. This enables differentiable network splitting that is end-to-end trainable. We validate the proposed method on controlled synthetic data, CelebA, and Taskonomy.