OptiSelect: How does the Optimizer Shape Data Curriculum?

Simin Fan; Alireza Abdollahpoorrostam; Martin Jaggi

cs.LGOct 2, 2026

OptiSelect: How does the Optimizer Shape Data Curriculum?

Authors: Simin Fan, Alireza Abdollahpoorrostam, Martin Jaggi

Organizations: EPFL

Abstract

Online data selection has demonstrated substantial efficiency gains for LLM pretraining by training on the most valuable candidates within each batch. Since a candidate's value is realized through its effective model update, principled selection should account for the optimizer step, which reshapes the raw gradient before it updates model parameters. We formalize this optimizer-aware selection paradigm as OptiSelect and present the first systematic study of how the optimizer shapes data selection. Our theory establishes a selection gain principle in which the advantage of online selection is governed by the discriminability of the optimizer-induced utility scores. We prove that sign-based and polar-tangential preconditioners of Lion and Muon would suffer from a discriminability collapse which caps attainable gains from OptiSelect, whereas diagonal-adaptive optimizers such as AdamW and Sophia admit strictly better upper bounds. The proposed principle also yields a quantitative derivation of the optimal candidate oversampling ratio. Pretraining experiments on 124M and 720M models are consistent with our theoretical analysis and show that AdamW's diagonal-adaptive scoring geometry remains the strongest scoring geometry even with Muon as optimizer. We further demonstrate that OptiSelect retains its benefits under data rephrasing, a technique used in modern data processing pipelines. Our findings provide theoretical foundations and practical guidance for co-designing optimizers and data selection in LLM pretraining.

Figures & tables

Figure 1: Selection gain ( Theorem 1 ) measured during 124M LLM pretraining. (a) Running medians at tuned

\eta^{\star}

: the measured gain mostly sits close to

c_{q}

for every optimizer; while Lion decays after proxy refresh. (b) Measured gain from learning rate sweep : the selection gain is flat in various learning rate

\eta

Panel of Figure 2 — Figure 2: Downstream loss on 124M and 720M models across optimizers. Models are trained on $4.2$ B/ $14.03$ B tokens under standard and OptiSelect paradigm, respectively. The intersection of two dotted curves denotes the processed-matched baseline with equal processed tokens under $\lambda=2$ .

Figure 3: Downstream Accuracy at 720M scale. OptiSelect outperforms standard training across all optimizers on in-domain and OoD benchmarks.

Figure 4: Scoring geometry ablations under AdamW and D-Muon optimizers . The matched geometry works best under AdamW , while both AdamW and SGD alternatives beat the matched scoring geometry on D-Muon .

Figure 5: Downstream loss curves across rephrased mixtures , where the benefits from selection are consistent.

Figure 6: Task accuracy on format rephrasing

\times

selection at 124M. (a) Accuracy against the rephrased mixture rates. (b) The selection gain at each mixture.

Figure 7: The candidate-multiplier sweep at 124M . (a) Downstream loss curves. (b) Selection gain over the standard . Open circles mark the theoretical values of

c_{q}\sigma_{U}

from Theorem 1 ; the dark-red curve draws the gain per streamed token

\Delta/\lambda

on the right axis, which peaks at

\lambda^{\star}=2

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Optimizer	Operator family	Linear in $\mathbf{x}$ ?	Magnitude information
Operators analyzed
SGD	Unpreconditioned (identity)	✓	✓
AdamW	Diagonal adaptive (second moment)	✓	✓
AdEMAMix	Diagonal adaptive (second moment)	✓	✓
Sophia	Diagonal adaptive (curvature)	✓ §	✓
Muon	Matrix sign / polar tangent	✓ ‡	Tangential; radial discarded

Appendix

Table 2: Properties of the frozen-state maps discussed in this appendix. Linearity guarantees additivity of the represented batch direction, not exact attribution under the full optimizer update. “Magnitude information” refers to retaining continuous amplitude information rather than only sign patterns; it does not imply norm preservation. § The Sophia row describes its unclipped scoring surrogate. ‡ The Muon row describes the ideal polar derivative; the implemented right-preconditioned proxy is also linear but retains radial components. ∗ The Lion / Signum rows describe the exact sign maps; their implemented fixed-sign surrogates are linear and retain candidate amplitude on coordinates with nonzero frozen signs.

Scale	$d_{\mathrm{model}}$	layers	heads	seq. len.	tokens/step
124M	768	12	12	512	32,768
720M	2048	12	16	512	131,072

Appendix

Table 3: Model architecture configurations.

Optimizer	Mode	val / dwn at $\eta=\{0.25,\,0.5,\,1,\,2\}\times$ default
AdamW	std	3.692 / 3.963	3.610 / 3.865	3.579 / 3.840	3.642 / 3.902
AdamW	sel	3.733 / 3.850	3.650 / 3.771	3.616 / 3.745	3.949 / 4.080
AdEMAMix	std	3.680 / 3.950	3.617 / 3.885	3.599 / 3.863	4.206 / 4.446
AdEMAMix	sel	3.724 / 3.855	3.659 / 3.780	3.635 / 3.763	4.484 / 4.646
Sophia	std	3.796 / 4.062	3.687 / 3.947	3.647 / 3.908	3.730 / 3.979
Sophia	sel	3.828 / 3.950	3.724 / 3.856	3.679 / 3.804	3.737 / 3.891

Appendix

Table 4: Downstream loss (“dwn”) and FineWeb Validation loss (“val”) across a sweep of learning rates at 124M. Bold marks the optimal learning rate on the downstream loss , which are taken for experiments.

Figure 8: Downstream loss in 124M learning-rate sweep. Stars

\star

mark the optimum for each optimizer.

