OPUS: Towards Efficient and Principled Data Selection in Large Language Model Pre-training in Every Iteration

Summary

OPUS proposes a dynamic data-selection framework for LLM pre-training that scores candidates in the geometry the optimizer actually follows, rather than in raw-gradient space. The central conceptual move is to replace the SGD-flavoured utility of prior online-selection work (e.g. GREATS) with an optimizer-induced utility $\eta_t \langle \mathbf{P}_t \nabla \mathcal{L}(z), \mathbf{g}_{\mathrm{val}}\rangle$ minus a redundancy penalty, where $\mathbf{P}_t$ is a linearised approximation of the AdamW or Muon preconditioner at step $t$. To make this practical, the authors combine three ingredients: (i) BENCH-PROXY, an in-corpus retrieval procedure that assembles a benchmark-aligned proxy shard so the proxy gradient is stable and low-variance; (ii) a ghost-outer-product + CountSketch estimator that avoids materialising per-sample gradients and reduces inner products to $O(d_{in}+d_{out})$ for diagonal AdamW; and (iii) Boltzmann sampling with an in-step redundancy term to preserve diversity. Empirically, the paper reports gains across FineWeb, FineWeb-Edu (including a strict "select from score 3 while baselines train on score 4+5" comparison), and continued pre-training on Qwen3-8B / SciencePedia, plus complementary explorations of contamination robustness and static toxic-content filtering, all at $\sim 4.7\%$ selection overhead.

The framing is timely and the ghost + CountSketch scaffolding is a real efficiency contribution. The main conceptual limitation is that the core inner product $\langle \mathbf{P}_t g_z, g_{\mathrm{val}}\rangle$ treats candidate and proxy asymmetrically (candidate preconditioned, proxy raw), which equals the intended one-step improvement only when $\mathbf{P}_t$ is symmetric — not obvious for Muon's dense Newton–Schulz-induced operator. Empirically the story is mostly strong but rests on single-seed numbers on 774M/1.5B GPT-2 models, with WSC swings that meaningfully move headline averages, and BENCH-PROXY is explicitly retrieval-aligned to the evaluation benchmarks. Executing the release confirms that the OPUS scoring path is real once patched, but the compute-matched Random baseline never finished under the reproduction's budget, the measured selection overhead sat at 6.95–6.98% (vs. 4.7%), and two critical bugs on BENCH-PROXY construction plus a broken evaluation harness mean the shipped scripts cannot produce the paper's numbers as-published.

Strengths

• Conceptual contribution. Framing utility in the optimizer-induced update geometry — with explicit linearised preconditioners for both AdamW (diagonal) and Muon (dense, via a frozen NS operator) — is a clean improvement over SGD-style raw-gradient scoring, and unifies the treatment of two very different optimizers under one scoring machinery.
• Efficiency ingredients. The ghost-outer-product + CountSketch reduction is well-motivated, and the near-single-digit selection overhead reported in Figure 4 is meaningful if it holds up: that is essentially the minimum bar an online-selection method must clear to be usable at scale.
• Empirical breadth. The sweep covers two model families, two optimizers, two corpora, from-scratch and continued pre-training, and a broad benchmark suite (10 in-domain + 6 OOD + domain-wise PPL) — more than most data-selection papers offer.
• Approximation-fidelity study. Explicitly measuring the gap between each approximation and its exact counterpart at three checkpoints (Table 11) is exactly the right auxiliary study, even if the specific late-checkpoint Taylor number needs re-verification.
• Complementary explorations. The contamination experiment (Fig. 5) and the static toxic-filter demonstration (Jigsaw + Enron mixed into an OPUS buffer) probe the same score in a very different regime and are a nice touch.
• Released OPUS scoring pipeline. The repository implements the OPUS scoring pipeline (BENCH-PROXY construction, ghost outer product via GCLinear, CountSketch projection, Boltzmann sampling) with the exact Appendix E defaults, and the six pipeline components in Figure 2 are each identifiable in the code. The OPUS arm trains cleanly once the shipped bugs are patched (val loss falls monotonically 5.52 -> 3.92 across 11 evaluations).

