PaperlayerView exact revision on arXiv

Source-linked AI summary

Elucidating the Design Space of Diffusion-Based Generative Models

Tero Karras, Miika Aittala, Timo Aila, Samuli Laine

arXiv:2206.00364v2cs.CVcs.AIcs.LGcs.NEstat.ML

TL;DR

Diffusion-model theory and practice often obscure which design choices can vary independently. The paper presents a modular design space spanning sampling, training, and score-network preconditioning, and reports improved quality and efficiency across CIFAR-10 and ImageNet-64.

  • Problem

    Prior diffusion-model derivations can obscure the available design space by making models appear as tightly coupled packages whose components cannot be modified independently.

  • Method

    The paper separates practical choices in sampling, score-network preconditioning, and training, including discretization, solver order, noise schedules, and noise-level sampling.

  • Results

    The design changes improve diffusion-model quality and efficiency, including ImageNet-64 FID from 2.07 to 1.55 with sampler replacement and to 1.36 after retraining, while CIFAR-10 uses 35 model evaluations.

  • Takeaways & Limitations

    The common framework supports targeted replacement and evaluation of individual diffusion-model components across pre-trained and retrained systems.

  • Takeaways & Limitations

    Stochastic sampling requires heuristic choices tuned case by case, and excessive noise injection can lose detail or produce oversaturated colors.

Abstract

from arXiv · show

We argue that the theory and practice of diffusion-based generative models are currently unnecessarily convoluted and seek to remedy the situation by presenting a design space that clearly separates the concrete design choices. This lets us identify several changes to both the sampling and training processes, as well as preconditioning of the score networks. Together, our improvements yield new state-of-the-art FID of 1.79 for CIFAR-10 in a class-conditional setting and 1.97 in an unconditional setting, with much faster sampling (35 network evaluations per image) than prior designs. To further demonstrate their modular nature, we show that our design changes dramatically improve both the efficiency and quality obtainable with pre-trained score networks from previous work, including improving the FID of a previously trained ImageNet-64 model from 2.07 to near-SOTA 1.55, and after re-training with our proposed improvements to a new SOTA of 1.36.

1 Introduction

The paper separates diffusion-model design choices that prior theory and practice often present as tightly coupled. It proposes changes to practical analysis, sampling, training, and network preconditioning, achieving improved FID and sampling efficiency.

  • Motivation: Diffusion-model literature can obscure independent design choices by presenting complete systems as tightly coupled packages.This makes it appear that modifying one component may break the overall system.
  • Design-space framework: The paper reframes diffusion models around tangible training and sampling objects, focusing on available degrees of freedom rather than only their statistical derivations.The framework targets models that learn score functions for Gaussian-noise-corrupted data through denoising score matching.
  • Sampling: The authors improve sampling through time discretization, higher-order Runge–Kutta integration, sampler schedules, and analysis of stochasticity.These changes reduce the number of sampling steps and can replace samplers in several widely used diffusion models.
  • Training and preconditioning: The paper analyzes input, output, and loss preconditioning, proposes a better noise-level distribution for training, and applies non-leaking augmentation to diffusion models.The work retains commonly used DDPM and NCSN architectures while targeting training dynamics.
  • Results: FID reached 1.79 for CIFAR-10 and 1.36 for ImageNet at 64×64 resolution.These are reported as record results obtained by combining the paper’s design changes.

2 Expressing diffusion models in a common framework

