# Task-Robust Pre-Training for Worst-Case Downstream Adaptation

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 69 references that are in the map.
- First posted on arXiv 21 Jun 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Certified Adversarial Robustness via Randomized Smoothing](https://mlsanity.com/paper/1902.02918/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [SimMIM: A Simple Framework for Masked Image Modeling](https://mlsanity.com/paper/2111.09886/) (2021)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [Pre-Trained Image Processing Transformer](https://mlsanity.com/paper/2012.00364/) (2020)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language](https://mlsanity.com/paper/2202.03555/) (2022)
- [Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation](https://mlsanity.com/paper/2012.07177/) (2020)
- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://mlsanity.com/paper/2112.09133/) (2021)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [On Evaluating Adversarial Robustness](https://mlsanity.com/paper/1902.06705/) (2019)

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

- [MedSeqFT: Sequential Fine-tuning Foundation Models for 3D Medical Image Segmentation](https://mlsanity.com/paper/2509.06096/) (2025)
- [Nonparametric Teaching of Attention Learners](https://mlsanity.com/paper/2602.20461/) (2026)
- [First-Order Softmax Weighted Switching Gradient Method for Distributed Stochastic Minimax Optimization with Stochastic Constraints](https://mlsanity.com/paper/2603.05774/) (2026)

## Abstract (arXiv)

Pre-training has achieved remarkable success when transferred to downstream tasks. In machine learning, we care about not only the good performance of a model but also its behavior under reasonable shifts of condition. The same philosophy holds when pre-training a foundation model. However, the foundation model may not uniformly behave well for a series of related downstream tasks. This happens, for example, when conducting mask recovery regression where the recovery ability or the training instances diverge like pattern features are extracted dominantly on pre-training, but semantic features are also required on a downstream task. This paper considers pre-training a model that guarantees a uniformly good performance over the downstream tasks. We call this goal as textitdownstream-task robustness. Our method first separates the upstream task into several representative ones and applies a simple minimax loss for pre-training. We then design an efficient algorithm to solve the minimax loss and prove its convergence in the convex setting. In the experiments, we show both on large-scale natural language processing and computer vision datasets our method increases the metrics on worse-case downstream tasks. Additionally, some theoretical explanations for why our loss is beneficial are provided. Specifically, we show fewer samples are inherently required for the most challenging downstream task in some cases.