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

Mar 8, 2026

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

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.

Fangxin Wang, Peyman Baghershahi, Langzhou He +3

Department of Computer Science, University of Illinois Chicago

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

Ask PDF

Jun 17, 2026

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

As Large Language Model (LLM) datasets scale to trillions of tokens, data selection has emerged as a critical frontier to filter out uninformative noise and construct adaptive learning trajectories. Beyond static heuristic filtering, advanced data selection methods for LLM training largely follow two paradigms, each with fundamental limitations. Influence-based methods provide principled bi-level objectives but require intractable inverse-Hessian computations, while excess-loss methods are computationally efficient but rely on a static reference model that becomes misaligned with the evolving proxy model during training. We propose BLADE (Bi-Level Adaptive Data sElection), a Hessian-free framework for data selection. BLADE reformulates the bi-level optimization problem underlying influence-based methods as a penalized single-level objective via Lagrange multipliers, avoiding inverse-Hessian computation while revealing a principled connection to excess-loss based data selection. The resulting objective recovers an excess-loss form but replaces the static reference model with a dynamic one that stays synchronized with training. Theoretically, we prove that this penalized formulation guarantees first-order convergence. For efficient online batch selection, we instantiate BLADE as a memoryless randomized block-coordinate Frank-Wolfe algorithm. Extensive experiments show that BLADE consistently outperforms state-of-the-art data selection baselines, providing a practical recipe for LLM training.

Jiaxing Wang, Deping Xiang, Jin Xu +9

University of Oxford

Training Data Selection Bilevel Optimization LLM Training Frank-Wolfe Algorithms

Ask PDF

May 29, 2026

Towards Efficient LLMs Annealing with Principled Sample Selection

The annealing phase is a pivotal convergence stage in LLM pre-training that ultimately determines final model quality. However, effectively selecting training data during this phase remains a key challenge. Current strategies rely on empirical heuristics, such as domain filtering or context extension, which lack a principled grounding in optimization theory. In this work, we characterize the annealing phase through the lens of the loss landscape's spectral geometry. We argue that optimal convergence requires gradient updates to satisfy heterogeneous constraints across different eigen-directions. Building on this insight, we formulate data selection as a problem of satisfying these directional constraints. To this end, we propose DiReCT (Directionally-Restrained Constrained Training), a novel framework that reformulates sample selection in the annealing stage as a constrained optimization problem. By imposing explicit directional constraints on per-sample gradients based on the spectral properties of the Hessian, DiReCT identifies samples that align with the optimal curvature-aware descent path. Extensive experiments across various model scales demonstrate that DiReCT consistently achieves state-of-the-art performance. For future research, code is available at https://github.com/xuyj233/Direct.

Yuanjian Xu, Jianing Hao, Wanbo Zhang +2

The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Microsoft Research Asia, Beijing, China · Fudan University, Shanghai, China

Language Model Pretraining Efficient Language Model Training Curvature-Aware Optimization Training Data Selection

Ask PDF

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