# Balancing Average and Worst-case Accuracy in Multitask Learning

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 53 references that are in the map.
- First posted on arXiv 12 Oct 2021.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Unsupervised Cross-lingual Representation Learning at Scale](https://mlsanity.com/paper/1911.02116/) (2019)
- [mT5: A massively multilingual pre-trained text-to-text transformer](https://mlsanity.com/paper/2010.11934/) (2020)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [XTREME: A Massively Multilingual Multi-task Benchmark for Evaluating Cross-lingual Generalization](https://mlsanity.com/paper/2003.11080/) (2020)
- [Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates](https://mlsanity.com/paper/1804.10959/) (2018)
- [Multi-Task Learning with Deep Neural Networks: A Survey](https://mlsanity.com/paper/2009.09796/) (2020)
- [Massively Multilingual Neural Machine Translation](https://mlsanity.com/paper/1903.00089/) (2019)
- [The Natural Language Decathlon: Multitask Learning as Question Answering](https://mlsanity.com/paper/1806.08730/) (2018)
- [Pareto Multi-Task Learning](https://mlsanity.com/paper/1912.12854/) (2019)
- [Distributionally Robust Stochastic Optimization with Wasserstein Distance](https://mlsanity.com/paper/1604.02199/) (2016)
- [Large-Scale Methods for Distributionally Robust Optimization](https://mlsanity.com/paper/2010.05893/) (2020)
- [Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models](https://mlsanity.com/paper/2010.05874/) (2020)
- [Distributionally Robust Language Modeling](https://mlsanity.com/paper/1909.02060/) (2019)
- [Multi-task Deep Reinforcement Learning with PopArt](https://mlsanity.com/paper/1809.04474/) (2018)
- [Routing Networks: Adaptive Selection of Non-linear Functions for Multi-Task Learning](https://mlsanity.com/paper/1711.01239/) (2017)
- [Examining and Combating Spurious Features under Distribution Shift](https://mlsanity.com/paper/2106.07171/) (2021)
- [A Hierarchical Multi-task Approach for Learning Embeddings from Semantic Tasks](https://mlsanity.com/paper/1811.06031/) (2018)
- [What Kind of Language Is Hard to Language-Model?](https://mlsanity.com/paper/1906.04726/) (2019)
- [Adaptive Scheduling for Multi-Task Learning](https://mlsanity.com/paper/1909.06434/) (2019)
- [Scheduled Multi-Task Learning: From Syntax to Translation](https://mlsanity.com/paper/1804.08915/) (2018)

## Built on by (most cited first, 10 of 10)

- [UniMax: Fairer and more Effective Language Sampling for Large-Scale Multilingual Pretraining](https://mlsanity.com/paper/2304.09151/) (2023)
- [Auto-Lambda: Disentangling Dynamic Task Relationships](https://mlsanity.com/paper/2202.03091/) (2022)
- [SELF-GUIDE: Better Task-Specific Instruction Following via Self-Synthetic Finetuning](https://mlsanity.com/paper/2407.12874/) (2024)
- [Robust Multi-Task Learning with Excess Risks](https://mlsanity.com/paper/2402.02009/) (2024)
- [Less is More: Selective Layer Finetuning with SubTuning](https://mlsanity.com/paper/2302.06354/) (2023)
- [AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs](https://mlsanity.com/paper/2506.00569/) (2025)
- [DaMo: Data Mixing Optimizer in Fine-tuning Multimodal LLMs for Mobile Phone Agents](https://mlsanity.com/paper/2510.19336/) (2025)
- [Multitask Learning Can Improve Worst-Group Outcomes](https://mlsanity.com/paper/2312.03151/) (2023)
- [MaD-Mix: Multi-Modal Data Mixtures via Latent Space Coupling for Vision-Language Model Training](https://mlsanity.com/paper/2602.07790/) (2026)
- [Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling](https://mlsanity.com/paper/2605.14350/) (2026)

## Abstract (arXiv)

When training and evaluating machine learning models on a large number of tasks, it is important to not only look at average task accuracy -- which may be biased by easy or redundant tasks -- but also worst-case accuracy (i.e. the performance on the task with the lowest accuracy). In this work, we show how to use techniques from the distributionally robust optimization (DRO) literature to improve worst-case performance in multitask learning. We highlight several failure cases of DRO when applied off-the-shelf and present an improved method, Lookahead-DRO (L-DRO), which mitigates these issues. The core idea of L-DRO is to anticipate the interaction between tasks during training in order to choose a dynamic re-weighting of the various task losses, which will (i) lead to minimal worst-case loss and (ii) train on as many tasks as possible. After demonstrating the efficacy of L-DRO on a small controlled synthetic setting, we evaluate it on two realistic benchmarks: a multitask version of the CIFAR-100 image classification dataset and a large-scale multilingual language modeling experiment. Our empirical results show that L-DRO achieves a better trade-off between average and worst-case accuracy with little computational overhead compared to several strong baselines.