Research map / Large language models / Language, safety and interpretability
Data Selection via Optimal Control for Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Formulates LM pre-training data selection as an optimal control problem solved via Pontryagin's Maximum Principle, and proposes PDS, which approximates the PMP conditions with a proxy LM and a learned data scorer to select CommonCrawl data.
Problem
Selecting high-quality pre-training data from massive web corpora so that LMs learn faster and perform better downstream, with theoretical grounding beyond heuristics.
Approach
Treat data quality weights as time-invariant control variables over the LM training dynamics, derive necessary optimality conditions via discrete PMP, and approximate them with a proxy-LM bi-level solver averaged over checkpoints, then distill scores into a data scorer for offline selection.
Key results
- PDS reaches 51.0 average accuracy on OLMo tasks at 1B vs 49.3 for conventional pre-training (470M: 48.2 vs 47.0).
- PDS achieves about 2.0x training-compute acceleration on the 1.7B model compared to conventional pre-training.
- In the data-constrained setting, PDS reduces data demand by 1.8x (selecting 1/4 of data and training 4 epochs).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Select pre-training data with PDS instead of conventional uniform sampling | 470M LM · OLMo evaluation tasks | average accuracy: 47.0 -> 48.2 | Helped |
| Select pre-training data with PDS instead of DSIR | 1B LM · OLMo evaluation tasks | average accuracy: 49.6 -> 51.0 | Helped |
| Select pre-training data with PDS instead of IF-Score | 470M LM · MMLU | 0-shot accuracy: 28.4 -> 28.9 | Helped |
| Select pre-training data with PDS instead of IF-Score | 470M LM · MMLU | perplexity: 31.1 -> 27.1 | Helped |
| Select pre-training data with PDS instead of IF-Score | 1B LM · MMLU | 0-shot accuracy: 30.7 -> 31.4 | Helped |
| Select pre-training data with PDS instead of IF-Score | 1B LM · MMLU | perplexity: 23.6 -> 20.5 | Helped |
| Use long-range training dynamics rather than single-step information for data scoring (ablation) | PDS · OLMo evaluation tasks | accuracy | Helped |
| Use early-stage rather than later-stage training checkpoints for data scoring (ablation) | PDS · OLMo evaluation tasks | accuracy | Helped |
| Use an 8.7M proxy LM (ablation) | PDS · OLMo evaluation tasks | average accuracy: 44.0; DSIR 43.8 | Mixed |
| Compute data quality scores based on Adam instead of GD (ablation) | PDS · Pilot experiments | performance | No clear effect |
| Select one quarter of the limited corpus with PDS and train for 4 epochs | PDS · Data-constrained setting, DCLM | test loss: 1.8 times | Helped |
| Use PDS-selected data instead of conventional pre-training data | 1.7B LM · OLMo evaluation tasks | training computation to reach comparable accuracy: 2.0 times acceleration | Helped |
Surprising findings
- A proxy LM of only 8.7M parameters yields downstream performance close to DSIR, suggesting small LMs mainly capture shallow quality patterns.
- Early-stage training dynamics (checkpoints 10K-50K) gave better selection than later-stage checkpoints (50K-100K).
Limitations
- PMP conditions are only necessary, so the solver does not guarantee optimal selection.
- Individual-importance formulation ignores dependence between data points, which can hurt diversity on duplicated corpora.
- Scaling-law extrapolations to 400B models are not empirically verified.
Experimental setup
Datasets: RedPajama CommonCrawl LIMA DCLM
Benchmarks: OLMo evaluation suite (HellaSwag, LAMBADA, Winogrande, OBQA, ARC-e, ARC-c, PIQA, SciQ, BoolQ) MMLU DCLM subset loss
Metrics: 0-shot accuracy average accuracy perplexity test loss training FLOPs speedup
Baselines: Conventional (RedPajama) pre-training RHO-Loss DSIR IF-Score
Builds on: Pontryagin's Maximum Principle Adam RHO-Loss DSIR LIMA
Closest papers by meaning
- Predictive Data Selection: The Data That Predicts Is the Data That Teaches (ICML 2025 · similarity 0.90)
- BLADE: Scalable Bi-level Adaptive Data Selection for LLM Training (NeurIPS 2026 · similarity 0.89)
- LLM Data Selection and Utilization via Dynamic Bi-level Optimization (ICML 2025 · similarity 0.89)
- MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models (NeurIPS 2024 · similarity 0.88)
- Get more for less: Principled Data Selection for Warming Up Fine-Tuning in LLMs (ICLR 2024 · similarity 0.88)
- GREATS: Online Selection of High-Quality Data for LLM Training in Every Iteration (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.