Large Diffusion Models Need More Than EMA

Chuyu Hu; Huanran Chen; Shitong Shao; Yinpeng Dong

cs.LGOct 16, 2026

Large Diffusion Models Need More Than EMA

Authors: Chuyu Hu, Huanran Chen, Shitong Shao, Yinpeng Dong

Organizations: Tsinghua University · Shanghai Qi Zhi Institute · The Hong Kong University of Science and Technology

Abstract

Large-scale diffusion models are commonly trained with a constant learning rate (LR) and exponential moving average (EMA), despite the strong performance of LR decay and Schedule-Free optimization in training large language models. We study whether EMA is uniquely suited to diffusion training, or whether stronger alternatives have been underexplored. Our theory shows that averaging a constant-step trajectory has an inherent optimization limitation, whereas LR decay and Schedule-Free optimization admit vanishing stationarity guarantees under standard nonconvex assumptions. Empirically, LR decay and EMA provide complementary benefits: combining decay with EMA further improves the standard baseline at most tested starting points, although the gain remains sensitive to when decay begins. Schedule-Free avoids this timing requirement and consistently outperforms EMA with a constant LR across training stages, achieving both lower fixed-evaluation objectives and better generation quality. Our analysis further shows that all three strategies substantially reduce parameter variability, yet stronger stabilization is not always better: EMA with a constant LR retains comparatively higher parameter variability, while aggressive decay combined with EMA achieves the lowest checkpoint variability but also weakens continued optimization. Schedule-Free substantially reduces parameter variability while preserving more late-stage progress, providing a better balance between parameter stability and continued optimization. Finally, the benefits of Schedule-Free extend beyond pretraining to ReFL and generator-side DMD, demonstrating its transferability across distinct diffusion optimization settings. Overall, our results suggest that large diffusion models need more than parameter averaging alone: effective optimization requires balancing variance reduction with continued optimization.

Figures & tables

Figure 1: Late-stage cosine decay is beneficial but timing-sensitive. (a) Five independent 40k-step cosine-decay branches from the same Constant LR + EMA trajectory. (b) Three-seed mean FID-5K for the matched baseline, decayed online model, and Decay + EMA . EMA further improves all five decay endpoints, while Decay + EMA outperforms the matched baseline at four of five starting points.

Figure 2: Generation quality and fixed-evaluation objectives for Constant LR + EMA and Schedule-Free AdamW . (a) Three-seed mean FID-5K: Schedule-Free is better at every checkpoint from 320k to 400k. (b–d) Fixed-evaluation denoising MSE, projection loss, and total loss: denoising MSE is similar, while Schedule-Free has lower projection and total losses.

Figure 3: Late-stage checkpoint variability versus FID-5K at 400k. Variability is the trace variance between consecutive checkpoints (

g=1

). Decay + EMA is most stable, whereas Schedule-Free AdamW achieves the best FID-5K.

Figure 4: Late-stage parameter movement and optimization progress. (a) Net

\ell_{2}

displacement over five adjacent 1k-step intervals. (b) Fixed-evaluation loss decrease relative to 400k. Decay + EMA minimizes displacement but does not improve the objective, whereas Schedule-Free AdamW retains continued progress. The comparison uses EMA sequences for constant LR and decay and the averaged

x

sequence for Schedule-Free; full results are in Appendix A.3 .

Figure 5: DMD optimizer ablation comparing all-AdamW with Schedule-Free applied only to the generator (Gen) or guidance model (Gui). (a) FID-5K: generator-side Schedule-Free improves all checkpoints, whereas guidance-side Schedule-Free does not. (b) Generator DMD objective

\mathcal{L}_{\mathrm{DMD}}

. (c) Fake-score diffusion loss

\mathcal{L}_{\mathrm{diff}}

; definitions are in Appendix A.6 .

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Optimizer	Base LR	Warmup	LR schedule	EMA
Constant LR + EMA	AdamW	1e-4	None	Constant	0.9999
Decay + EMA	AdamW	1e-4	None	cosine to 1e-6	0.9999
Schedule-Free	SF-AdamW	4e-4	5k linear	Schedule-Free	None

Appendix

Table 1: Main pretraining optimizer configurations.

Figure 6: Learning-rate schedules used in the main experiment. (a) Constant LR + EMA and independent 40k-step cosine-decay branches from different late-stage checkpoints of the same Constant-LR AdamW trajectory. Each branch decays the learning rate from $10^{-4}$ to $10^{-6}$ ; different dashed curves indicate different decay starting points. (b) Constant-LR AdamW uses $10^{-4}$ throughout, whereas Schedule-Free AdamW uses a base learning rate of $4\times 10^{-4}$ with a 5k-step linear warmup.

Method / sequence	Checkpoint variability (trace variance, $g=1$ )
Constant-LR online	7.1611709595e+00 (100%)
Schedule-Free averaged $x$	1.4373749495e-03 (0.0201%)
Constant-LR EMA	5.5516868830e-03 (0.0775%)
Cosine-decay model	7.3927e-04 (0.0103%)
Cosine-decay EMA	1.1807e-04 (0.0016%)