Weaknesses

• Preconditioning asymmetry. The final score contracts a preconditioned candidate against a raw proxy gradient (Eq. 15). This equals the intended $\langle \mathbf{P}_t g_z, g_{\mathrm{val}}\rangle$ objective only when $\mathbf{P}_t$ is self-adjoint. That is fine for AdamW's diagonal $\mathbf{P}$, but the paper's own derivation says Muon's preconditioner is dense and coordinate-mixing (via a frozen Newton–Schulz operator), and symmetry is not established. The redundancy term is built symmetrically in $\phi$, further muddying the geometry. This is a substantive modelling choice on which the paper's "optimizer-induced" thesis rests, and it deserves an explicit justification or ablation.
• Proxy leakage to the evaluation suite. BENCH-PROXY explicitly retrieves the top-$M$ pre-training documents most similar to the target benchmark validation set. Even though the documents come from the pre-training corpus, selection is systematically pulled toward benchmark-similar features. The paper offers an OOD control (Table 4), but the six OOD tasks (BBH, RACE-M/H, AX-b/g, StoryCloze) are close cousins of the in-BENCH-PROXY tasks; there is no held-out benchmark family (e.g. code, math) or BENCH-PROXY built from a disjoint benchmark set.
• Overclaim in the "ensuring valid gradient estimation" phrasing. Retrieval by frozen-encoder similarity gives a stable, in-distribution proxy — not an unbiased estimator of the true $\mathcal{D}_{\mathrm{val}}$ descent direction. No guarantee is provided.
• Section 3.1 vs Appendix E disagreement on the proxy. §3.1 and Appendix D define $\mathcal{D}_{\mathrm{proxy}}$ as a BENCH-PROXY retrieval shard from the pre-training corpus and warn that using raw benchmark validation as the proxy is unstable, but Appendix E states GPT-2 runs "use the validation split as the proxy set for scoring". Which pool was actually used per experiment is not clear, and the ambiguity changes what Table 5's "Std. proxy" vs "OPUS" ablation is measuring.
• Single-seed comparisons on close numbers. All main tables report point estimates with no seed variance. Many between-method gaps are $\le 1$ pt, and WSC (285 samples) alone lifts the Table 1 GPT-2 XL / Muon average by nearly a full point in OPUS's favour (48.08 vs 36.54–38.46 for every other method) — nearly half of the headline +2.2 gap. Without per-seed intervals the "+2.2%" headline is hard to interpret.
• "Matches or exceeds 2$\times$ baselines" is not universally true. In Table 10 (GPT-2 Large / FineWeb-Edu) Random(60B) on score 4+5 reaches 42.15, above OPUS(30B) score 4+5 (42.04) and OPUS(30B) score 3 (41.72). The framing generalises an efficiency win that doesn't hold in every setting.
• Model-scale limitation. From-scratch experiments run on GPT-2 Large (774M) / XL (1.5B) for 30B tokens — small by modern LLM standards — and CPT on Qwen3-8B uses only 1.5B tokens. The data-wall framing invokes trillion-token web-scale training, but the empirical support does not clearly extrapolate there.
• Frozen-NS approximation for Muon is not stress-tested. Freezing the Newton–Schulz operator across candidates linearises an inherently non-linear step and depends on candidate directions being close to the reference $\bar q_{t,L}$. Table 11's Muon-preconditioner Util. Err. row does not decompose by candidate diversity or refresh cadence.
• Missing / garbled derivations. Section 3.1 shows stub sentences with empty equation slots ("Substituting Eq. (5) into Eq. (4) yields" is not followed by an equation), so the reader cannot verify how the isotropic-Hessian assumption yields the specific $\eta_t$ / $\eta_t^2$ redundancy coefficient.
• Underspecified pseudocode. Algorithm 1 jumps from step 4 (candidate features) to lines 12–14 (sample, add, update history) with no visible inner-loop header, no line for computing $U_z^{(t)}$, and no initialisation of $\tilde B_t$ or $\Phi^{(t,r)}$; two end for statements are present against one for. The support and replacement rule for the Boltzmann softmax is also not stated.
• Notation inconsistencies. The desideratum "overhead remains modest as model dimension $m$ grows" mislabels $m$ (the sketch dimension) as the model dimension and inverts the scalability argument. Two Table 8 sub-block captions write $d = 8192$ where $m = 8192$ is intended.
• Approximation-fidelity claim vs. printed numbers. The prose says Taylor-surrogate Util. Err. "decreases as training progresses"; as printed in Table 11 the late-checkpoint value ($\approx 5.3\times 10^{-1}$) is the largest in the table.
• BBH inconsistency. "All benchmarks are evaluated at zero-shot settings" (§4.1) is contradicted by Appendix F's three-shot BBH protocol.
• CountSketch ranking overclaim. CountSketch unbiasedness is a first-moment identity and does not, by itself, imply preserved rankings at $1250\times$ compression.
• Qualitative-figure ambiguities. In the toxic-filter demonstration, the reported "mean score = +0.548" for rejected samples is the mean over the eight shown samples, not the 24-sample rejected group. The appendix's per-method sample panels also lack explicit method headers on some blocks, making cross-method score scales look like contradictions.

