Optimizers for Diffusion Models: A Controlled Benchmark

Arman Bolatov; Egor Shulgin; David Li; Abduragim Shtanchaev; Sebastian U. Stich; Maxim Panov; Eric Moulines; Peter Richtárik

cs.LGSep 19, 2026

Optimizers for Diffusion Models: A Controlled Benchmark

Authors: Arman Bolatov, Egor Shulgin, David Li, Abduragim Shtanchaev, Sebastian U. Stich, Maxim Panov, Eric Moulines, Peter Richtárik,

Organizations: CISPA Helmholtz Center for Information Security · KAUST · MBZUAI · EPITA

Abstract

Discrete diffusion models now match autoregressive language models on several benchmarks, while the question of how best to train them has received far less attention: the optimizer is inherited from one paper to the next and never compared. New optimizers, meanwhile, are validated almost exclusively on autoregressive pretraining, a different objective on a different loss surface. We present a controlled optimizer benchmark across four diffusion formulations, to our knowledge the first for discrete diffusion: seven optimizers (AdamW, Lion, Muon, SOAP, MARS, MARS-M, Schedule-Free) on masked diffusion (text8), uniform diffusion (QM9, and LM1B through the Gaussian duality) and Gaussian diffusion on images (CelebA-64), each on a task with published reference values. Every optimizer receives the same search protocol, and every winner is retrained at the full budget with three seeds. AdamW is a strong default but not always the right choice: it is beaten by a resolved margin on two of the four tasks, and the winner changes with the formulation, so the optimizer deserves the same care as the rest of the training recipe. Notably, methods validated on autoregressive language model pretraining transfer well: Muon, MARS-M and SOAP each beat the tuned AdamW on at least one diffusion formulation. The benchmark, all runs and every figure are reproducible end to end from the released code at https://github.com/armanbolatov/diffusion-baselines.

Figures & tables

	text8 (MDLM)	QM9 (UDLM)	LM1B (DUO)	CelebA-64 (DDPM)
	bpc $\downarrow$	PPL $\downarrow$	PPL $\downarrow$	FID $\downarrow$
MARS-M	1.8072 $\pm$ 0.0089	2.076 $\pm$ 0.014 ↓	62.2 $\pm$ 0.4 ↑	6.32 $\pm$ 0.62
SOAP	1.8073 $\pm$ 0.0100	2.066 $\pm$ 0.013 ↑	69.5 $\pm$ 1.5 ↓	10.52 $\pm$ 1.41 ↓
Muon	1.8126 $\pm$ 0.0040	2.069 $\pm$ 0.014 ↑	61.6 $\pm$ 0.5 ↑	8.58 $\pm$ 0.99 ↓
AdamW	1.8154 $\pm$ 0.0053	2.072 $\pm$ 0.013	68.2 $\pm$ 0.8	6.88 $\pm$ 0.71
Lion	1.8509 $\pm$ 0.0064 ↓	2.068 $\pm$ 0.014 ↑	diverges	11.52 $\pm$ 3.37 ↓

Table 1: Final metric after full-budget training, mean and standard deviation over three seeds. Lower is better; bold is the best per task. Cells are shaded on a diverging scale centered on AdamW, green better and red worse, the tint deepening with the size of the gap relative to each task’s spread; arrows mark the differences that the paired within-seed test resolves (beyond twice the standard error). Lion’s LM1B entry diverges on two of three seeds (Section 5.3 ).

Figure 1: text8 validation bits per character on a log scale, mean over three seeds with a

\pm 1

standard deviation band, against training steps (left) and cumulative training compute (right). The FLOP counts differ by at most 7% across optimizers, so the two panels nearly coincide; per-step wall-clock does not follow FLOPs (Table 6 ).

Figure 2: QM9 validation perplexity on a log scale, mean over three seeds with a

\pm 1

standard deviation band, against training steps (left) and cumulative training compute (right). The FLOP counts differ by at most 13% across optimizers; per-step wall-clock does not follow FLOPs (Table 6 , Appendix I ).

Figure 3: LM1B validation perplexity on a log scale, mean over three seeds with a

\pm 1

standard deviation band, against training steps (left) and cumulative training compute (right). Lion is omitted: it diverges after DUO’s curriculum switch at step 15k on two of three seeds, ending near

10^{4}

perplexity, and plotting it would flatten every other curve.

Figure 4: CelebA-64 monitor FID (left) and held-out denoising MSE (right) during training, mean over three seeds. Both are measured on the same runs: FID separates the optimizers while the MSE curves collapse onto each other. Lion’s FID is not monotone between 50k and 65k steps.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

task	params	shape	batch	ctx	steps	tune steps
text8 (MDLM)	16.6M	8L/256d/8h	512	256	30000	6000
QM9 (UDLM)	90.2M	12L/768d/12h	2048	32	25000	5000
LM1B (DUO)	137.1M	12L/768d/12h	512	128	30000	6000
CelebA-64 (DDPM)	78.7M	UNet, ch 128	128	64px	100000	20000

Appendix

Table 2: The four training tasks. Steps is the full benchmark budget, tune steps is the reduced budget used during hyperparameter search.

