OptiSelect: How does the Optimizer Shape Data Curriculum?

Introduction

• OptiSelect scores samples in an oversampled candidate pool by optimizer-aware marginal utility at each step, selecting the highest-utility batch for training.Because utility is realized through the induced model update, scoring accounts for the optimizer preconditioner that reshapes the raw gradient.Source: §1 · p. 1

Contributions

• The framework formalizes an optimizer-aware transformation operator and derives utility functions for five optimizers spanning distinct preconditioning geometries.The study considers AdamW, AdEMAMix, Sophia, Lion, and Muon.Source: §1 · p. 1
• Selection gain scales with utility-score discriminability: sign collapse bounds Lion’s gain and radial collapse bounds D-Muon’s, while AdamW and Sophia retain larger score spread.The same principle derives an optimal candidate oversampling ratio of λ⋆= 2.Source: §1 · p. 2
• Pretraining on 124M and 720M models agrees with the theory; AdamW’s scoring geometry prevails under D-Muon, and selection benefits persist under format rephrasing.Decoupling scoring geometry from the base optimizer shows that the optimizer-matched geometry is not necessarily optimal.Source: §1 · p. 2

Data selection for LLM pretraining

• LLM pretraining data selection is dominated by static corpus-curation pipelines, while online methods include MATES, GREATS, and OPUS; earlier curricula and sampling methods did not explicitly account for the optimizer.Static examples include DSIR, DCLM, FineWeb-Edu, DoReMi, and DoGE; OPUS computes optimizer-aware scores.Source: §2 · p. 2

Optimizer-aware influence

• Influence methods have progressed from influence functions, TracIn, and TRAK to Adam-state-corrected LESS and TrackStar and optimizer-aware OPUS, but prior work does not compare the influence geometries induced by different optimizers.OPUS applies its attribution scores to per-step sample selection during LLM pretraining.Source: §2 · p. 2

Optimizers for LLM pretraining

• Optimizer research covers diagonal-adaptive, curvature-aware, sign-based, and matrix-geometry methods, but existing benchmarks focus on standard pretraining and no prior work examines how optimizer algorithms affect data-selection preferences.Benchmark rankings can shift with model scale and batch size; this work studies optimizer effects on selection theoretically and empirically.Source: §2 · p. 2

3.1 DATA SCHEMES IN LLM PRETRAINING

• At each step, OptiSelect draws λB candidates and selects the top B by optimizer-aware influence for the model update.The optimizer also defines the attribution scores used for selection.Source: §3.1 · p. 3

3.2 OPTIMIZER-AWARE INFLUENCE SCORING

• Using a representative held-out validation set, OptiSelect scores optimizer-state-conditioned directions by predicted loss reduction, combining target alignment with redundancy approximated using an identity Hessian.The scoring operator need not coincide with the actual optimizer update map.Source: §3.2 · p. 3
• AdamW and AdEMAMix use inverse-root second-moment scaling, Sophia diagonal curvature, Lion its momentum buffer, and Muon a momentum-matrix derivative or cheaper proxy.The scoring operator is conditioned on optimizer state and varies by optimizer family.Source: §3.2 · pp. 3–4

Selection-gain principle

• Across 15 pretraining runs, normalized score gain stayed close to c_q and was nearly invariant across learning rates.Figure 1 also shows measured gain close to c_q across optimizers, with Lion decaying after proxy refresh.Source: Figure 1 · p. 4, §4 · p. 4
• Theorem 1 bounds OptiSelect’s gain through utility-score discriminability: absolute gain depends on score spread σ_U(t), while c_q depends on retained fraction q = 1/λ.The principle is analyzed under a Gaussian model of candidate utilities.Source: §4 · p. 4

Sign Collapse

• Lion’s frozen sign pattern preserves SGD score variance but cannot amplify discriminability, while momentum dominance can make the exact score’s spread collapse.This can leave the selector with little ranking information, approaching uniform selection.Source: §4 · p. 5

Diagonal-adaptive advantage

• Theorem 3 ranks exact and linearized sign-score spreads below AdamW and Sophia, which each have strictly better discriminability bounds than sign-based scoring.The comparison does not order AdamW against Sophia.Source: §4 · pp. 4–5

Radial collapse in the ideal polar geometry of Muon