Reproducibility & code

• OPUS scoring path runs after repair. Once patched, the shipped OPUS pipeline trains a GPT-2 XL / Muon run cleanly (val loss 5.52 -> 3.92 across 11 evaluations) with the Appendix E defaults, so the OPUS row itself is mechanically reproducible given data + compute.
• Measured selection overhead is 6.95–6.98%, not 4.7%. Instrumenting the OPUS training loop on a 2 A100 / 0.1B-token run yielded a per-step breakdown (42,733-row timing CSV) that resolves to opus_selection_total = 7.44% over all steps, 6.98% for step > 1000, a 1.075$\times$ wall-clock slowdown — a 48% relative overshoot vs. the paper's 4.7%. This is most likely a scale artefact (the paper measured at 8-GPU / 30B tokens, where the Muon optimizer step composes the denominator very differently) rather than a methodological error, so we call it inconclusive rather than refuted. The companion "over 3.5$\times$ naive slowdown" comparison could not be measured at all: the release contains no naive per-example-gradient scoring path to time against.
• Baseline scorers not shipped. The training entry point exposes only opus, ppl, and random strategies. QuRating, DSIR, GREATS, DCLM-FastText, FineWeb-Edu classifier, and UltraFineweb — every non-trivial baseline in Tables 1, 2, 9, 10 — must be stood up externally. Reproducing head-to-head numbers therefore depends on out-of-repo infrastructure.
• Two critical bugs on the BENCH-PROXY path. data/build_betr_proxy.py pre-filters candidates per file with hardcoded caps (keep_fraction=0.005, min_keep=2000, max_keep=20000) and no CLI knobs, so a 30M-token request produced only 11.07M tokens while reporting success; the per-file quota is also a different selection rule than the corpus-wide top-scored selection described in Appendix D. Separately, train.py's load_multi_file_proxy_tokens splits the proxy budget evenly per file regardless of file size, so a 30M-token request against a 29,997,151-token shard plus a 10,839-token spillover shard loaded only half the proxy. Two independent silent-failure bugs on the object the method is defined against strongly suggests the code has not been exercised end-to-end at the reported scale.
• Evaluation harness broken as shipped. eval/run_lighteval.sh set PYTHONPATH to eval/ so import model fails, resolves PROJECT_ROOT one directory too high so all nine benchmark datasets report "not found" while the script exits 0 and writes a results file, and eval/lighteval.py multiplies accuracy_norm by 100 twice yielding average accuracies of 640. Each fix is one or two lines, but together they mean the eval script as published could not have produced any number in the paper. The task list also covers only nine in-proxy tasks — ANLI, WSC, and the entire OOD suite (BBH three-shot, RACE-M/H, AX-b/g, StoryCloze) are not implemented; the paper simply refers users to OpenCompass, which is not vendored.
• Domain-wise PPL evaluator, contamination and toxic-filter drivers, fidelity-study references, and Figure plotting scripts are absent. The WebOrganizer topic classifier + per-domain 1K held-out shards + per-domain PPL scorer needed for Table 3 are not shipped (train.py only logs an aggregate validation loss). No script constructs the 80%-permuted contaminated FineWeb-Edu pool for Figure 5, and no Jigsaw / Enron loaders or OPUS-as-static-filter driver exist for the toxic-detection demonstration. The approximation-fidelity study (Table 11) requires exact-reference implementations (materialised Hessian action, exact per-sample rank-1 gradients, exact validation gradient at virtual parameters, per-candidate NS) that are not in the release. Plotting scripts for Figures 1, 4, 5, 6, and 7 are also absent.
• CPT pipeline. cpt/run_qwen3_cpt.sh expects a preprocessed SciencePedia bin directory, but the download / tokenisation / loss-mask pipeline for SciencePedia is not scripted, and the dataset is not on the HF Hub. No OlympicArena or SciAssess (3-shot chain-of-thought) evaluators are shipped. This makes the CPT results in Figure 3 / Figure 7 unreproducible from the release even in principle, without third-party data.
• Central 2.2% / 8$\times$ headline was not verified end-to-end. The compute-matched Random GPT-2 XL / Muon baseline needed for the headline comparison did not complete in the reproduction (a sequence of ENOSPC, preemption, and a double-launch incident capped it at step 322–791 of 6104), so there is no Random checkpoint to compare against. This is a compute limitation of the reproduction, not a paper fault, but it does mean that neither the +2.2% accuracy claim nor the 8$\times$ compute claim was confirmed or refuted in this executed pass.

