Analyzing and Improving the Training Dynamics of Diffusion Models

Tero Karras; Miika Aittala; Jaakko Lehtinen; Janne Hellsten; Timo Aila; Samuli Laine

Source-linked AI summary

Analyzing and Improving the Training Dynamics of Diffusion Models

Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine

arXiv:2312.02696v2Dec 5, 2023cs.CVcs.AIcs.LGcs.NEstat.ML

TL;DR

Diffusion-model training remains challenging because stochastic losses require accurate predictions across many noise levels, noise realizations, and conditioning inputs, while weight and update magnitudes can drift unevenly. The paper redesigns ADM layers to control these magnitudes and introduces post-hoc EMA, achieving ImageNet-512 FID 1.81 and enabling cheaper EMA studies.

Problem
Diffusion-model training remains challenging because stochastic losses require accurate predictions across many noise levels, Gaussian-noise realizations, and conditioning inputs.
Method
The paper redesigns ADM layers to control activation, weight, and update magnitudes, removes most group normalization, and introduces post-hoc EMA analysis.
Results
1.81 ImageNet-512 FID was achieved with modest guidance, improving the unguided 1.91 result and surpassing the previous record.
Takeaways & Limitations
The drop-in denoiser replacement is intended to support large-scale image generators, while post-hoc EMA makes extensive EMA studies substantially cheaper.
Takeaways & Limitations
The post-hoc EMA findings are preliminary and anecdotal, with detailed study left for future work.

Abstract

from arXiv · show

Diffusion models currently dominate the field of data-driven image synthesis with their unparalleled scaling to large datasets. In this paper, we identify and rectify several causes for uneven and ineffective training in the popular ADM diffusion model architecture, without altering its high-level structure. Observing uncontrolled magnitude changes and imbalances in both the network activations and weights over the course of training, we redesign the network layers to preserve activation, weight, and update magnitudes on expectation. We find that systematic application of this philosophy eliminates the observed drifts and imbalances, resulting in considerably better networks at equal computational complexity. Our modifications improve the previous record FID of 2.41 in ImageNet-512 synthesis to 1.81, achieved using fast deterministic sampling. As an independent contribution, we present a method for setting the exponential moving average (EMA) parameters post-hoc, i.e., after completing the training run. This allows precise tuning of EMA length without the cost of performing several training runs, and reveals its surprising interactions with network architecture, training time, and guidance.

1. Introduction

Diffusion models generate images by repeatedly denoising pure noise, with each step related to score matching and typically implemented using a U-Net with self-attention.

Diffusion models convert pure noise into novel images through repeated application of image denoising.

NVIDIA

The paper studies difficult, stochastic diffusion training dynamics and proposes magnitude-focused ADM modifications plus post-hoc EMA tuning. These contributions improve ImageNet-512 quality while preserving the ADM network’s overall structure and enabling more efficient EMA exploration.

Diffusion training is difficult because stochastic signals span noise levels, noise realizations, and conditioning inputs, while intermediate sampling errors can snowball.
Current designs may produce uneven parameter responses, complicating improvements through interactions among hyperparameters, architecture, and training setups.
The paper standardizes expected magnitudes of weights, activations, gradients, and updates through coordinated design choices addressing their interdependencies.
The modified ADM U-Net preserves the overall structure, achieves ImageNet-512 FIDs of 1.81 with guidance and 1.91 without guidance, and uses fast deterministic sampling.
Post-hoc EMA reconstructs networks with arbitrary EMA profiles from preintegrated training snapshots, enabling computationally cheaper EMA exploration.
The implementation and pretrained models are publicly available through the EDM2 repository.

2. Improving the training dynamics

The paper incrementally redesigns the ADM U-Net to control loss, activation, weight, and update magnitudes without changing its overall structure. These changes remove training drifts and substantially improve ImageNet-512 FID.