• In ideal polar Muon geometry, candidates differing only radially receive the same first-order score, while tangential differences are reweighted by the momentum spectrum.The practical right-preconditioned proxy can distinguish radial differences; the theorem gives no unconditional variance ordering against SGD or radial-collapse guarantee for that proxy.Source: §4 · p. 6

5.1 THE OptiSelect BENCHMARK: SELECTION GAIN ACROSS OPTIMIZERS AND SCALE

• OptiSelect improves downstream loss with AdamW, AdEMAMix, and Sophia at both scales, but not Lion; D-Muon’s gain shrinks by training’s end at 124M and persists at 720M.Despite limited selection gains, D-Muon consistently achieves the best averaged downstream accuracy.Source: Table 1 · p. 6, Figure 2 · p. 7, §5.1 · p. 7
• With AdamW, its matched scoring geometry performs best; with D-Muon, AdamW’s diagonal-adaptive geometry outperforms the matched geometry, while SGD scoring beats most alternatives.The AdamW second-momentum diagonal-adaptive operator is the best scoring geometry in both sweeps; a mismatched geometry can hurt selection performance.Source: Figure 4 · p. 7, §5.1 · p. 8

5.3 DOES SELECTION HELP WITH REPHRASING?

• OptiSelect improves downstream loss and average downstream accuracy across rephrasing mixtures, indicating that selection and rephrasing are not substitutes or in conflict.The experiments use mixtures with 0%, 50%, and 100% rephrased documents.Source: Figure 5 · p. 8, Figure 6 · p. 8, §5.3 · p. 8
• Across rephrasing mixtures, selection yields similar downstream-loss gains, but its accuracy benefits vary with the evaluation harness’s formatting pattern.Format mismatch is a confound for downstream accuracy, and its effect depends on the harness pattern.Source: §5.3 · p. 8, Figure 6 · p. 8

5.4 MORE CANDIDATES, BETTER PERFORMANCE?

• Under Gaussian candidate utilities, gain per streamed token has a unique maximum at λ⋆ = 2; selection gain grows monotonically with λ but not enough to offset extra costs past that peak.The analysis frames λ as a trade-off between per-step selection gain and the overhead of streaming and scoring more tokens.Source: §5.4 · p. 9, Figure 7 · p. 9

• AdamW’s variance normalization favors gradients large along rare dimensions, while Muon’s matrix-sign geometry favors rotations of singular subspaces over singular-value changes.The authors identify designing optimizers that adapt to the data curriculum as future work.Source: §6 · p. 9
• OptiSelect benefits 3 of 4 out-of-domain benchmarks but yields little Commonsense QA gain; it leads at equal applied tokens and trails at equal tokens read at both model sizes.Selection converts extra forward–backward compute into fewer optimizer steps for the same quality, not fewer tokens consumed; the authors recommend reporting both token matchings.Source: §6 · p. 9
• Evidence is limited to 124M and 720M models, single-seed runs, and reasoning proxies within small-model capabilities; Math, Code, and Agentic tasks remain for future study.Behavior at multi-billion scale is therefore extrapolated rather than directly tested.Source: §6 · p. 9

7 CONCLUSION

• OptiSelect formalizes optimizer-dependent online in-batch data selection as a training paradigm.This work studies how different optimization geometries affect data-selection preferences.Source: §7 · pp. 9–10
• The study combines theoretical analysis and empirical benchmarking to examine how optimization geometries affect data-selection preferences.The authors position these findings as foundations and insights for the next generation of pretraining optimizers co-designed with modern data deployment.Source: §7 · pp. 9–10

AI DISCLOSURE

• Generative AI tools assisted with paper polishing, proof writing, literature analysis, and coding.Source: AI DISCLOSURE · p. 11
• The authors did not use generative AI to formulate research or proof ideas, develop theoretical models or conceptual frameworks, or generate synthetic data; synthetic evaluation was not applicable.Source: AI DISCLOSURE · p. 11
• The authors reviewed all AI-assisted work and take responsibility for the final content, including AI-assisted text, claims, and artifacts.Source: AI DISCLOSURE · p. 11

ETHICS STATEMENT

• The study uses publicly available datasets and evaluation protocols, and the authors identify no additional ethical concerns specific to it.Source: ETHICS STATEMENT · p. 11

A EXTENDED RELATED WORK