The common framework expresses diffusion models through separable choices for noisy distributions, score estimation, trajectories, discretization, and denoiser parameterization. This exposes interchangeable components while preserving the sampling and training roles of each choice.

  • Noisy distributions: Diffusion models construct noisy distributions p(x; σ) by adding Gaussian noise to the data, with sufficiently large σ producing an approximately Gaussian distribution.Sampling can begin from high-noise Gaussian-like data and progressively denoise toward the data distribution.
  • Continuous dynamics: The probability-flow ODE evolves samples between noise levels while preserving the corresponding distributions, using the score as a vector field toward higher data density.Backward evolution moves samples toward the data distribution, whereas forward evolution moves them away as noise increases.
  • Continuous dynamics: The noise schedule σ(t) is a practical design choice with major implications, not merely a theoretically convenient parameterization.The framework explicitly separates the schedule from other sampling components.
  • Denoiser parameterization: The denoiser Dθ is trained through a network-derived formulation, with pre- and post-processing treated as part of its parameterization.The framework defines Dθ from skip and output terms applied to the raw network Fθ.
  • Discretization: Numerical sampling independently requires choosing an integration scheme and discrete time points, with the paper reporting advantages for a second-order solver over Euler’s method.The sampling trajectory is obtained by taking finite steps over discrete time intervals.
  • Evaluation: Figure 2 compares deterministic samplers by FID versus neural function evaluations across three pre-trained models, marking the lowest NFE within 3% of each curve’s best FID.NFE measures how many denoiser evaluations are needed to generate one image.
  • Modularity: The framework removes implicit dependencies among components, so individual formulas can be changed without requiring compensating changes elsewhere, although combinations can perform differently.Table 1 collects formulas for deterministic variants of widely used earlier methods in this common representation.

3 Improvements to deterministic sampling

The paper treats deterministic sampling as a modular numerical integration problem, improving solver order, noise schedules, and time-step placement to reduce evaluations while preserving image quality.

  • Modular sampling: The sampler can be evaluated independently of training details because the denoiser is treated as a black box.The authors test different samplers on pre-trained models to assess this separation.
  • Time-step placement: The time-step schedule shortens steps as σ decreases, with ρ = 7 selected because values from 5 to 10 perform better for image sampling.The analysis indicates that errors near σmin have a large impact.
  • Higher-order integration: Heun’s 2nd order method reaches Euler’s FID with considerably lower NFE.The correction step accounts for changes in dx/dt between adjacent time points and has O(h3) local error.
  • Trajectory curvature and schedules: The chosen schedule σ(t) = t and s(t) = 1 produces largely linear solution trajectories and reduces curvature-related discretization error.With this schedule, a single Euler step to t = 0 yields the denoised image Dθ(x; t).
  • Results: The deterministic changes reduce the NFE needed for high-quality results by 7.3× for VP, 300× for VE, and 3.2× for DDIM.The authors report 26.3 high-quality CIFAR-10 images per second on a single NVIDIA V100.

4 Stochastic sampling

The stochastic sampler combines higher-order deterministic integration with controlled noise injection, improving low-step sampling while exposing trade-offs from heuristic stochasticity choices.

  • Role of stochasticity: Stochasticity can correct errors from earlier sampling steps by driving samples toward the desired marginal distribution at each time.Discrete approximation of the Langevin term also introduces error, so stochasticity is not uniformly beneficial.
  • Sampler design: The proposed sampler alternates noise injection to a higher noise level with a backward ODE step using the denoiser.It combines explicit Langevin-like churn with the 2nd order deterministic ODE integrator.
  • Sampler design: The method evaluates the denoiser after noise injection, unlike Euler–Maruyama’s interpretation that retains the initial state for that evaluation.This distinction becomes significant when using large steps and low NFE.
  • Practical limitations: Excessive stochasticity gradually removes image detail and can produce oversaturated colors at very low and high noise levels.The authors attribute this degradation to a suspected nonconservative vector field induced by practical denoisers.
  • Practical limitations: Restricting stochasticity to σ levels in [S_tmin, S_tmax] and clamping added noise addresses the observed color drift heuristically.The churn factor is set as γ_i = S_churn/N within the selected noise range.
  • Results: The stochastic sampler outperforms previous samplers by a significant margin, especially at low step counts.Its heuristic parameters require case-by-case grid search, and the benefit of adaptive solvers remains open.
  • Results: Sampler improvements alone reduce the ImageNet-64 model’s FID from 2.07 to 1.55.The result approaches reported reference values of 1.48 for cascaded diffusion, 1.55 for classifier-free guidance, and 1.52 for StyleGAN-XL.

5 Preconditioning and training