Preliminary changes: The study evaluates incremental ADM U-Net modifications on latent ImageNet-512 images using deterministic sampling and reports FID without guidance.The exploration network has approximately 300M trainable parameters, and FID is computed from 50,000 generated images.
Preliminary changes: Hyperparameter tuning, loss reweighting across noise levels, and removing 32×32 self-attention reduce FID from 8.00 to 7.24.The loss is scaled by the reciprocal of the raw loss as a function of noise level to counter changing gradient feedback.
Preliminary changes: Architectural streamlining and cosine attention reduce FID from 7.24 to 6.96 while enabling 16-bit floating-point computation throughout the network.The changes remove additive biases, add a constant input channel, unify weight initialization, and normalize key and query vectors before attention dot products.
Standardizing activation magnitudes: Activation magnitudes grow uncontrollably in CONFIG C because residual paths accumulate contributions along long unnormalized signal paths.Adding group normalization to the main path degraded result quality, motivating an alternative magnitude-preserving design.
Standardizing activation magnitudes: 6.96 to 3.75: The activation-magnitude modification eliminates magnitude drift and substantially improves FID.The modification removes direct means for the network to learn changes in overall activation magnitudes.
Standardizing weights and updates: 3.75 to 3.02: Forced weight normalization and inverse square root learning-rate decay preserve activation and weight magnitudes during training.Weight vectors are normalized before each training step, while the decay schedule provides direct control over effective learning rate.

3. Post-hoc EMA

The paper introduces post-hoc EMA selection by reconstructing arbitrary averaging profiles from stored parameter snapshots, then uses this method to analyze how EMA length interacts with configurations, tensors, and training progress.

Motivation and overview: Post-hoc EMA selection enables dense evaluation of EMA lengths without specifying the profile before training.The method is designed to reveal interactions with network architecture, training time, and classifier-free guidance.
Power-function EMA profiles: Power-function profiles avoid weighting random initialization and automatically stretch with training duration.Their incremental update resembles traditional EMA, but its decay factor depends on the current training time.
Post-hoc reconstruction: The method stores two power-function averages at periodic snapshots and reconstructs arbitrary profiles using a least-squares linear combination.The stored profiles correspond to relative standard deviations of 0.05 and 0.10.
Post-hoc reconstruction: A few dozen snapshots are sufficient for virtually perfect EMA reconstruction, with reconstruction error experimentally decreasing as O(1/n^4).The figure illustrates the approximation using two averaged parameter vectors per snapshot.
EMA-length analysis: Optimal EMA length differs across configurations, while the optimum narrows toward the final configuration.Figure 5a reports FID as a function of EMA length, with shaded regions showing the minimum and maximum across three evaluations.
EMA-length analysis: In CONFIG B, tensors disagree about preferred EMA lengths and per-tensor changes can improve FID by up to 10%, from 7.24 to approximately 6.5.In CONFIG G, tensors agree on the optimum, and deviating from the common 13% EMA length is detrimental.
EMA-length analysis: The optimal EMA length shifts gradually toward relatively longer averages as training progresses.This shift occurs even though EMA length is already defined relative to total training duration.

4. Results

On ImageNet, the method achieves strong deterministic-sampling results across resolutions while exposing important interactions among guidance, EMA length, and model configuration. A small unconditional model can provide guidance for much larger conditional models, reducing guidance overhead.

Training behavior: Dropout improves results under overfitting, but is harmful in smaller configurations and therefore enabled only in larger configurations showing overfitting.Overfitting is identified when training loss decreases while validation loss and FID increase.
EMA and guidance: The optimal EMA length depends strongly on guidance strength, making post-hoc EMA useful for evaluating this interaction efficiently.The authors suggest that some prior differences between vanilla and guided results may reflect non-optimal EMA parameters.
ImageNet-512: FID 1.91 without guidance surpasses the previous 2.99 record on ImageNet-512, while modest guidance of 1.2 further improves FID to 1.81.The 1.81 result sets a new record for ImageNet-512.
Low-cost guidance: An XS unconditional model sufficiently guides even the largest XXL conditional model; enlarging the unconditional model does not improve results.This reduces the typical 2× computational overhead associated with guidance.
ImageNet-64: FID 1.58 improves the previous ImageNet-64 deterministic-sampling record of 2.22, reaching 1.33 through scaling.The L-sized model saturates this dataset.