Appendix

Table 2: Checkpoint variability over 20 consecutive saved checkpoints (checkpoint gap

g=1

), measured by total trace variance. Parentheses give the ratio to the Constant-LR online trajectory. Both model and EMA sequences are reported for cosine decay.

Run	Weights	400–401k	401–402k	402–403k	403–404k	404–405k	Mean
AdamW baseline	model	69.953	68.033	68.007	68.192	68.026	68.442
AdamW baseline	EMA	12.636	12.700	12.683	12.676	12.669	12.673
Schedule-Free	$x$	6.5643	6.5572	6.5492	6.5408	6.5326	6.5488
40k cosine decay	model	0.909307	0.843436	0.831407	0.827545	0.820779	0.846495
40k cosine decay	EMA	1.794585	1.624687	1.471350	1.333021	1.208285	1.486386

Appendix

Table 3: Net parameter displacement over five adjacent 1k-step late-stage windows. Each entry is the

\ell_{2}

distance

\|\theta_{t+1000}-\theta_{t}\|_{2}

, not its square. For the Constant-LR baseline and cosine decay, we report both model and EMA weights; Schedule-Free uses its evaluation-time averaged

x

sequence.

Method	400k	401k	402k	403k	404k	405k
Constant LR + EMA	.43738	.43736	.43732	.43729	.43726	.43723
Schedule-Free	.43461	.43460	.43457	.43454	.43451	.43447
40k Cosine Decay model	.43696	.43735	.43745	.43754	.43752	.43760
40k Cosine Decay + EMA	.43691	.43691	.43693	.43695	.43697	.43699

Appendix

Table 4: Fixed-evaluation total loss from 400k to 405k. All methods use identical fixed evaluation data, noise, and timestep sampling. Constant LR + EMA and Schedule-Free decrease throughout this interval, whereas neither the 40k Cosine Decay model nor its EMA improves in the post-decay regime.

Figure 7: Stability at an aggressive learning-rate scale. Standard AdamW with learning rate $4\times 10^{-4}$ and 5k-step linear warmup develops a loss spike and sustained NaNs after approximately 218.7k steps. Schedule-Free remains stable with the same rate and warmup; standard AdamW is stable at $10^{-4}$ . The inset enlarges the normal-loss range.

	Training reward	Orthogonal evaluation rewards
Run	HPSv3	Pick	CLIP	HPSv2.1	Aes	ImgR	HPSv3++	Uni IAA	Uni IQA	Uni ISTA
Z-Image base (no RL)	6.2544	22.8690	0.2754	0.2962	5.4118	0.9947	12.0641	60.9511	75.0656	37.7925
AdamW $1\times$ , wu0	6.3500	22.8766	0.2754	0.2971	5.4177	0.9959	12.1271	60.7050	74.8078	38.3209
SF-AdamW $4\times$ , wu5	6.8626	22.8921	0.2748	0.2995	5.4456	1.0150	12.3485	60.6995	75.2258	38.3189
SF-AdamW $20\times$ , wu5	9.6134	22.8341	0.2807	0.3123	5.7066	1.1197	12.0347	51.5663	67.4061	41.7878

Appendix

Table 5: ReFL results with HPSv3 as the training reward. HPSv3 is the directly optimized on-reward metric; all remaining columns are orthogonal reward models used only for cross-reward evaluation and provide no training signal. Fine-tuned runs are evaluated at checkpoint 100, and the Z-Image base model without reinforcement learning is included as a reference. “wu0” and “wu5” denote zero and five warmup steps, respectively. Bold indicates the best value in each column among the displayed rows.

Method	Ckpt	FID-5K $\downarrow$	sFID $\downarrow$	IS $\uparrow$	Precision $\uparrow$	Recall $\uparrow$
AdamW baseline	1k	10.126	35.310	248.559	0.7886	0.7395
Gen SF-AdamW + Gui AdamW	1k	9.988	35.733	236.856	0.7850	0.7401
Gen AdamW + Gui SF-AdamW	1k	10.199	35.584	256.577	0.7970	0.7210
AdamW baseline	2k	9.757	34.865	242.844	0.7814	0.7497
Gen SF-AdamW + Gui AdamW	2k	9.573	35.261	256.355	0.7950	0.7451
Gen AdamW + Gui SF-AdamW	2k	9.882	35.529	245.570	0.7848	0.7395

Appendix

Table 6: DMD generation quality under the all-AdamW baseline and the two component-wise Schedule-Free ablations. Every metric is evaluated on 5,000 generated images. Within each checkpoint, bold marks the best of the three configurations. Lower FID-5K and sFID and higher IS, precision, and recall are better.