The paper redesigns score-network preconditioning and training around explicit, modular choices, then evaluates their effects on diffusion-model quality and sampling behavior.

  • Preconditioning: Directly training Dθ is poorly conditioned because x combines clean signal and noise whose magnitude varies substantially with σ.The paper instead trains Fθ, from which Dθ is derived.
  • Preconditioning: The proposed σ-dependent skip connection lets the network estimate y, n, or an intermediate combination while scaling inputs, outputs, and noise conditioning.The functions cskip(σ), cin(σ), cout(σ), and cnoise(σ) control these transformations.
  • Preconditioning: The preconditioning choices are derived by targeting unit-variance inputs and training targets while minimizing amplification of network errors.The cnoise formula is selected empirically.
  • Training: Balanced loss weighting and targeted log-normal noise sampling substantially improve FID when combined with the proposed preconditioning.The training distribution prioritizes intermediate noise levels where loss reduction is most achievable.
  • Training: Augmentation regularization applies geometric transformations before noise addition and conditions Fθ on augmentation parameters so inference produces non-augmented images.The method is intended to prevent overfitting on smaller datasets.
  • Sampling interaction: Deterministic sampling was best for the improved CIFAR-10 training setup, whereas stochastic sampling remained beneficial for class-conditional ImageNet-64.The ImageNet-64 model achieved FID 1.36, compared with the previous record of 1.48.

6 Conclusions

The common framework exposes diffusion models as modular systems whose samplers and training components can be investigated independently. The resulting changes improve prior models and achieve strong CIFAR-10 and ImageNet-64 results, while higher-resolution transfer and sampler–objective interactions remain open issues.

  • 6 Conclusions: A common framework exposes a modular design for targeted investigation of individual diffusion-model components.The authors present this modularity as a way to explore the viable design space more systematically.
  • 6 Conclusions: Replacing samplers improved a pre-trained ImageNet-64 model from FID 2.07 to 1.55, while training improvements reached FID 1.36.The 1.36 result is reported as state of the art.
  • 6 Conclusions: The approach also obtained state-of-the-art CIFAR-10 results using 35 model evaluations, deterministic sampling, and a small network.The paper positions these changes as orthogonal to several high-resolution extensions.
  • 6 Conclusions: Many parameter values may need readjustment for higher-resolution datasets, and the interaction between stochastic sampling and the training objective remains unresolved.These are identified as boundaries and future-work questions.
  • Societal impact: The project consumed approximately 250MWh on an in-house NVIDIA V100 cluster, and improved sample quality may amplify disinformation, stereotypes, and harmful biases.These concerns are presented in the societal-impact discussion.

A Additional results

Additional experiments compare the proposed samplers and training configurations with earlier diffusion models across ImageNet-64, CIFAR-10, FFHQ, and AFHQv2. The evaluations emphasize FID as a function of sampling cost and use controlled comparisons across configurations.

  • ImageNet-64: Pre-trained ADM ImageNet-64 samples are compared using DDIM, iDDPM, and the proposed samplers in deterministic and stochastic settings.Separate figures show corresponding results after training from scratch with the improved configuration.
  • Additional datasets: Earlier Song et al. samplers and training configurations are compared with the proposed approach across unconditional and class-conditional CIFAR-10, FFHQ, and AFHQv2.The same latent codes are reused across training configurations and ODE choices.
  • Quantitative evaluation: Tables 3 and 4 summarize deterministic and stochastic sampling results across datasets using the NFE comparisons shown in Figures 2 and 4.NFE measures the number of neural function evaluations used for sampling.

B Derivation of formulas

