Visual Summary | AutoData: Agentic Search for Pre-training Data Selection

Summary (Overview)

AutoData is an agentic framework that frames pre-training data selection as a search over executable selection algorithms, extending autonomous AI research from model and training-code optimization to data engineering.
The framework uses an AIDE-style LLM agent that iteratively proposes, executes, and refines data selection recipes based on proxy model validation feedback (val-bpb or CORE accuracy).
Within ~200 search steps on a small GPT-2 proxy model (125M), AutoData discovers recipes that outperform human-designed pipelines including DCLM, perplexity filtering, RegMix, and default ClimbMix ordering.
Discovered recipes transfer across model scales (125M to 1.3B parameters), achieving the best validation bits-per-byte (val-bpb) with statistically significant improvements over all baselines at most scales.
Analysis reveals that AutoData discovers composite scoring rules paired with diversity-preserving selection mechanisms (e.g., tournament selection, stratified quotas, Gumbel noise), moving beyond single-feature ranking and threshold filtering.

Introduction and Theoretical Foundation

Motivation

The performance of large language models depends on both model architecture and data. While recent autoregressive research agents (e.g., for NanoChat speedrun tasks) have automated model and training-code optimization, they treat training data as static. This overlooks a first-order lever: data. A key motivating example is that replacing FineWeb-Edu with NVIDIA-ClimbMix reduces wall-clock GPT-2 training time by 27%—a larger gain than most architecture-level improvements at that scale.

Problem Statement

Existing human-designed data curation methods combine deduplication, quality classification, perplexity-based filtering, importance sampling, and domain mixture optimization. However, these manually designed methods rely on fixed combinations of heuristics, bottlenecking both the design space and iteration speed.

Theoretical Framing

AutoData frames data selection as a heuristic engineering problem over per-document features. Each document is characterized by lexical statistics, categorical labels, perplexity, and LLM-annotated quality signals. The agent searches over a program space of scoring, stratification, and stochastic selection rules, discovering feature interactions automatically through iterative refinement with validation feedback.

Methodology

Problem Formulation

Given a candidate document pool $\mathcal{D} = \{d_i\}_{i=1}^{N}$ with document-level features $\{\mathbf{x}_i\}_{i=1}^{N}$ and a fixed budget $B$ , a selection algorithm is a function:

f: \left(\mathcal{D}, \{\mathbf{x}_i\}_{i=1}^{N}, B\right) \to \mathcal{S}, \quad \mathcal{S} \subset \mathcal{D}, |\mathcal{S}| = B

AutoData searches over the space $\mathcal{F}$ of selection algorithms to maximize empirical training performance:

f^{*} = \arg \max_{f \in \mathcal{F}} J\big(\mathcal{A}_{\theta}\big(f(\mathcal{D}, \{\mathbf{x}_i\}, B)\big)\big)

where $\mathcal{A}_{\theta}(S)$ is the model obtained by pre-training with fixed hyperparameters $\theta$ on subset $S$ , and $J(\cdot)$ is an evaluation function.

Feature Bank

Four complementary feature axes are provided to the agent:

Feature Type	Description
Lexical	Token length, distinct $n$ -gram ratios for $n \in \{1,\ldots,5\}$
Categorical	Topic and format labels from WebOrganizer classifiers (24×24 = 576 joint categories)
Perplexity	Per-document bits-per-byte under Qwen2.5-0.5B-Base reference model
LLM annotation	Gemini-3-Flash annotations: `n_factual` (incorrect factual claims), `n_rsteps` (inferential reasoning steps), `n_rerrors` (invalid reasoning steps)

Self-Evolving Loop

At each iteration:

The LLM agent proposes a candidate selection function $f \in \mathcal{F}$
The program is executed in three stages: (i) construct training subset $S = f(\mathcal{D}, \{\mathbf{x}_i\}, B)$ ; (ii) pre-train a proxy language model with fixed hyperparameters $\theta$ ; (iii) evaluate on a held-out validation set
The score is returned as feedback, guiding subsequent proposals

Experimental Setup

Data pool: NVIDIA ClimbMix (400B-token English corpus, 553,155,584 documents across 6,542 shards)
Selection budget: $B = 14,374,266$ documents (~2.6% of corpus), matching NanoChat re-training task scale
Proxy model: Depth-8 GPT-2 with target param-data-ratio = 10 (~0.42B training tokens), ~10 min per run on one H100 GPU
Subsample protocol: Each strategy evaluated with 4 proxy training runs (varying seed only) on the same selected subsample, run in parallel on 4×H100 GPUs
Evaluation signals: val-bpb (held-out ClimbMix shard) and CORE (unweighted mean of centered accuracy across 22 standard tasks)

