LIT-714: Analyzing and Improving the Training Dynamics of Diffusion Models

Karras et al. (2023) — ARXIV-2312.02696 (CVPR 2024). Read in full as NOTE-370, appendices included.

Key takeaways

• The diagnosis is magnitude drift, measured. In the ADM U-Net as run under EDM, activation magnitudes on the main path grow without bound and show no sign of stabilising by the end of training; weight magnitudes grow too. Adding normalization to the main path was tried and made results “significantly” worse. The fix is to make every layer preserve magnitude on expectation, without looking at the activations.

• The ablation ladder, ImageNet-512 latent, ~300M parameters, FID without guidance (each row cumulative; each number the minimum of three FID evaluations of one model):

  config	change	FID
  A	EDM baseline (ADM)	8.00
  B	retuned hyperparameters, learned per-noise-level loss weight, no 32×32 attention	7.24
  C	no biases, cosine attention, simplified group norm	6.96
  D	magnitude-preserving learned layers (weight normalization, no learned gain)	3.75
  E	forced weight normalization + inverse-sqrt LR decay	3.02
  F	group norms removed, ¼ as many pixel norms	2.71
  G	magnitude-preserving SiLU, Fourier features, sums and concatenation	2.56

  D is the step: normalizing each output channel’s weight vector before use severs the link from weight magnitude to activation magnitude. It then makes weights grow faster (the update is orthogonal to the weight, so each step lengthens it), which silently decays the effective learning rate, and E removes that by renormalizing the stored weights after every step while still normalizing on use — the second normalization is what keeps Adam’s second-moment estimate on the tangent-plane step.

• Scaled up, it sets records at far lower cost. EDM2-XXL (1.5B): FID 1.91 unguided, 1.81 guided on ImageNet-512, against 2.99 / 2.41 before, with 63 deterministic function evaluations where earlier diffusion entries used 250–1000 stochastic ones. EDM2-S (280M, 102 Gflops) already beats the previous unguided record. ImageNet-64 in pixel space: 2.22 (EDM) → 1.58 at matched complexity, 1.33 scaled.

• Post-hoc EMA is a least-squares fit over stored averages. Replace the exponential profile with a power function t^γ, which gives the initial weights zero weight and stretches with training length; track two averages (σ_rel 0.05 and 0.10) and save both every ~8M images; afterwards, solve a small linear system for the combination of snapshots whose profile best matches any target profile. Reconstruction error falls about as O(1/n⁴) in the number of snapshots, and 160–512 snapshots in fp16 give “nearly perfect” reconstruction over σ_rel ∈ [0.015, 0.25].

• What it reveals is that there is no EMA length to carry over. The optimum moves across configs B–G; it narrows as the architecture improves, and per-tensor sweeps explain why — in B, moving one tensor’s EMA alone could improve FID ~10%, in G the tensors agree. It moves slowly longer through training. It moves with guidance strength “very strongly”. And it moves with the metric: FID and FD_DINOv2 disagree on the optimum unguided (13% vs 19%) and wildly at guidance 1.4 (2% vs 14%), each rating the other’s optimum as poor. They also disagree on the guidance weight (1.4 vs 1.9).

• Two cheap guidance findings. An XS unconditional model (125M) guides the XXL conditional model as well as any larger one does, cutting guidance overhead by almost half; and the authors suggest that the gap between guided and unguided results in prior work may partly be an artefact of an EMA length tuned for one and used for both.

Standing in the anthology

Filed from #290’s promoted list (item 10), and the curation entry that sent it here says why it was unforced: it is the successor of LIT-075, the source of SOTA-188, and neither of its two contributions was anywhere in this record — no file said “post-hoc EMA” or “magnitude-preserving”.

It now sources two practices of its own, both Proposed: SOTA-428 (choose the EMA length after training) and SOTA-431 (magnitude-preserving layers with forced weight normalization). Neither is Active, for the same reason: one group, one architecture family, and the authors’ own discussion asks whether the layer changes would help a DiT. It is not added as a source to SOTA-188, whose recommendation it does not test; it corrects one sentence there instead — EDM’s loss weighting equalises the noise levels at initialization only, and stops doing so as training proceeds. It also answers half of the capacity question SOTA-424 records as open.

DP-010 applies to the headline in a specific way. “1.81” is guided and “1.91” is not, and the gap between them rests on an EMA length the paper itself shows is metric-specific; the finding that should travel is Appendix A.5, where the two Fréchet distances pick EMA lengths seven-fold apart.

The records are set against quoted numbers, not reruns; only EDM is retrained. On ImageNet-512 (Table 2), ADM from LIT-699 is listed at 23.24 FID unguided and 7.72 guided, and its upsampling variant ADM-U at 9.96 / 3.85; DiT-XL/2 from LIT-448 at 12.03 / 3.04; the U-ViT L of simple diffusion (LIT-660) at 3.54 / 3.02; and VDM++ (LIT-692), the same U-ViT trained with Kingma and Gao’s weighting, at 2.99 / 2.65, the unguided record EDM2 then displaces. All four ran stochastic samplers at 250 or 256 evaluations, against EDM2’s 63 deterministic ones. On ImageNet-64 (Table 3) the pixel-space EDM2-S (1.58, deterministic) is set beside ADM’s 2.07 and VDM++‘s 1.43, both stochastic at 250 and 511 evaluations, so the deterministic and stochastic columns are not a like-for-like contest.

Two later papers use it. sCM (LIT-790) re-trains EDM2 in TrigFlow form as its ImageNet-512 teacher, with a corrected latent normalization, and beats the published numbers at XL and XXL, 1.80 against 1.85 and 1.73 against 1.81. Its 1.5B two-step sCD student reaches 1.88 against that teacher’s 1.73, and it adopts this paper’s learned per-time loss weight. STARFlow (LIT-785) quotes EDM2-XXL at 1.25 against its own 3.00 at ImageNet-512. That 1.25 is not among this paper’s numbers, whose guided headline is 1.81, and STARFlow does not say where it comes from.

The record holds no practice for the thing post-hoc EMA refines — sample from an average of the weights at all. The paper calls it “indispensable” and supplies no arm without it, so it cannot source that trunk; the gap is written down in NOTE-370 rather than filled with adoption.