5. Discussion and future work

The improved denoiser is intended as a drop-in ADM replacement that supports further architectural studies. Post-hoc EMA makes extensive EMA investigations feasible, while applicability beyond the tested setting remains open.

Discussion: The improved denoiser is a drop-in replacement for the widely used ADM network.The authors hope it will find widespread use in large-scale image generators.
Discussion: Reducing entanglement among training aspects should make local U-Net modifications easier to study without breaking other components.
Future work: Whether the magnitude-focused methodology benefits other diffusion architectures, such as RIN and DiT, or other application areas remains an open question.The authors note that related work has received relatively little attention outside ImageNet classifiers.
Discussion: Post-hoc EMA reduced the cost of producing some plots from a thousand GPU-years to a GPU-month.The authors expect cheap EMA data to support studies of EMA’s role and principled EMA-length selection.

A.2. Quality vs. compute

The improved model reaches comparable quality more quickly across parameter, sampling, and training-cost measures. Dropout and guidance choices depend on model size, while post-hoc EMA reduces sensitivity to learning-rate decay and enables cheaper unconditional guidance.

Quality vs. compute: Across parameter, sampling, and training-cost metrics, the model reaches the same quality much quicker and continues improving.The sampling-cost advantage is especially pronounced because the model uses fewer score function evaluations.
Quality vs. compute: Early cleanup in CONFIG B improves both convergence and execution speed during equal-length wall-clock training runs.Other tracked metrics were largely consistent with FID, motivating the paper’s standardization on FID.
Quality vs. compute: A very small unconditional model reduces guided-diffusion computational cost by almost 50%, while large unconditional models are not useful for classifier-free guidance.The reported EMA lengths apply to both conditional and unconditional models.
Quality vs. compute: With post-hoc EMA, learning-rate decays tref ∈[30k, 160k] yield FIDs within 10% of the optimum.With EMA fixed at 13%, varying tref increased FID by as much as 72% in the tested range.

A.5. Fréchet distances using DINOv2

FID and FDDINOv2 can prefer substantially different EMA lengths and guidance strengths, so EMA and guidance should be tuned for the chosen quality metric. The architectural changes also improve FDDINOv2 comparably to FID.

Metric setup: FDDINOv2 is supplied as an alternative Fréchet-distance metric because its feature space aligns better with human preferences than InceptionV3’s.The evaluation uses 50,000 generated images and all 1,281,167 available real images.
EMA and metric dependence: Without guidance, FID and FDDINOv2 prefer clearly different EMA lengths, with FDDINOv2 favoring longer EMA.Each metric considers the other’s optimum a poor choice.
EMA and metric dependence: With guidance strength 1.4, FID and FDDINOv2 prefer EMA lengths of 2% and 14%, respectively.Both metrics agree that modest guidance helps, but they strongly disagree on the preferred EMA length.
Guidance dependence: FDDINOv2 prefers guidance strength 1.9, whereas FID prefers 1.4 and considers 1.9 clearly excessive.
Architectural comparison: The CONFIG B-to-G changes improve FDDINOv2 at least as significantly as the improvement observed with FID.The authors associate FDDINOv2’s preference with better global coherency in a cursory image assessment.
Magnitude analysis: Activation and weight magnitudes are analyzed across encoder and decoder resolution buckets using maximum or mean aggregation.The analysis identifies trainable tensors and measures immediate pre-nonlinearity operation outputs, with aggregation choices targeting extreme and average behavior.
Magnitude analysis: The magnitude computation scales per-feature L2 norms so unit-normal distributed activations or weights have expected magnitude 1 regardless of tensor dimensions.Maximum aggregation captures potential extremes, while mean aggregation better reflects average behavior; their qualitative behavior is similar.

B. Architecture details