optimizer	axis	values
AdamW	beta1	$\{$ $0.5$ , $0.8$ , $0.9$ , $0.95$ , $0.99$ , $0.999$ $\}$
	beta2	$\{$ $0.5$ , $0.8$ , $0.9$ , $0.95$ , $0.99$ , $0.999$ $\}$
Lion	beta1	$\{$ $0.5$ , $0.8$ , $0.9$ , $0.95$ , $0.99$ , $0.999$ $\}$
	beta2	$\{$ $0.5$ , $0.8$ , $0.9$ , $0.95$ , $0.99$ , $0.999$ $\}$
Muon	adamw_lr	$\{$ $3\!\cdot\!10^{-6}$ , $1\!\cdot\!10^{-5}$ , $3\!\cdot\!10^{-5}$ , $1\!\cdot\!10^{-4}$ , $3\!\cdot\!10^{-4}$ , $1\!\cdot\!10^{-3}$ , $3\!\cdot\!10^{-3}$ , $0.01$ , $0.03$ $\}$
	momentum	$\{$ $0$ , $0.2$ , $0.5$ , $0.8$ , $0.9$ , $0.95$ , $0.99$ , $0.999$ $\}$

Appendix

Table 3: Optimizer-specific search axes, the union of the value sets searched across studies; a few studies searched narrower sets on some axes (MARS’s smallest gamma on text8 only, narrower adamw_lr and momentum in the early CelebA-64 studies), and the edge selections in the text are relative to the set each study searched. Every optimizer additionally sweeps the three shared axes listed in the text.

task	optimizer	axis	searched	common
text8	Schedule-Free	weight_decay	$\{0,\ldots,10\}$ (7)	(6)
text8	MARS	lr	$\{0.0001,\ldots,0.01\}$ (5)	(8)
qm9	MARS	lr	$\{0.0001,\ldots,0.01\}$ (5)	(8)
celeba64	Muon	lr	$\{1e-05,\ldots,0.01\}$ (7)	(8)
celeba64	MARS	lr	$\{0.0001,\ldots,0.01\}$ (5)	(8)
celeba64	MARS-M	lr	$\{3e-05,\ldots,0.003\}$ (5)	(8)

Appendix

Table 4: Every study whose shared axes differed from the common grid, generated from the values each study actually searched.

task	optimizer	lr	wd	clip	other
text8 (MDLM)	AdamW	0.001	0.3	0	beta1=0.8 , beta2=0.8
	Lion	0.0003	3	0	beta1=0.8 , beta2=0.9
	Muon	0.01	0.1	1	adamw_lr=0.001 , momentum=0.5 , nesterov=False
	Schedule-Free	0.01	1	0	beta1=0.95 , beta2=0.8
	SOAP	0.003	0.3	0	precondition_frequency=10 , beta1=0.95 , beta2=0.95
	MARS	0.001	1	0	gamma=0.003 , beta1=0.9 , beta2=0.9

Appendix

Table 5: Selected hyperparameters. These are the configurations retrained at the full budget with three seeds.

	text8 (MDLM)		QM9 (UDLM)		CelebA-64 (DDPM)
	time	mem	time	mem	time	mem
AdamW	1.00	1.00	1.00	1.00	1.00	1.00
Lion	1.00	0.99	0.99	0.99	1.05	0.98
Muon	1.14	0.99	1.08	0.99	2.04	0.98
Schedule-Free	1.00	1.00	1.01	1.00	1.01	1.00
SOAP	1.31	1.03	1.52	1.10	2.90	1.02

Appendix

Table 6: Median time per training step and peak memory, both relative to AdamW on the same task. Values above 1 are more expensive; the tint grows with the overhead.

Figure 5: CelebA-64 faces, DDIM-100 as in the headline FID protocol, from the same four initial noise draws for every optimizer, one column per optimizer. DDIM is deterministic, so each column is the image that optimizer’s model assigns to the shared latents; the near-identical faces across optimizers show that the models learned nearly the same latent-to-image map.

optimizer	sample (first 64 of 256 characters)
AdamW	x iat sedokungu wolozzuy portiuoudeyoimoo qamigu kamea nlako xan
Lion	x ikprszdokungu woeozzly lortipoudeymimooljtmigubkamea nlako xao
Muon	ddicture only a low low short of a common pamigabiam a n and x o
Schedule-Free	x ist sedobunin woeozzly corti ouceymimoo jamigu kamea stano xao
SOAP	xxiatriedokunin colozzen port piusecormoo quilgaliam a alkelax o
MARS	x iap szdokungu joeozzly korti oudeymioxolqtmigubkamea nlako xao

Appendix

Table 7: Unconditional text8 samples from each optimizer’s model, 256 sampling steps.