Recommended Changes

Essential.

• Justify the candidate/proxy preconditioning asymmetry. Add a paragraph in §3.1 explaining why $\langle \mathbf{P}_t g_z, g_{\mathrm{val}}\rangle$ is retained as the target when $\mathbf{P}_t$ (Muon) is not obviously symmetric, and ablate the alternative (preconditioning the proxy) on at least one Muon setting. (Maps to the preconditioning-asymmetry weakness.)
• Report per-seed variance for Tables 1, 2, 4, 5, 8, 9, 10. At least three seeds for the headline configurations, with std / bootstrap CIs on both per-benchmark and average columns. (Maps to single-seed comparisons and the WSC-driven margin.)
• Add a stricter OOD test for BENCH-PROXY. Report Table 4 with a held-out benchmark family (e.g. code / math) that is disjoint from the tasks used to build BENCH-PROXY, or run OPUS with a BENCH-PROXY built from a disjoint benchmark set and evaluate on the current suite. (Maps to the proxy-leakage weakness.)
• Restore the missing derivation and pseudocode details. Print the intermediate Taylor expansion that yields the $\eta_t$ / $\eta_t^2$ coefficients under $\mathbf{H}_{\mathrm{val}} \approx \mathbf{I}$, and add the inner-loop header, utility-computation line, initialisations of $\tilde B_t$ and $\Phi^{(t,r)}$, and the softmax support / replacement rule to Algorithm 1. (Maps to missing-derivations and underspecified-pseudocode weaknesses.)
• Reconcile the §3.1 vs Appendix E proxy definition. State, per experiment, whether $\mathcal{D}_{\mathrm{proxy}}$ was the BENCH-PROXY retrieval shard or the raw validation split, and correct the ablation labels in Table 5 to match. (Maps to the "§3.1 vs Appendix E disagreement" weakness.)
• Fix the shipped BENCH-PROXY and evaluation code. Expose data/build_betr_proxy.py's per-file caps as CLI knobs (or switch to true corpus-wide top-$M$), fix load_multi_file_proxy_tokens to allocate the proxy budget by file size, and fix eval/run_lighteval.sh (PYTHONPATH, PROJECT_ROOT) and the double $\times 100$ multiplier in eval/lighteval.py. (Maps to the two Reproducibility bullets on BENCH-PROXY bugs and the broken eval harness.)
• Ship baseline scorers or their outputs. Vendor scoring code for the eight comparison baselines (QuRating, DSIR, GREATS, DCLM-FastText, FineWeb-Edu classifier, UltraFineweb, PPL, Random), or release the selected-shard token IDs / masks so a third party can rerun training on the same data selections. (Maps to the baseline-scorers Reproducibility bullet.)
• Vendor the evaluation harness. Add ANLI, WSC, and the six OOD benchmarks (with the three-shot BBH curated subtask list) to the shipped evaluator, either by extending lighteval.py or by pinning an OpenCompass commit + config. (Maps to the eval-harness-not-shipped Reproducibility bullet.)
• Report overhead at more than one compute scale. Either publish per-step timing at both the paper's 8-GPU / 30B-token setup and a smaller cluster, or fully specify the hardware / precision / step-count configuration under which "4.7%" was measured, so downstream users can predict overhead at their own scale. (Maps to the measured-overhead-mismatch Reproducibility bullet.)
• Verify the late-checkpoint Taylor error in Table 11. Either correct the number if garbled, or add Spearman rank correlation between exact and approximated scores at each checkpoint so "ranking-preserving" is empirically grounded. (Maps to the approximation-fidelity-claim weakness.)

