PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Diffusion Models Need More Than EMA

Chuyu Hu, Huanran Chen, Shitong Shao, Yinpeng Dong

arXiv:2610.22945v1cs.LG

TL;DR

The paper examines whether EMA is uniquely suited to large-scale diffusion training or whether LR decay and Schedule-Free are stronger alternatives. It analyzes their optimization trade-offs and compares them empirically, finding that Schedule-Free offers consistent gains without requiring a prespecified decay point.

  • Problem

    It is unclear whether constant LR with EMA has a genuine optimization advantage for large-scale diffusion or whether LR decay and Schedule-Free have been underexplored.

    Source: §1 · p. 2

  • Method

    The paper analyzes EMA, LR decay, and Schedule-Free through a shared variance-reduction perspective and compares them in large-scale diffusion pretraining.

    Source: §1 · pp. 1–2

  • Results

    Schedule-Free consistently outperforms CONSTANT LR + EMA across training stages without a prespecified decay point, while DECAY + EMA can improve the baseline depending on when decay begins.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Effective optimization must balance reducing trajectory variability with preserving continued progress; Schedule-Free provides that balance without choosing a decay point in advance.

    Source: §1 · p. 2, §4 · p. 9

    EMA averaging does not in general eliminate the residual optimization error of the underlying constant-step trajectory.

    Source: §2.2 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks whether constant-LR training with EMA is genuinely better suited to large-scale diffusion than LR decay or Schedule-Free optimization. It analyzes their stability–progress trade-offs and compares them in large-scale pretraining.

1 INTRODUCTION

  • Optimization recipes have changed less than diffusion architectures, despite LR decay being standard in large-scale language-model pretraining and Schedule-Free showing strong performance in that regime.The paper notes that improving the training recipe can reduce compute for a strong model without changing its architecture, objective, or data.Source: §1 · p. 1
  • The central question is whether EMA has a diffusion-specific optimization advantage or whether LR decay and Schedule-Free have been underexplored.Source: §1 · p. 2
  • EMA, LR decay, and Schedule-Free reduce persistent optimization variability differently, with distinct trade-offs between parameter stability and continued progress.EMA averages a constant-step trajectory, decay shrinks updates, and Schedule-Free averages an increasingly long trajectory while retaining large-step optimization.Source: §1 · p. 2
  • Averaging a constant-step trajectory has an optimization limitation, whereas LR decay and Schedule-Free admit vanishing stationarity guarantees under standard nonconvex assumptions.Source: §1 · p. 2
  • Schedule-Free consistently outperforms CONSTANT LR + EMA across training stages without a prespecified decay point, while DECAY + EMA can improve the standard recipe depending on when decay begins.DECAY + EMA achieves the lowest parameter variability, but weakens late-stage progress; Schedule-Free retains more progress while remaining stable.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries explain how constant-step noise, EMA, LR decay, and Schedule-Free affect optimization variability. Their central distinction is how each reduces variability while preserving—or suppressing—continued optimization.

2.1 WHY CONSTANT-STEP SGD HAS A NOISE FLOOR?

  • Under constant-step SGD, optimization error vanishes with training time but a noise floor remains, so the trajectory fluctuates around rather than settles at the optimum.Fresh noise is injected at every step even as previous error contracts.Source: §2.1 · pp. 2–3

2.2 EMA AVERAGES A NOISY TRAJECTORY

  • EMA averages past iterates to reduce fluctuations, but does not change the online dynamics or generally eliminate residual optimization error.Variance reduction does not imply correction of the underlying constant-step trajectory's optimization bias.Source: §2.2 · pp. 3–4

2.3 LEARNING-RATE DECAY SHRINKS THE NOISY TRAJECTORY

  • LR decay shrinks both updates and injected noise, making optimization error and parameter variance vanish asymptotically while reducing the size of further optimization steps.In practice, the paper studies finite late-stage cosine decays and their sensitivity to the decay starting point.Source: §2.3 · p. 4

2.4 SCHEDULE-FREE COMBINES TRAJECTORY AVERAGING WITH LARGE-STEP OPTIMIZATION

  • Schedule-Free keeps constant-scale optimization steps while averaging an increasingly long trajectory, stabilizing evaluation parameters and admitting a vanishing stationarity guarantee.Its variability reduction retains large-step optimization dynamics rather than suppressing them as aggressively as LR decay.Source: §2.4 · p. 5