The appendix details the ADM/EDM baseline, latent-diffusion preprocessing, U-Net components, attention implementation, initialization, loss, and configuration-specific analyses. It also documents comparisons involving dropout, EMA, and FID/FDDINOv2 evaluation.

Configuration analysis: The appendix includes configuration diagrams and hyperparameters for the architectural changes, alongside analyses of dropout, EMA length, and FID/FDDINOv2.
EDM baseline: The baseline is ADM implemented in EDM, trained in VAE latent space for 2^19 iterations with batch size 4096.Traditional EMA uses a half-life of 50M images, corresponding to 12k training iterations.
Preconditioning: The denoiser preconditions raw-network inputs and outputs using cin, cout, cskip, and cnoise, with σdata representing training-data standard deviation.
Latent diffusion: ImageNet-512 images are encoded into latents with a pretrained VAE and decoded after generation, with fresh latent samples drawn during training.The encoder represents each image as a distribution of latents rather than a strictly bijective mapping.
Architecture walkthrough: The ADM architecture uses a U-Net with encoder-decoder blocks, skip connections, class and noise conditioning, residual paths, and self-attention at resolutions of 32×32 and below.The network expands inputs to 192 channels and contracts the final activation to four channels.
Self-attention: Attention uses one head per 64 incoming channels, with dot products and softmax computed in 32-bit floating point.
Initialization and loss: Most layers use He’s uniform initialization, while final residual, attention, and output-block weights and biases are initialized to zero.The zero initialization makes the initial denoiser behave as Dθ(x, σ) = cskip(σ)x.
Initialization and loss: The training objective is the EDM denoising score-matching loss for clean images y sampled from the training set and Gaussian noise n.

B.2. Minor improvements (CONFIG B)

CONFIG B improves the EDM baseline by retuning training hyperparameters and adding adaptive loss weighting across noise levels. The configuration also introduces architectural and averaging changes documented for subsequent variants.

Training hyperparameters: Halving the batch size to 2048, doubling αref to 0.0002, and reducing Adam β2 to 0.99 sped convergence, especially early in training.These changes also helped somewhat toward the end of training.
Loss weighting: Because static λ(σ) cannot maintain balanced gradient magnitudes as losses evolve, the method treats continuous noise levels as an infinite multitask-learning problem.The overall loss assigns each noise level σ a weighted contribution λ(σ)L(Dθ; σ).
Training hyperparameters: The improved CONFIG B baseline uses FID 7.24 with 291.8M parameters, 100.4 Gflops model complexity, and 7.59 Tflops sampling cost.The configuration trains for 2147.5 million images with minibatch size 2048.
Loss weighting: The uncertainty-based formulation downweights tasks with high predicted uncertainty while penalizing uncertainty, encouraging σi to decrease.The approach uses an uncertainty function u(σ) for the continuous noise-level setting.
Architectural changes: Cosine attention makes attention weights depend only on query-key directions, curbing uncontrolled growth and enabling 16-bit floating point throughout self-attention.This addresses the quadratic scaling of attention weights with query and key magnitudes.

B.5. Controlling effective learning rate (CONFIG E)

CONFIG E controls effective learning rate by preventing raw weight magnitudes from drifting during optimization. Forced weight normalization removes the implicit, unequal learning-rate decay caused by weight growth and makes explicit scheduling meaningful.

Weight growth: Weight norms increase after training steps because tangent-plane gradients move weights farther from the origin, with the effect also appearing in CONFIG C.The phenomenon occurs unless the weight gradient is zero.
Forced weight normalization: Forced weight normalization constrains optimization to a weight hypersphere and eliminates the drift in raw weight magnitudes.The method re-normalizes weights after each training step.
Effective learning rate: Without correction, growing weights impose implicit learning-rate decay that can drift uncontrollably and unequally between layers.Larger weights reduce the effective learning rate under the constrained dynamics.
Effective learning rate: After forced normalization, the effective learning rate is directly proportional to α, so CONFIG E uses an explicit inverse-square-root decay schedule.The constant schedule in CONFIGS A–D is recovered when tref = ∞.
Optimizer compatibility: Standard weight normalization remains necessary because it makes Adam estimate variance from tangent-plane updates rather than erased normal components.Using only forced normalization corrupts variance estimates and produces smaller, uncontrolled updates.

