# Adaptive Training Distributions with Scalable Online Bilevel Optimization

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 84 references that are in the map.
- First posted on arXiv 20 Nov 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models](https://mlsanity.com/paper/2206.04615/) (2022)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Learning to Reweight Examples for Robust Deep Learning](https://mlsanity.com/paper/1803.09050/) (2018)
- [MentorNet: Learning Data-Driven Curriculum for Very Deep Neural Networks on Corrupted Labels](https://mlsanity.com/paper/1712.05055/) (2017)
- [Data Shapley: Equitable Valuation of Data for Machine Learning](https://mlsanity.com/paper/1904.02868/) (2019)
- [Exploring the Limits of Weakly Supervised Pretraining](https://mlsanity.com/paper/1805.00932/) (2018)
- [Estimating Training Data Influence by Tracing Gradient Descent](https://mlsanity.com/paper/2002.08484/) (2020)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [Not All Samples Are Created Equal: Deep Learning with Importance Sampling](https://mlsanity.com/paper/1803.00942/) (2018)
- [How Good is Your Tokenizer? On the Monolingual Performance of Multilingual Language Models](https://mlsanity.com/paper/2012.15613/) (2020)
- [A Brief Review of Domain Adaptation](https://mlsanity.com/paper/2010.03978/) (2020)
- [Fast AutoAugment](https://mlsanity.com/paper/1905.00397/) (2019)
- [Revisiting Few-sample BERT Fine-tuning](https://mlsanity.com/paper/2006.05987/) (2020)
- [Approximation Methods for Bilevel Programming](https://mlsanity.com/paper/1802.02246/) (2018)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [Unsupervised Domain Clusters in Pretrained Language Models](https://mlsanity.com/paper/2004.02105/) (2020)
- [Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules](https://mlsanity.com/paper/1905.05393/) (2019)

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

- [Data Selection via Optimal Control for Language Models](https://mlsanity.com/paper/2410.07064/) (2024)
- [Task-Adaptive Pretrained Language Models via Clustered-Importance Sampling](https://mlsanity.com/paper/2410.03735/) (2024)
- [Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining](https://mlsanity.com/paper/2502.06733/) (2025)
- [Dynamic Gradient Alignment for Online Data Mixing](https://mlsanity.com/paper/2410.02498/) (2024)
- [Careful with that Scalpel: Improving Gradient Surgery with an EMA](https://mlsanity.com/paper/2402.02998/) (2024)
- [A Bayesian Approach to Data Point Selection](https://mlsanity.com/paper/2411.03768/) (2024)
- [Soup-of-Experts: Pretraining Specialist Models via Parameters Averaging](https://mlsanity.com/paper/2502.01804/) (2025)
- [Online Nonconvex Bilevel Optimization with Bregman Divergences](https://mlsanity.com/paper/2409.10470/) (2024)
- [Fair Bayesian Data Selection via Generalized Discrepancy Measures](https://mlsanity.com/paper/2511.07032/) (2025)
- [Influence Guided Sampling for Domain Adaptation of Text Retrievers](https://mlsanity.com/paper/2601.21759/) (2026)
- [Detectability in Diversity: Improved Canary Crafting for Privacy Auditing in One Run](https://mlsanity.com/paper/2605.27292/) (2026)

## Abstract (arXiv)

Large neural networks pretrained on web-scale corpora are central to modern machine learning. In this paradigm, the distribution of the large, heterogeneous pretraining data rarely matches that of the application domain. This work considers modifying the pretraining distribution in the case where one has a small sample of data reflecting the targeted test conditions. We propose an algorithm motivated by a recent formulation of this setting as an online, bilevel optimization problem. With scalability in mind, our algorithm prioritizes computing gradients at training points which are likely to most improve the loss on the targeted distribution. Empirically, we show that in some cases this approach is beneficial over existing strategies from the domain adaptation literature but may not succeed in other cases. We propose a simple test to evaluate when our approach can be expected to work well and point towards further research to address current limitations.