3 VARIANCE REDUCTION IN DIFFUSION MODELS

The experiments compare constant LR + EMA, LR decay + EMA, and Schedule-Free across large-scale diffusion pretraining and other training regimes. Schedule-Free improves generation quality while balancing parameter stability with continued optimization; decay + EMA can improve the baseline, but its gains depend on timing.

3.1 EXPERIMENTAL SETUP: MODERN LARGE-SCALE DIFFUSION PRETRAINING

  • The controlled pretraining comparison uses approximately 675M-parameter SiT-XL/2 + REPA and evaluates parameter variability, optimization quality, and generation quality.The model, data, objective, batch size, and evaluation protocol are fixed across the main comparisons while the optimization recipe changes.Source: §3.1 · p. 5

3.2 LR DECAY CAN FURTHER IMPROVE CONSTANT LR + EMA, BUT THE GAIN IS TIMING-SENSITIVE

  • DECAY + EMA improves FID-5K at four of five tested starting points, but slightly hurts at one; EMA further improves all five decayed endpoints over their online models.At the 360k→400k branch, EMA changes FID-5K from 9.4895 without EMA to 9.1248 with EMA.Source: §3.2 · p. 6, Figure 1 · p. 6

3.3 SCHEDULE-FREE CONSISTENTLY IMPROVES OVER CONSTANT LR + EMA

  • Schedule-Free achieves lower FID-5K than CONSTANT LR + EMA at every checkpoint from 320k to 400k and a lower total fixed-evaluation loss throughout the evaluated interval.At 800k steps, it also improves FID-50K from 2.0770 to 1.8735, sFID from 4.539 to 4.479, and Inception Score from 265.6 to 272.7.Source: §3.3 · pp. 6–7, Figure 2 · p. 7

3.4 SCHEDULE-FREE OFFERS A BETTER TRADE-OFF THAN CONSTANT LR + EMA AND DECAY + EMA

  • At 400k, DECAY + EMA has the lowest checkpoint variability, but Schedule-Free has the best FID-5K and retains more late-stage optimization progress.FID-5K is 9.0504 for Schedule-Free, 9.1248 for DECAY + EMA, and 9.2139 for CONSTANT LR + EMA; after 400k, the post-decay EMA objective no longer improves.Source: §3.4 · pp. 7–8, Figure 3 · p. 7, Figure 4 · p. 8

3.5 BROADER VALIDATION ACROSS DIFFUSION TRAINING REGIMES

  • Schedule-Free improves the directly optimized HPSv3 reward in ReFL, while generator-side Schedule-Free improves DMD FID-5K at every evaluated checkpoint and lowers both generator and fake-score losses.In ReFL, HPSv3 rises from 6.3500 to 6.8626 at 4× LR and to 9.6134 at 20× LR; applying Schedule-Free only to DMD guidance instead yields worse FID-5K than the all-AdamW baseline at all three checkpoints.Source: §3.5 · pp. 8–9, Figure 5 · p. 9

4 CONCLUSION AND PRACTICAL RECOMMENDATIONS

The paper frames EMA, LR decay, and Schedule-Free as different ways to reduce trajectory variability, with distinct effects on continued optimization. Its recommendations favor Schedule-Free when the decay point is uncertain and decay with EMA when the training endpoint is fixed.

4 CONCLUSION AND PRACTICAL RECOMMENDATIONS

  • When the decay point is uncertain, prefer Schedule-Free; when the training endpoint is fixed, retaining EMA while annealing near the end can further improve the constant-LR + EMA recipe.The benefit of decay depends on when annealing begins, so it is most useful when the remaining training budget is known in advance.Source: §4 · p. 9
  • Choose a strategy for both stability and continued optimization rather than selecting solely for the lowest parameter variability.The paper characterizes EMA as averaging, LR decay as smaller updates, and Schedule-Free as strong stabilization without a prespecified decay point.Source: §4 · p. 9