• Prior work either selects data with the optimizer treated as given or studies optimizers with the data stream treated as given; OptiSelect addresses their intersection.The paper says the two literatures barely cite each other.Source: A EXTENDED RELATED WORK · p. 17

A.1 ONLINE DATA SELECTION AND CURRICULUM LEARNING

• Earlier curricula and static filters score losses, gradients, or corpus quality, whereas OPUS (Wang et al., 2026) scores online with optimizer-aware influence but validates only AdamW and Muon.FinePhrase rephrasing composes with online selection, whose gain is essentially unchanged at 0%, 50%, and 100% rephrased shares.Source: §A.1 · p. 17

A.2 INFLUENCE FUNCTIONS AND DATA ATTRIBUTION

• LESS (Xia et al., 2024) and TrackStar (Chang et al., 2025) account for Adam’s adaptive state, while OptiSelect prospectively compares optimizer-induced influence geometries.Unlike retrospective attribution, OptiSelect scores candidates before their update with the optimizer state frozen.Source: §A.2 · pp. 17–18

A.3 MODERN OPTIMIZERS FOR LLM PRETRAINING

• Standard-training optimizer benchmarks omit selection, although the benefit each optimizer draws from selection has its own ordering, stable across scale and unpredicted by standard-training quality.The cited benchmark literature describes families including diagonal-adaptive, sign-based, polar, and Kronecker-factored optimizers.Source: §A.3 · p. 18

B BENCHMARKED OPTIMIZERS: UPDATE RULES AND FROZEN-STATE OPERATORS

• Frozen-state scoring separates candidate-dependent directions from fixed terms; SGD’s resulting operator is the identity map on raw gradients.The framework distinguishes linear scoring surrogates from exact nonlinear optimizer updates.Source: B BENCHMARKED OPTIMIZERS: UPDATE RULES… · p. 19, §B.1 · p. 19

B.2 ADAMW

• AdamW uses a frozen-denominator diagonal map that upweights coordinates with smaller second moments; AdEMAMix has the same functional form at its own frozen second moment.The maps need not yield identical scores because the optimizers can follow different parameter trajectories and second-moment estimates.Source: §B.2 · p. 20, §B.3 · p. 20

B.4 SOPHIA

• Sophia’s scoring surrogate freezes its curvature estimate and omits clipping, weighting coordinates inversely by estimated curvature subject to a numerical floor.Its per-sample decomposition is exact for the linear surrogate, not generally for the clipped optimizer update.Source: §B.4 · p. 21

B.5 LION

• Lion and Signum use fixed-sign linear scoring surrogates, but Lion’s exact sign-score collapse bound does not directly apply to its implemented surrogate.Lion’s candidate dependence can be rare sign flips under the stated momentum-dominance and Gaussian assumptions; Signum matches Lion’s operator form under matched coefficients.Source: §B.5 · p. 21, §B.6 · p. 22

B.7 D-MUON

• Muon’s ideal polar derivative scores singular-subspace rotations and new column-space directions, while the implemented right-preconditioned proxy retains radial components.The proxy agrees with the ideal derivative on the column-space complement only under full-column-rank reference matrices and zero regularization.Source: §B.7 · p. 22
• Table 2 distinguishes exact Lion/Signum sign maps from amplitude-retaining fixed-sign surrogates, Sophia’s unclipped surrogate, and Muon’s ideal polar derivative versus implemented proxy.The implemented Muon proxy is linear and retains radial components; magnitude information means continuous amplitude rather than norm preservation.Source: Table 2 · p. 23

B.8 CLASSIFICATION OF THE FROZEN-STATE OPERATORS

• Frozen linear scoring maps decompose batch directions exactly into per-sample directions, and fixed references also yield linear scalar scores.This holds for SGD, AdamW, AdEMAMix, and Sophia’s unclipped scoring operator.Source: §B.8 · p. 23
• Exact clipped Sophia, coordinate-wise sign maps, and finite Muon matrix-sign are nonlinear; linear proxies decompose only the chosen scoring map, not the optimizer update.The Muon analysis proxy is first-order, while deployed sign and Muon proxies are heuristic approximations.Source: §B.8 · p. 23

C.1 PROOF OF THEOREM 1