Figure 8: Late-stage generation quality for Constant LR + EMA , Schedule-Free AdamW , and 40k Decay + EMA . (a) FID-5K, where lower is better. (b) Inception Score, where higher is better. The Constant LR + EMA and Schedule-Free AdamW curves follow continuous training trajectories. Green markers denote endpoints of independent 40k-step cosine-decay branches and should not be interpreted as a continuous trajectory. Schedule-Free AdamW outperforms Constant LR + EMA on both metrics at every checkpoint from 320k to 400k.

Figure 9: Fixed-evaluation objectives for Constant LR + EMA , Schedule-Free AdamW , and 40k Decay + EMA . (a) Denoising MSE is similar across the three strategies. (b) Schedule-Free AdamW maintains a distinctly lower projection loss. (c) Consequently, Schedule-Free AdamW also achieves the lowest total fixed-evaluation loss throughout the evaluated interval. The cosine-decay values are endpoints of independent 40k-step branches; dashed segments connect them only as a visual guide.

Method	Checkpoint / branch	FID-5K	IS	sFID
Schedule-Free AdamW	320k	9.3722	232.5443	35.1116
	340k	9.2903	236.6247	35.1530
	360k	9.1916	238.2731	35.1622
	380k	9.1237	240.8229	35.1441
	400k	9.0504	245.7993	35.1334
Constant LR + EMA	320k	9.4972	227.6105	35.0881

Appendix

Table 7: Complete generation-quality results for the main experiment, including Constant LR AdamW without EMA and 40k-step cosine-decay AdamW evaluated with and without EMA. The first three methods follow continuous trajectories; every cosine-decay result is the endpoint of an independent branch from a different late-stage checkpoint. FID-5K entries are means over three sampling seeds. Lower FID-5K and sFID and higher Inception Score are better.

Timestep	Constant + EMA	Schedule-Free	Relative improvement
random $t$	0.766211	0.765308	0.12%
$t=0.02$	1.016053	1.013539	0.25%
$t=0.10$	0.949475	0.947750	0.18%
$t=0.20$	0.868994	0.867809	0.14%
$t=0.40$	0.796613	0.795586	0.13%
$t=0.60$	0.709609	0.708319	0.18%

Appendix

Table 8: Timestep-wise denoising MSE. Relative improvement is for Schedule-Free versus Constant + EMA.

Figure 10: Selected class-conditional ImageNet-1K samples from the REPA-trained SiT-XL/2 model optimized with Schedule-Free AdamW at $4\times$ the baseline learning rate after 800k training steps. Images are generated at $256\times 256$ resolution.

Figure 11: Additional fixed-prompt ReFL qualitative comparison (prompt set 1). The rows compare the base model, AdamW, and Schedule-Free AdamW at two learning-rate scales; the complete prompts are printed above the columns.

Figure 12: Additional fixed-prompt ReFL qualitative comparison (prompt set 2), using the same row ordering as Figure 11 .

Figure 13: Additional fixed-prompt ReFL qualitative comparison (prompt set 3), using the same row ordering as Figure 11 .

Figure 14: Additional fixed-prompt ReFL qualitative comparison (prompt set 4), using the same row ordering as Figure 11 .

Figure 15: Selected class-conditional ImageNet samples from the best-performing DMD optimizer variant after 3k generator steps. Schedule-Free AdamW is applied to the generator, while the guidance model retains AdamW. The figure shows an $8\times 8$ grid of samples at $256\times 256$ resolution.

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 205 papers

ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models
May 18, 2026Aaron DefazioLanguage Model Pretraining Learning Rate Scheduling
FAIR at Meta Super-Intelligence Labs
Just on Time: Token-Level Early Stopping for Diffusion Language Models
Feb 11, 2026Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi +7Diffusion Language Model Inference Diffusion Language Model Decoding
SoftServe Inc.
LESS Is More: Mutual-Stability Sampling for Diffusion Language Models
Jun 15, 2026Amr Mohamed, Guokan Shang, Michalis VazirgiannisMasked Diffusion Models Diffusion Language Model Decoding
1MBZUAI · 2Ecole Polytechnique

May 18, 2026J/K move · Enter open · S save

ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models

Aaron Defazio

FAIR at Meta Super-Intelligence Labs

Language Model Pretraining Learning Rate Scheduling LLM Training

Open paper AskPDF arXiv

Schedule-Free Learning has shown promise as a practical anytime training method for machine learning, showing success across dozens of standard benchmark problems. However, strong performance for LLM training has only been demonstrated at small scales. We identify a number of fixes necessary to scale up Schedule-Free Learning to larger batch sizes and model sizes, and present a learning-rate-free and schedule-free method (ScheduleFree+) for training large language models which greatly outperforms Warmup-Stable-Decay (WSD) schedules. We also demonstrate that Schedule-Free Learning is most effective for long duration training, and at 1000 tokens per parameter, it outperforms SOTA schedules by 31%. Schedule-Free Learning provides a theoretical foundation for the use of model averaging and checkpoint merging during pretraining.

Abstract

Deep Learning Optimization

Explore similar work

ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models

Just on Time: Token-Level Early Stopping for Diffusion Language Models

LESS Is More: Mutual-Stability Sampling for Diffusion Language Models