# Data Selection via Optimal Control for Language Models

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

## Key facts

- Cited by 38 papers in the map (data through 14 Oct 2026).
- Builds on 66 of its 87 references that are in the map.
- First posted on arXiv 9 Oct 2024.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [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)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [DeepSeek LLM: Scaling Open-Source Language Models with Longtermism](https://mlsanity.com/paper/2401.02954/) (2024)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [Are Emergent Abilities of Large Language Models a Mirage?](https://mlsanity.com/paper/2304.15004/) (2023)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)

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

- [NVILA: Efficient Frontier Visual Language Models](https://mlsanity.com/paper/2412.04468/) (2024)
- [MiniPLM: Knowledge Distillation for Pre-Training Language Models](https://mlsanity.com/paper/2410.17215/) (2024)
- [Language Models Improve When Pretraining Data Matches Target Tasks](https://mlsanity.com/paper/2507.12466/) (2025)
- [QuaDMix: Quality-Diversity Balanced Data Selection for Efficient LLM Pretraining](https://mlsanity.com/paper/2504.16511/) (2025)
- [How Learning Rate Decay Wastes Your Best Data in Curriculum-Based LLM Pretraining](https://mlsanity.com/paper/2511.18903/) (2025)
- [Group-Level Data Selection for Efficient Pretraining](https://mlsanity.com/paper/2502.14709/) (2025)
- [DataRater: Meta-Learned Dataset Curation](https://mlsanity.com/paper/2505.17895/) (2025)
- [Data-Efficient RLVR via Off-Policy Influence Guidance](https://mlsanity.com/paper/2510.26491/) (2025)
- [Adaptive Defense against Harmful Fine-Tuning for Large Language Models via Bayesian Data Scheduler](https://mlsanity.com/paper/2510.27172/) (2025)
- [Influence-Preserving Proxies for Gradient-Based Data Selection in LLM Fine-tuning](https://mlsanity.com/paper/2602.17835/) (2026)
- [IDEAL: Data Equilibrium Adaptation for Multi-Capability Language Model Alignment](https://mlsanity.com/paper/2505.12762/) (2025)
- [Data Efficacy for Language Model Training](https://mlsanity.com/paper/2506.21545/) (2025)
- [Optimsyn: Influence-Guided Rubrics Optimization for Synthetic Data Generation](https://mlsanity.com/paper/2604.00536/) (2026)
- [Domain2Vec: Vectorizing Datasets to Find the Optimal Data Mixture without Training](https://mlsanity.com/paper/2506.10952/) (2025)
- [Shift is Good: Mismatched Data Mixing Improves Test Performance](https://mlsanity.com/paper/2510.25108/) (2025)
- [DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models](https://mlsanity.com/paper/2603.26164/) (2026)
- [MASS: Mathematical Data Selection via Skill Graphs for Pretraining Large Language Models](https://mlsanity.com/paper/2503.14917/) (2025)
- [Nonparametric Data Attribution for Diffusion Models](https://mlsanity.com/paper/2510.14269/) (2025)
- [Mining Useful General Data for Low-Resource Domain Adaptation](https://mlsanity.com/paper/2511.07380/) (2025)
- [TRIM: Token-wise Attention-Derived Saliency for Data-Efficient Instruction Tuning](https://mlsanity.com/paper/2510.07118/) (2025)
- [Towards Next-Generation LLM Training: From the Data-Centric Perspective](https://mlsanity.com/paper/2603.14712/) (2026)
- [Demystifying Data Organization for Enhanced LLM Training](https://mlsanity.com/paper/2605.30334/) (2026)
- [LAMDAS: LLM as an Implicit Classifier for Domain-specific Data Selection](https://mlsanity.com/paper/2509.06524/) (2025)
- [RL-Guided Data Selection for Language Model Finetuning](https://mlsanity.com/paper/2509.25850/) (2025)
- [Rethinking Data Mixing from the Perspective of Large Language Models](https://mlsanity.com/paper/2604.07963/) (2026)
- [Unifying and Optimizing Data Values for Selection via Sequential Decision-Making](https://mlsanity.com/paper/2502.04554/) (2025)
- [Measuring Fine-Grained Relatedness in Multitask Learning via Data Attribution](https://mlsanity.com/paper/2505.21438/) (2025)
- [Data Selection for LLM Alignment Using Fine-Grained Preferences](https://mlsanity.com/paper/2508.07638/) (2025)
- [Holdout-Loss-Based Data Selection for LLM Finetuning via In-Context Learning](https://mlsanity.com/paper/2510.14459/) (2025)
- [Alchemist: Unlocking Efficiency in Text-to-Image Model Training via Meta-Gradient Data Selection](https://mlsanity.com/paper/2512.16905/) (2025)

## Abstract (arXiv)

This work investigates the selection of high-quality pre-training data from massive corpora to enhance LMs' capabilities for downstream usage. We formulate data selection as a generalized Optimal Control problem, which can be solved theoretically by Pontryagin's Maximum Principle (PMP), yielding a set of necessary conditions that characterize the relationship between optimal data selection and LM training dynamics. Based on these theoretical results, we introduce PMP-based Data Selection (PDS), a framework that approximates optimal data selection by solving the PMP conditions. In our experiments, we adopt PDS to select data from CommmonCrawl and show that the PDS-selected corpus accelerates the learning of LMs and constantly boosts their performance on a wide range of downstream tasks across various model sizes. Moreover, the benefits of PDS extend to ~400B models trained on ~10T tokens, as evidenced by the extrapolation of the test loss curves according to the Scaling Laws. PDS also improves data utilization when the pre-training data is limited, by reducing the data demand by 1.8 times, which helps mitigate the quick exhaustion of available web-crawled corpora. Our code, model, and data can be found at https://github.com/microsoft/LMOps/tree/main/data_selection.