• Under local smoothness, comparing one-step loss reductions is equivalent up to controlled second-order error to comparing retained batches’ expected mean utilities.The proof connects loss reduction to utility expectations and evaluates uniform sampling through candidate utilities and their order statistics.Source: §C.1 · pp. 24–25
• At σU(t)=0, utilities are identical and both sampling rules have zero selection gain; for a finite candidate pool, fixed-q cq,N=cq+O(N^-1/2), and cq is invariant under positive affine utility-score transformations.Source: §C.1 · pp. 25–26
• For fixed retained fraction q and positive mean utility, the relative one-step benefit is 1 + cq CV[U(·, t)] + o(1), subject to negligible remainder terms.The comparison is against uniform sampling for a fixed optimizer and scoring geometry; score spread alone does not establish ranking fidelity or alignment with the validation objective.Source: §C.1 · pp. 26–27
• The loss-reduction argument requires the same frozen linear map to score candidates and form the batch direction; nonlinear updates need attribution error O(η_t^2) for the conclusion to extend.For deployed fixed-sign and right-preconditioned proxies, the theorem describes the additive scoring geometry’s predicted gain unless update mismatch is separately bounded.Source: §C.1 · p. 27

D DETAILED PROOFS FOR Discriminability (§4)

• Appendix D provides detailed proofs and supporting assumptions for the discriminability theorems, distinguishing fixed-sign linear scoring from the exact nonlinear Lion score.The fixed-diagonal spread identity applies to the former, whereas the momentum-dominance bound for the exact score does not itself bound variance of the implemented linearized score.Source: D DETAILED PROOFS FOR Discriminability… · p. 28, U Lion · p. 31

D.1 SIGN COLLAPSE (THEOREM 2)

• For any fixed sign pattern, the linearized sign score has the same variance as the SGD score, so it is spread-neutral.This fixed linear diagonal variance identity does not apply directly to the exact nonlinear Lion score.Source: §D.1 · pp. 28–29

Proof of Proposition 2. Step 1: define coordinate-wise sign flips

• Under Gaussian candidate modeling, coordinate sign-flip indicators are independent, and momentum dominance lets the flip probability be bounded using Gaussian tails.The proof defines each flip event as a candidate gradient reversing the momentum buffer’s sign at that coordinate.Source: Proof of Proposition 2. Step 1: define… · p. 30, U Lion · pp. 30–31

U Lion

• For exact Lion scoring, candidate-dependent variation comes from coordinate-wise sign reversals, and the proof establishes the stated collapse bound under momentum dominance.The same conclusion applies to Signum after replacing β1 with its momentum coefficient µ.Source: U Lion · pp. 29–31, §D.1 · p. 29

D.2 DIAGONAL-ADAPTIVE ADVANTAGE OF AD A MW AND SO P H I A (THEOREM 3)

• AdamW preconditioning gives a spread ratio of at least 1 − O(ε′) under weighted variance normalization and non-constant second moments.The assumption is described as empirically mild because gradient clipping caps incoming squared gradients and the validation proxy emphasizes rare-feature coordinates; without the normalization, the conclusion can fail.Source: §D.2 · pp. 31–32

πSophia

• Sophia’s unclipped score amplifies spread through squared inverse curvature on flat directions, given analogous curvature normalization and ε ≪ min_j γ[ht]j.The amplification concentrates on directions with small curvature estimates where validation-relevant weights lie; the condition is measurable during training.Source: §D.2 · p. 32, πSophia · p. 32

D.3.1 PRELIMINARIES: THE IDEAL POLAR MAP

• For full-column-rank A, the ideal polar factor msign(A) uniquely maximizes ⟨A, D⟩F over ∥D∥op ≤ 1.The result assumes the ideal polar-factor map, not the conventional primary matrix-sign function or the finite Newton–Schulz computation.Source: §D.3.1 · p. 33

D.3.2 POLAR DERIVATIVE AND RADIAL GEOMETRY

• The ideal polar derivative has an exact radial nullspace: radial perturbations leave the polar factor unchanged, while the derivative is self-adjoint.For arbitrary candidates differing radially, equality is guaranteed only for first-order scores at the frozen state.Source: §D.3.2 · pp. 33–34

D.3.3 SCORE VARIANCE AND PROOF OF THEOREM 4