A.1 MAIN PRETRAINING CONFIGURATION

  • The main pretraining comparison uses approximately 675M-parameter SiT-XL/2 for class-conditional 256 × 256 ImageNet-1K generation with flow matching and REPA.The main trajectories run up to 800k steps; constant-LR AdamW uses 10−4, while Schedule-Free uses a 4 × 10−4 base rate with a 5k-step warmup.Source: §A.1 · p. 13, §A.2 · p. 13, Table 1 · p. 13

A.3 LATE-STAGE OPTIMIZATION METRICS

The late-stage diagnostics track checkpoint variability, retained parameter movement, and objective progress to characterize optimization behavior. Checkpoint variability is a finite-window spread measure, while net displacement captures endpoint change rather than the full trajectory.

A.3.1 CHECKPOINT VARIABILITY

  • At checkpoint gap g = 1, the diagnostic uses 20 adjacent saved checkpoints to emphasize local parameter variability.For cosine decay only, the diagnostic continues from the decay endpoint at a fixed 10^-6 terminal rate.Source: §A.3.1 · p. 14
  • Checkpoint variability is total trace variance across 20 consecutive saved checkpoints, and can include both stochastic fluctuations and systematic movement.The finite-window statistic is not identified with pure theoretical stochastic variance.Source: Table 2 · p. 14, §A.3.1 · p. 14

A.3.2 LATE-STAGE NET PARAMETER DISPLACEMENT

  • Net displacement is the ℓ2 distance between checkpoints 1,000 optimization steps apart, measuring retained endpoint change rather than path length.The reported late-stage diagnostic covers five adjacent 1k-step windows.Source: Table 3 · p. 15, §A.3.2 · p. 15
  • The Constant-LR online trajectory retains the largest late-stage movement, while Schedule-Free moves less than Constant LR + EMA but more than both cosine-decay sequences.Within cosine decay, the model sequence moves less than EMA; explicit decay suppresses late-stage mobility most strongly.Source: §A.3.2 · p. 15, Table 3 · p. 15

A.3.3 FIXED-EVALUATION LOSS PROGRESS

The section compares late-stage objective progress under fixed evaluation inputs and tests stability at a larger learning-rate scale. Schedule-Free continues to reduce loss, while the post-decay cosine runs do not improve in the measured interval.

A.3.3 FIXED-EVALUATION LOSS PROGRESS

  • From 400k to 405k, Constant LR + EMA decreases from 0.43738 to 0.43723 and Schedule-Free decreases from 0.43461 to 0.43447, while neither 40k Cosine Decay sequence improves.All methods use identical fixed evaluation data, noise, and timestep sampling.Source: Table 4 · p. 15, §A.3.3 · p. 16

A.4 LEARNING-RATE STABILITY ABLATION

  • At 4 × 10^-4 with a 5k-step warmup, standard AdamW develops a loss spike and sustained NaNs after approximately 218.7k steps, whereas Schedule-Free remains stable.Standard AdamW is stable at 10^-4 in this setting.Source: Figure 7 · p. 16, §A.4 · p. 16

A.3.3 FIXED-EVALUATION LOSS PROGRESS

  • Fixed-evaluation loss and net displacement are complementary diagnostics: one tests objective improvement, while the other measures retained parameter change over a fixed window.The fixed evaluation uses identical data, noise, and timestep samples to reduce evaluation randomness.Source: §A.3.3 · pp. 15–16

A.5 REFL EXPERIMENTAL DETAILS

The appendices specify ReFL and DMD experimental settings and provide theoretical results on SGD stationarity and EMA variance reduction. The reported theory includes a vanishing stationarity rate for a suitable learning-rate choice and lower stationary variance for EMA in the stated quadratic setting.

A.5 REFL EXPERIMENTAL DETAILS

  • ReFL is run on Z-Image-Turbo with the official configuration fixed except for optimizer, learning rate, and warmup, using HPSv3 as the sole training reward.At checkpoint 100, runs are evaluated on the same 1,000 fixed prompts and seeds; other listed reward models are used only for cross-reward evaluation.Source: §A.5 · pp. 16–17, Table 5 · p. 17

A.6 DMD EXPERIMENTAL DETAILS

  • The DMD setup compares all-AdamW training with component-wise Schedule-Free ablations, changing either the generator or guidance optimizer while holding remaining settings fixed.The generator-side ablation uses averaged x weights, while the guidance-side ablation leaves the generator on baseline AdamW.Source: §A.6 · p. 17