B.7. Magnitude-preserving fixed-function layers (CONFIG G)

CONFIG G extends magnitude preservation beyond learned layers to fixed-function operations, aiming to standardize activations throughout the network. The resulting activation magnitudes closely approach the target in practice while retaining learned gains where deviations are needed.

Magnitude-preserving fixed-function layers: CONFIG G extends magnitude-preserving design to the remaining fixed-function layers, completing the treatment of learned and non-learned operations.The changes are defined by comparing the CONFIG F and CONFIG G architectures.
Activation standardization: The design aims to standardize every network input and preserve standardized activation magnitude through every operation.If both conditions hold, all activations throughout the network are standardized.
Activation standardization: The standardization assumptions are not strictly true in practice, but CONFIG G activations are nevertheless surprisingly close to the ideal.This is observed in the activation-mean plot for CONFIG G.
Blending operations: Residual and main paths use a blend factor, with t = 0.3 giving the residual path 30% and the main path 70% in encoder, decoder, and self-attention blocks.The embedding network instead uses t = 0.5 for equal noise-level and class-label contributions.
Learned gain: A learned scalar gain is introduced where standardization must allow activation magnitude to deviate from one, including the final output and encoder-decoder scaling points.The gain is initialized to zero, making the residual output equal to the input at initialization.

C.5. Synthesizing novel EMA profiles after training

Post-hoc EMA synthesis reconstructs a desired averaging profile by linearly combining model snapshots collected during training. The combination weights are obtained by minimizing the L2 mismatch between the synthesized and target response functions.

Problem definition: The method approximates a novel EMA average from stored snapshots because the precise target average cannot be calculated from those snapshots alone.The snapshots are combined using learned coefficients xi.
Problem definition: The coefficients are chosen so the snapshots’ approximate response function matches the desired response function as closely as possible.The approximation is expressed as an inner product between snapshot response functions and the coefficient vector.
Least-squares solution: A least-squares objective minimizes the L2 distance between target and synthesized response functions, yielding the matrix equation Ax − b = 0 and solution x = A^-1b.The matrix and vector contain inner products between corresponding response functions.
Implementation: The implementation supports multiple target response functions simultaneously, with matrix columns representing targets and rows representing input snapshots.Standard numerical integration can compute the required inner products for arbitrary EMA profiles.
Implementation: Training tracks two weighted averages at σrel = 0.05 and 0.10, saving snapshots every 8 million images and storing 160–512 pre-averaged models.The snapshots use 16-bit floating point to conserve disk space.

D. Implementation details

The implementation uses deterministic EDM sampling, mixed-precision training, standardized ImageNet preprocessing, and protocol-based evaluation. Model complexity is estimated from traced tensor operations, while the paper also documents practical numerical safeguards and societal costs.

Sampling: Sampling uses EDM’s second-order deterministic sampler with N = 32, σmin = 0.002, σmax = 80, and ρ = 7.The same sampler settings are used for ImageNet-512 and ImageNet-64 experiments.
Guidance: Classifier-free guidance combines conditional and unconditional denoisers, with w = 1 disabling guidance and w > 1 strengthening it.The number of model evaluations is counted independently of the guidance weight.
Numerical implementation: Training stores parameters in FP32 but evaluates the model with FP16 parameters and activations, retaining FP32 for embedding-related layers.Gradients are forced finite and activations are clamped to [−256, +256] to mitigate rare FP16 overflows.
Evaluation: FID uses 50,000 generated images and all available real images, evaluated three times with the minimum reported.The protocol avoids augmentation and reports variation of typically ±2% across evaluations.
Model complexity estimation: Model complexity is estimated from traced PyTorch operations, with the largest ImageNet-512 model requiring 545.50 Gflops for convolution.For prior work without published complexity, the estimates are best-effort and believed accurate within 10%.