Research map / Image generation / Fast diffusion sampling
Analyzing and Improving the Training Dynamics of Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Redesigns the ADM diffusion U-Net into a magnitude-preserving architecture (EDM2) with controlled weight and update magnitudes, and introduces post-hoc EMA reconstruction from stored snapshots, reaching FID 1.81 on ImageNet-512.
Problem
Diffusion model training is uneven and inefficient because activation, weight, and update magnitudes drift uncontrollably in the ADM U-Net, and the optimal EMA length is hard to tune since it requires separate training runs.
Approach
Make every layer magnitude-preserving on expectation (dividing by weight norms, forced weight normalization to unify effective learning rate, fixed-function gains), and reconstruct arbitrary power-function EMA profiles post hoc as least-squares combinations of stored averaged snapshots.
Key results
- EDM2-XXL reaches FID 1.91 without guidance and 1.81 with guidance on ImageNet-512, improving the previous record of 2.41 (StyleGAN-XL) with deterministic sampling.
- The EDM2-S configuration reduces FID from 8.00 (ADM-based baseline) to 2.56 on ImageNet-512 at equal complexity, beating VDM++ at 2.99 without guidance.
- On ImageNet-64, EDM2 reaches FID 1.33 with 63 NFE, improving the deterministic-sampling record of 2.22 (EDM).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply minor training improvements, including hyperparameter tuning and disabling 32×32 self-attention (ablation) | ADM diffusion model · ImageNet-512 | FID: 8.00 to 7.24 | Helped |
| Streamline the architecture, including removing biases, changing initialization and positional encoding, and using cosine attention (ablation) | CONFIG B · ImageNet-512 | FID: 7.24 to 6.96 | Helped |
| Add magnitude-preserving normalization to learned layers (ablation) | CONFIG C · ImageNet-512 | FID: 6.96 to 3.75 | Helped |
| Add forced weight normalization and inverse square root learning-rate decay (ablation) | CONFIG D · ImageNet-512 | FID: 3.75 to 3.02 | Helped |
| Replace group normalization with fewer pixel normalization layers (ablation) | CONFIG E · ImageNet-512 | FID: 3.02 to 2.71 | Helped |
| Add magnitude-preserving fixed-function layers (ablation) | CONFIG F · ImageNet-512 | FID: 2.71 to 2.56 | Helped |
| Add group normalization to the main path (ablation) | ADM diffusion model · ImageNet-512 | FID | Hurt |
| Enable dropout when the training loss falls but validation loss and FID begin increasing (ablation) | EDM2 larger configurations (M–XXL) · ImageNet-512 | FID | Helped |
| Use an XS unconditional model to guide an XXL conditional model rather than a larger unconditional model (ablation) | EDM2 conditional model · ImageNet-512 | FID | No clear effect |
| Add modest classifier-free guidance (1.2) | EDM2-XXL · ImageNet-512 | FID: 1.91 to 1.81 | Helped |
| Scale EDM2 model capacity from S to XXL | EDM2 · ImageNet-512 | FID (no CFG): 2.56 to 1.91 | Helped |
| Scale EDM2 model capacity from S to L | EDM2 · ImageNet-64 | FID (deterministic): 1.58 to 1.33 | Helped |
Surprising findings
- The optimal EMA length depends very strongly on classifier-free guidance strength, suggesting prior guidance comparisons may be affected by non-optimal EMA.
- A small unconditional XS model suffices to guide the largest XXL conditional model, with larger unconditional models giving no gain.
Limitations
- Stochastic sampling was not tested, so results may improve further but are left for future work.
- Post-hoc EMA reconstruction is approximate and per-tensor EMA choice was not explored beyond a single experiment.
- Evaluation is limited to ImageNet class-conditional generation; generality to text-to-image or other architectures (DiT, RIN) is not shown.
Experimental setup
Datasets: ImageNet-512 ImageNet-64
Benchmarks: ImageNet-512 ImageNet-64
Metrics: FID gigaflops per evaluation number of function evaluations
Baselines: ADM ADM-U DiT-XL/2 RIN U-ViT VDM++ StyleGAN-XL EDM
Builds on: ADM EDM Kendall multi-task loss weight normalization DiT
Closest papers by meaning
- Multi-Architecture Multi-Expert Diffusion Models (AAAI 2024 · similarity 0.90)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.90)
- Neural Diffusion Models (ICML 2024 · similarity 0.90)
- Elucidating the Exposure Bias in Diffusion Models (ICLR 2024 · similarity 0.90)
- A Simple Early Exiting Framework for Accelerated Sampling in Diffusion Models (ICML 2024 · similarity 0.90)
- CacheQuant: Comprehensively Accelerated Diffusion Models (CVPR 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.