The appendix reframes diffusion models around marginal distributions and derives deterministic and stochastic formulations, then analyzes sampling, training, and practical implementation details. These formulations support modular sampler comparisons and score-matching derivations, while exposing implementation issues in prior samplers.

  • B.1 Original ODE / SDE formulation from previous work: Song et al.’s SDE uses drift and diffusion coefficients to define variance-preserving and variance-exploding perturbation processes.The associated marginal distribution is obtained by integrating perturbation kernels over the data distribution.
  • B.1 Original ODE / SDE formulation from previous work: The probability-flow ODE preserves the same time-dependent marginal distribution while using only the initial noise image as randomness.The original formulation is expressed through f and g, which correspond to terms in the ODE but are less practically central than the marginals.
  • B.2 Our ODE formulation (Eq. 1 and Eq. 4): The reformulation defines the ODE directly from the marginal-distribution parameters σ(t) and s(t), eliminating the need for f(t) and g(t).Changing σ(t) reparameterizes time, while changing s(t) reparameterizes the state variable x.
  • B.2 Our ODE formulation (Eq. 1 and Eq. 4): The canonical ODE interpretation shows that different realizations are reparameterizations of one underlying probability-flow trajectory.Setting s(t) = 1 recovers Eq. 1, while Eq. 4 retains the additional signal-scaling schedule.
  • B.3 Denoising score matching (Eq. 2 and Eq. 3): Denoising score matching reduces to independent convex optimization at each noisy sample, yielding a uniquely identified closed-form ideal denoiser.The closed-form expression can be computed in practice for small datasets, as demonstrated for CIFAR-10.
  • B.4 Evaluating our ODE in practice (Algorithm 1): The practical ODE algorithm substitutes the trained denoiser Dθ for the ideal denoiser after rewriting the score under time-dependent signal scaling.The resulting expressions correspond to lines 4 and 7 of Algorithm 1.
  • B.5 Our SDE formulation (Eq. 6): The stochastic formulation combines score-based denoising with noise injection, scaling Langevin exploration according to the current noise standard deviation.This baseline treats the data manifold as effectively spread out by the Gaussian blurring at the current noise level.

C.3.3 iDDPM preconditioning and training

The iDDPM checkpoint is restricted to a discrete set of 1000 noise levels, requiring adaptations when using the paper’s samplers. The appendix describes resampling, snapping stochastic times to supported levels, and reusing compatible training definitions.

  • C.3.3 iDDPM preconditioning and training: The iDDPM preconditioning formula matches the VP preconditioning formula, allowing the appendix to reuse the corresponding training loss with σ = u_j.The secondary L_vlb term is not incorporated into this reuse statement.
  • C.3.3 iDDPM preconditioning and training: The pre-trained ADM checkpoint contains 296 million parameters and supports only M = 1000 discrete noise levels.This restriction creates practical challenges for evaluating samplers at arbitrary time grids.
  • C.3.3 iDDPM preconditioning and training: DDIM requires resampling the discrete noise levels to construct a shorter time grid when N ≠ M.The appendix contrasts a divisibility-constrained resampling rule with a more flexible index-based scheme.
  • C.3.3 iDDPM preconditioning and training: The stochastic sampler snaps each proposed intermediate time to the nearest supported iDDPM noise level.This is implemented by selecting the u_j minimizing the distance to the proposed time.
  • C.3.3 iDDPM preconditioning and training: The imported checkpoint can run with Algorithms 1 and 2 after these adaptations, using only its predicted noise output and ignoring its variance output.The network is imported directly as Fθ(·) under the definitions in Table 1.

D.1 Truncation error analysis and choice of discretization parameters

The analysis links sampling-step accuracy to the ODE solver and noise-level discretization. Heun’s method and a polynomial schedule with ρ = 7 provide the preferred practical tradeoff for image generation.

  • Truncation error: Numerical ODE solvers approximate the diffusion trajectory, so larger steps accumulate local errors into global truncation error.Reducing global error requires minimizing the largest local truncation error through the solver and sampling times.
  • Euler analysis: Euler’s local RMSE reaches approximately 0.56 at noise levels σi ≤ 0.5 under uniform σ-step discretization.The observed error is much smaller at high noise levels, motivating smaller steps as σ decreases.
  • Heun’s method: Heun’s correction accounts for changes in dx/dt between time points and yields cubic local-error convergence at one additional Dθ evaluation per step.Its error curves are generally lower than Euler’s under the same setup.
  • Runge–Kutta variants: The general second-order Runge–Kutta family varies the evaluation point through α, with α = 1 recovering Heun’s method.The family has the same computational cost, but its variants incur different approximation errors.
  • Runge–Kutta variants: Although α = 1.1 performed slightly better experimentally, the authors fixed α = 1 because overshooting the target is theoretically difficult to justify.They leave varying α during sampling for future work.

E.1 Image degradation due to excessive stochastic iteration

