DataRater: Meta-Learned Dataset Curation

TL;DR

DeepMind paper (Calian, Farquhar, Silver, Dean, et al.) that turns the data-curation pipeline into a meta-learned component: a small transformer “DataRater” scores each training point, and its parameters are optimized by backpropagating through inner-model updates to minimize held-out loss on the same dataset. With MixFlow-MG bilevel-optimization tricks, they meta-train a ~50M DataRater against a population of ~400M inner Chinchilla LMs in feasible compute and then use top-K filtering at batch time to train downstream 50M–1B LMs. On low-quality corpora (the Pile, C4/noclean) DataRater filtering gives net FLOPS savings at matched performance — i.e. the filter pays for itself — and discard-fraction selected at 50M transfers cleanly to 1B.

Key claims

The curation pipeline is reformulated as a continuous bilevel optimization: filtering weights come from a softmax over a learned scorer s_θ(x), and θ is updated by meta-gradients of held-out training loss after T inner SGD steps [§2.1, §2.2].
Inner and outer losses share the same functional form (next-token cross-entropy on a disjoint held-out slice of the same dataset) — the DataRater is not trained against a downstream benchmark, just against “more efficient training on the original distribution” [§2.2 Meta-learning objective].
Meta-training one DataRater costs ~10× the FLOPS of training a single 1B LM, but amortizes across all downstream uses of the filtered dataset [§3.3].
Compute-to-match-baseline: DataRater filtering yields large net savings on the Pile and C4/noclean; on the more pre-filtered C4 the savings disappear or invert (training on the filtered set does not always reach the baseline) [Fig. 1, §3.2].
Discard-fraction optimum is dataset-dependent but model-size-invariant: 10% for C4, 25% for C4/noclean, 75% for the Pile, identical across 50M / 150M / 400M / 1B target models [§3.4, Fig. 4].
Cross-scale generalization: a single DataRater meta-trained with 400M inner models improves training across 50M–1B target scales on 11/14 input-dataset × scale × metric combinations [§3.5, Fig. 6].
Qualitative inspection: low-scored Pile samples concentrate on encoding errors, OCR garbage, all-caps English, whitespace-heavy text, high-entropy lists/tables, SSH keys, and multilingual content (the Pile is >97% English) [§3.6, Fig. 8].
The DataRater also implicitly re-weights the mixture-level composition of the Pile as the discard fraction rises, not just point-level filtering [§3.6, Fig. 7b].

Method

DataRater is a non-causal transformer s_θ(x) (≈50M params) that scores tokens / sequences; per-batch weights are obtained by softmax over scores within an oversampled minibatch of size B/(1-α), where α is the discard fraction, and then top-K filtering keeps the highest-scored B. Meta-optimization backpropagates loss on a held-out slice of the same dataset through T inner SGD steps of a population of inner LMs (~8 copies, ~400M params each), with periodic resets of inner models to span learning-progress stages and per-inner-model Adam meta-gradients averaged into θ. The bilevel computation uses MixFlow-MG (Kemaev et al. 2025) — mixed-mode differentiation + gradient checkpointing — to keep HBM tractable while backpropagating through several inner updates. At inference, top-K filtering at batch level is equivalent to thresholding the per-point score against the empirical CDF, so the same filter parameterizes a massively parallel data-processing pipeline (Apache Beam).

Results

On the Pile, DataRater filtering produces net FLOPS savings approaching ~40% at the 1B scale on the headline metric, after accounting for both DataRater meta-training (one-time) and the per-step filtering compute overhead [Fig. 1, §3.3].
C4/noclean: substantial savings on NLL and downstream HellaSwag/PIQA/ARC-E/SIQA/CommonsenseQA at 25% discard, transferred from a 50M selection sweep [Fig. 4, Fig. 6].
C4: the filter does not consistently match baseline — C4 is already heavily pre-filtered, so the learned filter has little headroom and in some cells underperforms [Fig. 1, §3.3 (left of “undefined” entries)].
Discard-fraction optima are stable across 50M → 1B inner sizes (50M sweep is sufficient to set α for 1B training) [Fig. 4].
Cross-scale transfer: 11/14 (dataset × scale × metric) cells improve over the unfiltered baseline; NLL gains are more consistent than downstream-accuracy gains, which show higher variance [§3.5, Fig. 6].
Training the DataRater itself: ~10% of the FLOPS of one 1B LM training run [§3.3].

Why it’s interesting

DataRater is the cleanest learned-filter counterpart to A Bitter Lesson for Data Filtering on the wiki: where the Stanford paper argues that at high compute the best filter is no filter (because human heuristics are scale-fragile), DataRater argues that the filter just has to be learned — and reports the strongest gains exactly where the bitter-lesson paper concedes filtering helps least at small scale (the Pile, C4/noclean, the low-quality regime). The two papers don’t have a clean head-to-head — DataRater never tests in the >>1B regime that the bitter-lesson paper claims dominates — but they stake out the two ends of the curation-gate debate that Synthetic Training Data has been collecting evidence on. It’s also a procedural cousin of Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training (also automated mixture/selection optimization with held-out validation as the loss) and of Shaping capabilities with token-level data filtering (also sub-document granularity, but DataRater’s gradient signal is task-agnostic where token-filtering’s is capability-shaping), and supplies a third concrete recipe for the future pretraining-data-curation concept page that the synthetic-training-data TL;DR has been flagging.