optimizer	samples (SMILES)		valid
AdamW	C#CC(=O)C1NC1C#N	CC1CCC(=O)C(C)O1	100%
Lion	C#CC(CO)C(=O)C#N	CC1CCC(=O)C(C)O1	100%
Muon	C#CC(CO)C(=O)C#N	CC1CCC(=O)C(C)O1	99%
Schedule-Free	C#CC(=O)C1NC1C#N	CC1CCC(=O)C(C)O1	98%
SOAP	C#CC(=O)C1NC1C#N	CC1CCC(=O)C(C)O1	99%
MARS	C#CC(=O)C1NC1C#N	CC1CCC(=O)C(C)O1	99%

Appendix

Table 8: Unconditional QM9 molecules from each optimizer’s model, with the fraction of its 1024 samples (as in Table 9 ) that parse as valid SMILES.

optimizer	valid $\uparrow$	unique $\uparrow$	novel $\uparrow$	FCD $\downarrow$
AdamW	99.6%	99.6%	8.6%	1.60
Lion	99.6%	99.5%	7.5%	1.69
Muon	99.3%	99.5%	8.8%	1.69
Schedule-Free	98.1%	99.6%	12.9%	1.47
SOAP	99.1%	99.3%	9.1%	1.66
MARS	99.3%	99.9%	7.0%	1.72

Appendix

Table 9: QM9 sample quality over 1024 molecules per optimizer (one seed’s weights, matching the evaluation protocol, one fixed sampling seed), discussed in the QM9 results. Valid is the fraction that parses under RDKit [ RDKit Project, 2026 ] , unique the distinct fraction among the valid, novel the fraction of valid molecules absent from the training set, and FCD the Frechet ChemNet Distance [ Preuer et al., 2018 ] to 10k training-set molecules.

optimizer	sample (first 64 characters, median-length sentence)
AdamW	you shouldn ’ t eat water or bangladeshi bins if they want to ge
Lion	film the financialonero war or. corporations proposition she the
Muon	you don ’ t get water or fast, or if they want to study it and e
Schedule-Free	bharatiya pcs mohammed toro brothers or wedding corporations pro
SOAP	sbelle ’ blows shown brothers or wedding corporations as she pic
MARS	you shouldn ’ t imagine naming or amiss proposition if they want

Appendix

Table 10: LM1B samples from each optimizer’s model (raw final weights of one seed), length-conditioned to the median validation sentence length and decoded with 1000 ancestral steps under the uniform-state reverse process. The Lion model sampled here is one of the two seeds that diverge after the curriculum switch, so its row samples the failed model.

Preference Optimization for Diffusion Models

See all 17 papers

Scaling Beyond Masked Diffusion Language Models
Feb 16, 2026Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang +4Diffusion Language Models Diffusion Models
NVIDIA, Santa Clara, USA · Department of Computer Science, Cornell Tech, NYC, USA · Department of Computer Science, Cornell University, Ithaca, USA +2
The Diffusion Duality, Chapter II: $Ψ$ -Samplers
Feb 24, 2026Justin Deschenaux, Caglar Gulcehre, Subham Sekhar SahooDiffusion Models Diffusion Sampling
EPFL, Lausanne, Switzerland · Microsoft AI · Cornell Tech, NY
Enhancing Diffusion Language Models with Autoregressive Post-Training Weights
Oct 6, 2026Yiming Qin, Ke Wang, Amel Abdelraheem +2Diffusion Language Model Training Model Merging
EPFL, Lausanne, Switzerland

Feb 16, 2026J/K move · Enter open · S save

Scaling Beyond Masked Diffusion Language Models

Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang, Justin Deschenaux+3

NVIDIA, Santa Clara, USA · Department of Computer Science, Cornell Tech, NYC, USA · Department of Computer Science, Cornell University, Ithaca, USA+2

Diffusion Language Models Diffusion Models Diffusion Language Model Training Diffusion Language Model Benchmarking Masked Diffusion Models Diffusion Language Model Inference

Open paper AskPDF arXiv

Diffusion language models are a promising alternative to autoregressive models due to their potential for faster generation. Among discrete diffusion approaches, Masked diffusion currently dominates, largely driven by strong perplexity on language modeling benchmarks. In this work, we present the first scaling law study of uniform-state and interpolating discrete diffusion methods. We also show that Masked diffusion models can be made approximately 12% more FLOPs-efficient when trained with a simple cross-entropy objective. We find that perplexity is informative within a diffusion family but can be misleading across families, where models with worse likelihood scaling may be preferable due to faster and more practical sampling, as reflected by the speed-quality Pareto frontier. These results challenge the view that Masked diffusion is categorically the future of diffusion language modeling and that perplexity alone suffices for cross-algorithm comparison. Scaling all methods to 1.7B parameters, we show that uniform-state diffusion remains competitive on likelihood-based benchmarks and outperforms autoregressive and Masked diffusion models on GSM8K, despite worse validation perplexity. We provide the code, model checkpoints, and video tutorials on the project page: http://s-sahoo.github.io/scaling-dllms

Abstract

Preference Optimization for Diffusion Models

Explore similar work

Scaling Beyond Masked Diffusion Language Models

The Diffusion Duality, Chapter II: ΨΨΨ-Samplers

Enhancing Diffusion Language Models with Autoregressive Post-Training Weights

The Diffusion Duality, Chapter II: $Ψ$ -Samplers