• The general score-variance identity requires no coordinate-uncorrelation assumption, but coordinate formulas require pairwise uncorrelated entries in the fixed singular-vector basis.There is no unconditional ordering against SGD; for the stated example the variance ratio is α²/s², exceeding one when 0 < s < α.Source: §D.3.3 · p. 35

D.4 RIGHT-PRECONDITIONED MU O N PROXY AND EXACT ERROR

• The right-preconditioned Muon proxy does not generally preserve radial invariance, even though the ideal polar map does.At ε = 0 it agrees with the derivative on the column-space complement but generally not on the parallel block; finite-score variance claims also require retaining rescaling and controlling the local Taylor approximation.Source: §D.4 · p. 36

D.5 OPTIMAL CANDIDATE MULTIPLIER (LEMMA 1)

• The candidate-multiplier efficiency curve is symmetric under λ ↦ λ/(λ − 1), and strict unimodality places its extremum at λ = 2.The symmetry gives f(1.5) = f(3) and f(1.25) = f(5).Source: §D.5 · p. 37

E.2 TRAINING ARGUMENTS

• Both standard-training and OptiSelect runs at 124M and 720M were trained from scratch for 128k steps, consuming 4.2B tokens total, with global batches of 32,768 and 131,072 tokens, respectively.Source: §E.2 · p. 38

E.4 LEARNING RATE SWEEP

• The optimal learning rate transfers to OptiSelect: all seven optimizers choose the same rate under standard training and selection at 124M, and five optimizers’ 720M standard-training rates are applied directly.The 720M search used four learning-rate grids at 16 GPUs, 256 sequences per step, and 2k steps (262M tokens).Source: §E.4 · p. 38, Figure 8 · p. 39, Figure 9 · p. 39

E.5 PROXY REFRESH INTERVAL

• At 124M with AdamW, proxy-refresh benefits saturate by ∆=1000, so the study sets that interval as default; refreshing is described as essentially free.At ∆=500, the reported result is 3.7178 vs. 3.7165 nats; caching trades per-step cost against staleness bias.Source: §E.5 · p. 40, Table 5 · p. 40

F EXTENDED EXPERIMENTAL RESULTS

• The extended experimental analysis covers results from Section 5 at both 124M and 720M model scales.Source: F EXTENDED EXPERIMENTAL RESULTS · p. 41

F.1 COMPLETE SCORES AT 124M

• At 124M, OptiSelect improves downstream loss for every optimizer while largely preserving optimizer rankings; AdamW and AdEMAMix lead, and held-out FineWeb-HQ rankings remain unchanged.Lion and D-Muon rank lower despite improving raw downstream performance.Source: §F.1 · p. 41, Table 6 · p. 41
• At equal trained tokens, OptiSelect beats standard training on in-domain and out-of-domain accuracy, but trails at equal read tokens; AdEMAMix leads in-domain and AdamW out-of-domain.Lion underperforms the standard baseline on average in-domain accuracy, while Sophia underperforms on average out-of-domain accuracy.Source: §F.1 · p. 41, Table 7 · p. 42

F.2 COMPLETE SCORES AT 720M

• The 720M results are presented as downstream and held-out FineWeb-HQ losses plus per-task zero-shot accuracy scores.The supplied passages identify these result tables and figures but do not state their comparative findings.Source: Table 8 · p. 43, Figure 12 · p. 43, Table 9 · p. 43

G.1 DETAILS ON THE DOWNSTREAM TASKS

• Downstream evaluation groups tasks by prompt format: QA-formatted tasks end with an answer slot, while completion-formatted tasks provide a textual stem continued by each option.ARC-Easy, ARC-Challenge, PIQA, SciQ, CommonSense QA, and Openbook QA are QA-formatted; HellaSwag, WinoGrande, and WSC are completion-formatted.Source: §G.1 · p. 44

G.2 EXTENDED RESULTS ON REPHRASING

• OptiSelect consistently outperforms standard training in average accuracy across rephrasing mixtures on in-domain proxy tasks and out-of-domain benchmarks.Table 10 reports zero-shot accuracy (%).Source: §G.2 · p. 44, Table 10 · p. 44
• Both standard training and OptiSelect achieve their best performance across all nine tasks at a 50% rephrasing rate.Source: §G.2 · p. 44
• The 50% rephrasing mixture combines original FineWeb and FinePhrase documents at a 1:1 ratio.Source: §G.2 · p. 44