B.1.1 STATIONARY DISTRIBUTION ON QUADRATIC OBJECTIVES

  • The quadratic-objective theorem considers one-dimensional SGD with constant-step updates, a positive quadratic coefficient, Gaussian gradient noise, and deterministic initialization.The supplied passage states the theorem setting but does not include its stationary-distribution expression.Source: §B.1.1 · p. 18

B.1.2 SMOOTH NONCONVEX STATIONARITY

  • For smooth nonconvex optimization, choosing η = Θ(T^-1/2) gives an average squared-gradient norm of O(T^-1/2), with the same order for the best iterate.The stated result assumes the setup of the smooth nonconvex SGD theorem.Source: §B.1.2 · p. 19

B.2.1 STATIONARY VARIANCE REDUCTION

  • Under the stated one-dimensional quadratic setting, EMA reduces the stationary variance of the evaluation iterate.The theorem compares the stationary variance of the EMA iterate with that of the online iterate.Source: §B.2.1 · p. 20

B.2.2 A LIMITATION OF EMA: PERSISTENT OPTIMIZATION BIAS

The appendix shows that EMA can reduce variability without removing persistent optimization bias, while decaying learning rates and Schedule-Free have vanishing stationarity guarantees under stated nonconvex assumptions. Late-stage comparisons find better generation quality for Schedule-Free than constant-LR EMA, alongside lower fixed-evaluation loss.

B.2.2 A LIMITATION OF EMA: PERSISTENT OPTIMIZATION BIAS

  • In the nonquadratic strongly convex example, EMA preserves the stationary mean bias, reducing variance around the mean without moving it to the optimizer.For the stated step-size range, SGD has a unique invariant distribution whose mean is strictly greater than the minimizer.Source: §B.2.2 · p. 21

B.3.1 EXACT DYNAMICS UNDER LEARNING-RATE DECAY AND EMA

  • Under the inverse-time learning-rate schedule, the EMA iterate has mean, variance, and expected objective all of order O(t^-1).The appendix states that EMA preserves the contraction obtained under learning-rate decay; the closed-form result uses an inverse-time schedule.Source: §B.3.1 · p. 25

B.3.2 SMOOTH NONCONVEX STATIONARITY WITH LEARNING-RATE DECAY

  • With summable decaying learning rates under the stated smooth nonconvex assumptions, the expected squared gradient norm at a learning-rate-weighted iterate vanishes.The minimum expected squared gradient norm over the iterates also tends to zero.Source: §B.3.2 · pp. 26–27

B.4.1 QUADRATIC VARIANCE REDUCTION

  • For the quadratic Schedule-Free setting, the appendix establishes that the averaged sequence has mean E[xt] = O(t^-1).The stated result assumes a one-dimensional quadratic objective, Gaussian noise, and deterministic initial values.Source: §B.4.1 · p. 27

C.1 ADDITIONAL LATE-STAGE COMPARISONS

  • From 320k to 400k, Schedule-Free AdamW beats constant-LR EMA on both FID-5K and Inception Score; its fixed-evaluation total loss is also lowest.The cosine-decay markers are endpoints of independent branches, not a continuous trajectory.Source: Figure 8 · p. 28, §C.1 · p. 28, Figure 9 · p. 29

B.4.2 SMOOTH NONCONVEX STATIONARITY

  • The Schedule-Free nonconvex guarantee concerns gradient-evaluation points yt and does not claim last-iterate convergence of xT.The cited theorem assumes an L-smooth objective bounded below.Source: §B.4.2 · p. 28

C.1 ADDITIONAL LATE-STAGE COMPARISONS

  • Schedule-Free maintains distinctly lower projection loss than constant-LR EMA and 40k decay plus EMA, while denoising MSE is similar across the three.The lower projection loss accompanies the lowest total fixed-evaluation loss over the evaluated interval.Source: Figure 9 · p. 29

C.2 MAIN FID-5K RESULTS

Across five checkpoints, Schedule-Free improves FID-5K over constant-LR EMA and has higher Inception Score. Cosine-decay results are separate 40k-step branches, so their outcomes depend on the starting checkpoint.