Repeated stochastic iteration at a fixed noise level degrades generated images, with failure modes depending on noise level. The sampler therefore uses noise-dependent corrections to limit low-noise iteration and compensate for high-noise degradation.

  • Image degradation: Repeated Langevin iteration damages image distributions despite theoretically moving them toward p(x; σ) under conservative denoising fields.The observed degradation depends on the dataset and noise level.
  • Low-noise behavior: At σ below roughly 0.2, images oversaturate after about 2k iterations and become fully corrupted thereafter.Setting Stmin > 0 prevents stochastic sampling at very low noise levels.
  • High-noise behavior: At high noise levels, omitting standard-deviation correction makes images increasingly abstract and colorless at high iteration counts.Setting Snoise > 1 counteracts this tendency efficiently.
  • Sampler ablations: For fixed NFE, stochastic sampling generally has an optimal Schurn, after which FID degrades.The best settings depend on the model and case, and were selected by grid search.
  • Sampler ablations: The best results across the tested cases use all stochastic corrections, although Snoise versus Stmin,tmax matters more depending on the case.The comparison is reported for pre-trained networks from Song et al. and Dhariwal and Nichol.

F Implementation details

The implementation combines a unified, extensively validated codebase with FID evaluation on 50,000 generated images. Training uses conditional geometric augmentation to regularize score networks without leaking transformations into sampled outputs.

  • Implementation: The authors rewrote the implementation and tested samplers, pre-trained models, architectures, training configurations, and evaluation against prior results.The implementation used PyTorch 1.10.0, CUDA 11.4, and CuDNN 8.2.0 on NVIDIA DGX-1 systems.
  • Evaluation: FID is computed between 50,000 generated images and all available real images, with NFE directly representing the dominant network-evaluation cost of sampling.The evaluation uses the StyleGAN3-provided Inception-v3 model without augmentation such as x-flips.
  • Augmentation: The augmentation pipeline applies six geometric transformations before noise addition and conditions the network on their parameters.Each training image receives independently sampled transformations, with fixed activation probabilities except for always-enabled x-flips.
  • Augmentation: Conditioning on augmentation parameters creates auxiliary denoising tasks that provide varied training samples and prevent overfitting to individual examples.The authors speculate that the auxiliary tasks also help the main task because denoising is similar across augmentation choices.
  • Augmentation: Sampling sets the augmentation condition to zero, and the authors report no leakage of geometric transformations into generated images.They therefore allow the augmentation constants to be chosen freely within the tested setup.

F.3 Training configurations

The training configurations progressively modify optimization, capacity, preconditioning, noise and loss design, and augmentation while retaining established network families. ImageNet-64 additionally uses substantially larger-scale hardware and mixed-precision training.

  • Configurations: The baseline reproduces Song et al.’s VP and VE setups, while the final configuration combines the paper’s improved training components.Each model is trained on 200 million images, or approximately 400,000 iterations with batch size 512.
  • Configurations: The faster comparison configuration increases GPU parallelism and batch size, disables gradient clipping, and raises the CIFAR-10 learning rate to 0.001.The learning rate is ramped during the first 10 million images.
  • Network capacity: The capacity configuration removes 4×4 layers and doubles 16×16-layer capacity to improve expressiveness while reducing overfitting from the removed layers.The authors identify 16×16 capacity as critical for high-quality results.
  • Training design: The final configuration replaces original preconditioning, noise distribution, and loss weighting, then enables augmentation regularization.VP and VE retain DDPM++ and NCSN++ architectures respectively, differing mainly in architecture rather than the improved training setup.
  • ImageNet-64: ImageNet-64 training uses 32 NVIDIA Ampere GPUs, batch size 4096, and mixed-precision FP16/FP32 computation.Trainable parameters remain FP32 but are cast to FP16 during Fθ evaluation except in embedding and self-attention layers.
  • Conditioning: Class conditioning and augmentation conditioning enter the DDPM++ and NCSN++ networks as additional inputs combined with the noise-level embedding.Class labels are represented as scaled one-hot vectors, while augmentation parameters are mapped through fully connected layers.
  • ImageNet-64: Class-conditional ImageNet-64 uses the unchanged ADM architecture with approximately 296 million trainable parameters.Compared with DDPM++, it is shallower but uses more channels and self-attention layers.
Loading 2206.00364v2…