Empirical Validation / Results

Main Results (Table 1)

The table below shows AutoData vs. human-designed baselines across five model scales (depth = 8 to 24). Mean over three runs, with sample standard deviation in subscript. Bold = best mean; shaded cells = statistically significant improvement over that method (p < 0.05, paired t-tests).

Method	d=8 (125M)	d=12 (286M)	d=16 (537M)	d=20 (897M)	d=24 (1.3B)
val-bpb ↓
Random uniform	$0.9484_{\pm 0.0002}$	$0.8477_{\pm 0.0003}$	$0.7806_{\pm 0.0002}$	$0.7346_{\pm 0.0004}$	$0.7055_{\pm 0.0011}$
DCLM-Baseline	$0.9954_{\pm 0.0001}$	$0.8886_{\pm 0.0001}$	$0.8202_{\pm 0.0003}$	$0.7766_{\pm 0.0006}$	$0.7499_{\pm 0.0001}$
PPL filter	$0.9605_{\pm 0.0001}$	$0.8638_{\pm 0.0000}$	$0.7983_{\pm 0.0001}$	$0.7545_{\pm 0.0003}$	$0.7267_{\pm 0.0001}$
RegMix	$0.9550_{\pm 0.0009}$	$0.8529_{\pm 0.0005}$	$0.7848_{\pm 0.0002}$	$0.7401_{\pm 0.0002}$	$0.7124_{\pm 0.0002}$
AutoData (CORE)	$0.9529_{\pm 0.0001}$	$0.8474_{\pm 0.0001}$	$0.7791_{\pm 0.0001}$	$0.7334_{\pm 0.0001}$	$0.7064_{\pm 0.0001}$
AutoData (VAL-BPB)	$0.9475_{\pm 0.0000}$	$0.8473_{\pm 0.0001}$	$0.7802_{\pm 0.0001}$	$0.7349_{\pm 0.0004}$	$0.7065_{\pm 0.0003}$
CORE ↑
Random uniform	$0.1041_{\pm 0.0125}$	$0.1403_{\pm 0.0034}$	$0.2053_{\pm 0.0065}$	$0.2364_{\pm 0.0011}$	$0.2609_{\pm 0.0058}$
DCLM-Baseline	$0.1043_{\pm 0.0028}$	$0.1362_{\pm 0.0065}$	$0.2006_{\pm 0.0111}$	$0.2245_{\pm 0.0027}$	$0.2470_{\pm 0.0019}$
PPL filter	$0.0933_{\pm 0.0051}$	$0.1449_{\pm 0.0034}$	$0.1926_{\pm 0.0029}$	$0.2276_{\pm 0.0080}$	$0.2535_{\pm 0.0051}$
RegMix	$0.0999_{\pm 0.0080}$	$0.1508_{\pm 0.0036}$	$0.2025_{\pm 0.0105}$	$0.2319_{\pm 0.0072}$	$0.2567_{\pm 0.0081}$
AutoData (CORE)	$0.1142_{\pm 0.0023}$	$0.1430_{\pm 0.0094}$	$0.2009_{\pm 0.0015}$	$0.2389_{\pm 0.0072}$	$0.2727_{\pm 0.0128}$
AutoData (VAL-BPB)	$0.1037_{\pm 0.0030}$	$0.1568_{\pm 0.0021}$	$0.1959_{\pm 0.0004}$	$0.2372_{\pm 0.0084}$	$0.2647_{\pm 0.0085}$

Key Findings

Human-designed curation pipelines struggle on a well-curated pool: DCLM-Baseline and PPL filtering regress val-bpb across all scales. RegMix performs better on GPT-2 depth=12 CORE but does not consistently improve either metric.
AutoData transfers across model scales: Recipes significantly outperform all baselines on val-bpb from depth 8 to depth 20. At depth 24, performance is similar to default ClimbMix, indicating a limit to cross-scale transfer.

Recipe Analysis

Algorithm 1 (val-bpb) reserves one-third of budget as a random backbone for coverage, then uses a composite score:

$s_{div} = \frac{1}{2}[\tanh\tilde{z}(div_{avg}) + \tanh\tilde{z}(div_5)]$
$s_{ppl} = \tanh\tilde{z}(-\log ppl)$ (gated by diversity)
$s_{len} = -|\tanh\tilde{z}(\log tok)|$
$s_{rat} = -|\tanh\tilde{z}(chars/tok)|$
Final score: $s = s_{div} + s_{ppl}\frac{s_{div}+1}{2} + s_{len} + s_{rat}$

Scores are shrinkage-centered within (topic, format) cells, then Gumbel tournament selection (slates of 3).

Algorithm 2 (CORE) reserves 50% as random backbone, uses pool-adaptive triangular kernel scoring on length, diversity, perplexity, and ratio features, adds a source-neighborhood prior, and uses Gumbel noise with Top- $B$ selection on the repair set.

Feature Usage Analysis (Table 2)

Feature	GPT-5.5	Gemini-3-Pro	Opus-4.7
Length	200	200	200
N-gram	200	199	200
Perplexity	200	199	200
Topic	178	2	187
Format	177	3	172

Feature Space Ablation (Table 3)

Feature subset	BPB ↓	Δσ ↑
Random sampling	0.95181 ± 0.00077	—
Lexical (2)	0.95043	+1.8
Perplexity (1)	0.95037	+1.9
Categorical (2)	0.95041	+1.8
LLM annotations (3)	0.95021	+2.1
All features (8)	0.95048	+1.7

LLM-annotation features yield the best result, suggesting semantic signals (factual errors, reasoning errors, reasoning steps) provide useful structure beyond cheap statistics.

Operation Usage Analysis (Table 4)

Axis	Sub-type	GPT-5.5	Gemini-3-Pro	Opus-4.7
Score aggregation	Single feature	0.0%	0.0%	0.0%
	Linear composite	13.0%	96.0%	85.5%
	Multiplicative composite	86.5%	0.5%	1.0%
	Gated/conditional	0.5%	3.5%	13.5%
Selection rule	Threshold	1.5%	6.0%	0.5%
	Top-B	6.5%	6.5%	12.0%
	Tournament	91.0%	0.0%	0.0%
	Stratified quota	1.0%	87.5%	87.5%

Theoretical and Practical Implications

Theoretical Implications

Data engineering as optimization: AutoData reframes pre-training data curation as an optimization problem over executable recipes, extending autonomous AI research to the data that shapes model learning.
Feature interaction discovery: The agent automatically discovers composite scoring rules and diversity-preserving mechanisms, showing that effective selection strategies move beyond single-feature ranking and threshold filtering.
Cross-agent diversity: Different LLM agents (GPT-5.5, Gemini-3-Pro, Claude-Opus-4.7) exhibit distinct search behaviors—GPT-5.5 prefers multiplicative composites and tournament selection, while Gemini and Opus favor linear composites and stratified quotas—revealing that the search space supports multiple valid solution styles.

Practical Implications

Cost-effective data curation: AutoData discovers effective selection recipes within an overnight search (~200 steps) on a small proxy model, making it practical for real-world use.
Transferability: Recipes discovered on small proxies transfer to larger scales (up to 1.3B parameters), suggesting that small-scale proxy search can inform larger pre-training runs.
Feature engineering guidance: The ablation study shows that a small, well-chosen feature set (even 1-2 features) can match the full 8-feature set, while LLM annotations offer additional headroom for future work.
Limitations to consider: CORE gains are not always statistically significant (high variance), and at depth 24 the advantage diminishes, indicating limits to cross-scale transfer.

Conclusion

AutoData demonstrates that data engineering is a natural next frontier for autonomous AI research. Instead of treating curation as a fixed, manually designed preprocessing pipeline, AutoData makes data selection searchable: agents discover how to score documents, combine signals, and preserve diversity using direct validation feedback.

Key takeaways:

Agentic search over executable selection algorithms is an effective and practical paradigm for pre-training data engineering
Discovered recipes outperform human-designed baselines on val-bpb with statistical significance across model scales
The approach extends autonomous research from model-side optimization to data engineering

Future directions identified by the authors:

Testing generalization to other corpora, domains, and larger model scales
Improving proxy objectives for agentic data selection (CORE vs. val-bpb yield different improvements)
Scaling to richer LLM annotations and larger data pools
Understanding the cost–quality trade-off across feature sources