C.2 MAIN FID-5K RESULTS

  • Each cosine-decay result is a separate 40k-step branch from a late-stage checkpoint, not a continuation of one trajectory.The branch is compared with the Constant LR + EMA checkpoint at the same ending step.Source: §C.2 · p. 29, Table 7 · p. 30
  • The reported FID-5K entries are means over three sampling seeds; lower FID-5K and sFID and higher Inception Score are better.Source: Table 7 · p. 30
  • FID-5K reductions of 0.0784 to 0.1635 and higher Inception Score at all five checkpoints favor Schedule-Free over Constant LR + EMA.The cosine-decay branches show timing sensitivity: four of five improve endpoint FID-5K, while the branch starting at 320k is worse at 360k (9.3163 versus 9.2700).Source: §C.2 · pp. 29–30, Table 7 · p. 30

C.3 TIMESTEP-WISE DENOISING ERROR

Under fixed evaluation data and noise, Schedule-Free has lower denoising MSE than Constant + EMA at every tested timestep and with random timestep sampling. The relative reductions are modest but consistently positive.

C.3 TIMESTEP-WISE DENOISING ERROR

  • The paired evaluation reuses the same data and noise for both methods and varies only the timestep.It tests t ∈ {0.02, 0.10, 0.20, 0.40, 0.60, 0.80, 0.90, 0.96} plus randomly sampled timesteps.Source: §C.3 · p. 30
  • Schedule-Free reduces denoising MSE by 0.05% to 0.25% across eight fixed timesteps and by 0.12% under random timestep sampling.The largest reduction is at t = 0.02; reductions remain positive at t = 0.90 and t = 0.96.Source: §C.3 · pp. 30–31, Table 8 · p. 31
  • The direction of the denoising-MSE improvement is consistent across the diffusion trajectory, rather than being driven by one timestep.Source: §C.3 · p. 31

D DISCUSSION

The discussion cautions that optimizer rankings can depend on the surrounding training regime and that the cross-study reversal cannot be attributed to scale alone. Additional figures provide qualitative samples from pretraining, ReFL, and generator-side DMD.

D DISCUSSION

  • Schedule-Free's optimizer ranking differs across studies: the benchmark pairs competing optimizers with late-stage decay, while this SiT+REPA setting preserves its base learning rate after warmup.The authors say the relative ranking can depend strongly on the surrounding training regime.Source: D DISCUSSION · p. 31
  • The cross-study reversal cannot be attributed to model or dataset scale because the studies also differ in formulation, architecture, objective, learning-rate protocol, budget, and data domain.A controlled study varying model and data scale while independently retuning optimizers is left for future work.Source: D DISCUSSION · pp. 31–32

E.1 MAIN PRETRAINING: SCHEDULE-FREE ADAMW

  • Figure 10 shows representative class-conditional samples from the 800k-step Schedule-Free checkpoint across varied semantic categories, complementing quantitative generation metrics.Source: §E.1 · p. 32, Figure 10 · p. 33

E.2 REFL QUALITATIVE COMPARISONS

  • The ReFL figures compare the base model, AdamW at 1×, and Schedule-Free AdamW at 4× and 20× across fixed prompts; Figure 15 shows generator-side Schedule-Free DMD samples.The DMD figure applies Schedule-Free AdamW to the generator while retaining AdamW for the guidance model.Source: §E.2 · p. 34, Figure 11 · p. 34, Figure 12 · p. 35

E.3 DMD: SCHEDULE-FREE GENERATOR SAMPLES

Figure 15 shows samples from the DMD ablation variant that achieves the lowest FID-5K among three variants at 3k generator steps, with Schedule-Free AdamW on the generator and baseline AdamW on the guidance model.

E.3 DMD: SCHEDULE-FREE GENERATOR SAMPLES

  • At 3k generator steps, Schedule-Free AdamW on the generator with baseline AdamW on the guidance model achieves the lowest FID-5K among three variants.Figure 15 shows samples from this best-performing configuration.Source: §E.3 · p. 38
  • Figure 15 provides a qualitative view of the generator-side Schedule-Free improvement discussed in Section 3.5.Source: §E.3 · p. 38
  • Complete optimizer settings for the configuration are given in Appendix A.6.Source: §E.3 · p. 38
Loading 2610.22945v1…