Suggested.

• Qualify the "2$\times$ baseline" efficiency framing. Enumerate the settings where Random(60B) is or is not surpassed by OPUS(30B) — Table 10 GPT-2 Large / FineWeb-Edu is one where it is not. (Maps to the "matches-or-exceeds" weakness.)
• Add a Muon frozen-NS sensitivity analysis. Sweep refresh cadence, or compare frozen-NS vs per-candidate NS on a small candidate pool, and report how Util. Err. scales with candidate diversity. (Maps to the frozen-NS weakness.)
• Spell out the CountSketch cost derivation. State the hashing assumptions (independent hash/sign per mode) needed to reach $O(d_{in}+d_{out})$ for diagonal-preconditioned rank-1 gradients. (Maps to the CountSketch-cost weakness.)
• Fix notation. Replace "model dimension $m$" with "model parameter dimension $d$" (or "sketch dimension $m$" as appropriate), and correct "$d = 8192$" to "$m = 8192$" in the Table 8 sub-block captions. (Maps to the notation-inconsistency weakness.)
• Qualify the CountSketch ranking claim. Weaken "preserving the ranking signal" to an empirical claim on the evaluated settings, or provide a variance / concentration argument. (Maps to the CountSketch-ranking-overclaim weakness.)
• Clarify the toxic-filter figure. Label the "mean score" as the mean over the eight shown samples, add method headers to every appendix qualitative-example block, and note that score scales are not comparable across methods. (Maps to the qualitative-figure ambiguities weakness.)
• Qualify the main-text "zero-shot" statement. Add the "except BBH (three-shot on a curated subset)" caveat where the zero-shot claim is first made. (Maps to the BBH inconsistency weakness.)
• Ship the contamination, toxic-filter, fidelity-study, domain-PPL, and CPT drivers. Add the 80%-token-permutation pipeline, Jigsaw / Enron loaders, exact-reference scoring for the fidelity study, the WebOrganizer + per-domain PPL evaluator, and a SciencePedia downloader / tokenizer / loss-mask generator plus OlympicArena / SciAssess evaluators. Also add plotting scripts for Figures 1, 4, 5, 6, 7. (Maps to the Reproducibility bullets on qualitative experiments, domain PPL, and the CPT pipeline.)
• Acknowledge scale scope. State explicitly in §4.1 or §5 that from-scratch results are up to 1.5B params / 30B tokens, and that trillion-token / larger-model extrapolation is left to future work. (Maps to the model-scale-limitation weakness.)