Figure 9: Downstream loss in 720M learning-rate search. Stars

\star

mark the optimum for each optimizer.

Figure 10: Ablation on Proxy-refresh Frequency. AdamW at

\eta=10^{-3}

10

k steps, varying only the refresh interval

\Delta

; stars mark the best cell. (a) The downstream loss degrades monotonically proxy-fresh frequency; while (b) Held-out FineWeb-HQ loss moves in an opposite way. (c) Selection entropy falls as

\Delta

grows, i.e. a stale reference makes the scores look more decisive while aligning them with a stale target. (d) Step time is flat to within

0.03

s (

6\%

) : the refresh is amortized over

\Delta

steps and never dominates.

refresh $\Delta$ (steps)	$\mathbf{500}$	$\mathbf{1000}$	$5000$	$10000$
downstream loss $\downarrow$	3.7165	3.7178	3.7421	3.8244
FineWeb-HQ val. loss	3.6265	3.6265	3.6185	3.6052
selection entropy	0.84	0.84	0.81	0.79
$t_{\mathrm{step}}$ (s)	0.51	0.50	0.50	0.48

Appendix

Table 5: Proxy refresh sweep, AdamW at

\eta=10^{-3}

10

k steps at 124M.

t_{\mathrm{step}}

is measured as median wall-clock per optimizer step.

Figure 11: Downstream loss and Held-out Fineweb-HQ loss compared across all optimizers at 720M. (i) Standard training , (ii) OptiSelect .

	Held-out FineWeb-HQ loss $\downarrow$			Downstream loss $\downarrow$
Optimizer	std@ $128$ k	sel@ $128$ k	$\Delta$	std@ $128$ k	sel@ $128$ k	$\Delta$
AdamW	3.107	3.094	-0.013	3.361	3.241	-0.119
AdEMAMix	3.108	3.093	-0.014	3.410	3.274	-0.135
Sophia	3.198	3.179	-0.019	3.486	3.458	-0.028
Lion	3.133	3.116	-0.016	3.417	3.399	-0.018
D-Muon	3.204	3.195	-0.008	3.455	3.471	+0.015

Appendix

Table 6: Downstream loss and Held-out Fineweb-HQ loss on 124M models.

Figure 12: Downstream loss and Held-out Fineweb-HQ loss compared across all optimizers at 720M. (i) Standard training , (ii) OptiSelect .

	Held-out FineWeb-HQ			Downstream probe
Optimizer	std@ $107$ k	sel@ $107$ k	$\Delta$	std@ $107$ k	sel@ $107$ k	$\Delta$
AdamW	2.704	2.714	+0.010	2.995	2.891	-0.104
AdEMAMix	2.705	2.710	+0.005	3.004	2.908	-0.096
Sophia	2.745	2.736	-0.009	3.049	3.026	-0.023
Lion	2.729	2.731	+0.002	3.018	3.022	+0.004
D-Muon	2.696	2.684	-0.012	2.993	2.944	-0.049

Appendix

Table 8: Downstream loss and Held-out Fineweb-HQ loss on 720M models.

Training Data Selection

Training data selection studies how to choose, filter, weight, or schedule examples to improve learning while reducing data and compute costs.

See all 187 papers

Filter-then-Weight: Online Data Selection and Reweighting for LLM Fine-Tuning
Mar 8, 2026Fangxin Wang, Peyman Baghershahi, Langzhou He +3Data Selection for Language Model Fine-Tuning Training Data Selection
Department of Computer Science, University of Illinois Chicago
BLADE: Scalable Bi-level Adaptive Data Selection for LLM Training
Jun 17, 2026Jiaxing Wang, Deping Xiang, Jin Xu +9Training Data Selection Bilevel Optimization
University of Oxford
Towards Efficient LLMs Annealing with Principled Sample Selection
May 29, 2026Yuanjian Xu, Jianing Hao, Wanbo Zhang +2Language Model Pretraining Efficient Language Model Training
The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Microsoft Research Asia, Beijing, China · Fudan University, Shanghai, China

Mar 8, 2026J/K move · Enter open · S save

Filter-then-Weight: Online Data Selection and Reweighting for LLM Fine-Tuning

Fangxin Wang, Peyman Baghershahi, Langzhou He, Henry Peng Zou+2

Department of Computer Science, University of Illinois Chicago

Data Selection for Language Model Fine-Tuning Training Data Selection LLM Fine-Tuning

Open paper AskPDF arXiv

Gradient-based data selection offers a principled framework for estimating sample utility in large language model (LLM) fine-tuning, but existing methods are mostly designed for offline settings. They are therefore less suited to online fine-tuning, where data arrives sequentially, sample utility is step-dependent, and the effective update geometry is shaped by adaptive optimizers. We propose an optimizer-aware framework for gradient-based online data selection and reweighting in LLM fine-tuning. Our key idea is to view online selection not as static sample ranking, but as shaping the next target-oriented update under the current optimizer state. We formulate this as an optimizer-aware update-matching problem, establish its connection to second-order target utility, and show why subset-level construction must account for interactions and redundancy among selected samples. Based on this view, we develop a two-stage Filter-then-Weight algorithm that first filters geometrically useful candidates and then optimizes their coefficients. To make the framework practical for LLMs, we introduce a factorized outer-product gradient representation and optimized matrix computations for long-context data. Experiments show that our method consistently improves convergence and downstream performance over existing online data selection baselines under the same data budget.

Abstract

Training Data Selection

Explore similar work

Filter-then-Weight: Online Data Selection and Reweighting for LLM Fine-Tuning

BLADE: Scalable Bi-level Adaptive Data Selection for LLM Training

Towards Efficient LLMs Annealing with Principled Sample Selection