Elucidating the Design Space of Diffusion-Based Generative Models

Tero Karras, Miika Aittala, Timo Aila, Samuli Laine

2022 · NeurIPS

Elucidating the Design Space of Diffusion-Based Generative Models

Problem

Framing

Diffusion models mixed solver choice, noise schedule, stochasticity, and network parameterization into monolithic recipes. This paper factorizes that design space, then replaces each weak choice with a better one: Heun-based sampling, tuned noise schedules, and score-network preconditioning. It reaches FID 1.79 on class-conditional CIFAR-10 and 1.36 on ImageNet-64 with 35 NFE.

Currently Used Methods

Foundational

@DenoisingDiffusionProbabilisticModels2020 — discrete-time diffusion with weighted variational training.
- Limitation in context: fixed design bundle hides sampler and preconditioning tradeoffs.
@songScoreSDE2020 — score-SDE framework with predictor-corrector and probability-flow ODE.
- Limitation in context: default discretization and schedules are not sampling-efficient.
@DenoisingDiffusionImplicitModels2020 — deterministic DDIM trajectories for faster generation.
- Limitation in context: Euler discretization leaves quality on the table at low NFE.
@nicholImprovedDDPM2021 — improved DDPM with learned variances and cosine schedule.
- Limitation in context: training and sampling choices remain tightly coupled.
@dhariwalDiffusionBeatGANs2021 — ADM architecture that pushes diffusion to GAN-level fidelity.
- Limitation in context: architecture improves quality, but solver design stays underexplored.

Proposed Method

Architecture

EDM keeps the backbone family unchanged: DDPM++ for VP, NCSN++ for VE, and ADM for ImageNet. The main architectural change is external preconditioning: the denoiser is expressed as a skip-connected wrapper around a noise-conditioned backbone $F_\theta$ .

D_\theta(\mathbf{x}; \sigma) = c_{\mathrm{skip}}(\sigma)\,\mathbf{x} + c_{\mathrm{out}}(\sigma)\,F_\theta\!\left(c_{\mathrm{in}}(\sigma)\,\mathbf{x}; c_{\mathrm{noise}}(\sigma)\right)

Loss / Objective

The objective trains the preconditioned denoiser over noise levels sampled from a log-normal distribution.

\mathcal{L}(\theta) = \mathbb{E}_{\mathbf{y},\mathbf{n},\sigma}\left[\lambda(\sigma)\,\left\| D_\theta(\mathbf{y}+\mathbf{n};\sigma)-\mathbf{y} \right\|^2\right],\quad \mathbf{y}\sim p_{\mathrm{data}},\; \mathbf{n}\sim \mathcal{N}(\mathbf{0},\sigma^2\mathbf{I})

c_{\mathrm{skip}}(\sigma)=\frac{\sigma_{\mathrm{data}}^2}{\sigma^2+\sigma_{\mathrm{data}}^2},\quad c_{\mathrm{out}}(\sigma)=\frac{\sigma\,\sigma_{\mathrm{data}}}{\sqrt{\sigma^2+\sigma_{\mathrm{data}}^2}},\quad c_{\mathrm{in}}(\sigma)=\frac{1}{\sqrt{\sigma^2+\sigma_{\mathrm{data}}^2}},\quad \lambda(\sigma)=\frac{1}{c_{\mathrm{out}}(\sigma)^2}=\frac{\sigma^2+\sigma_{\mathrm{data}}^2}{\sigma^2\sigma_{\mathrm{data}}^2}

Sampling Rule / Algorithm

For deterministic sampling, EDM integrates the probability-flow ODE with Heun's second-order method on a noise schedule $\{\sigma_i\}$ .

\dot{\mathbf{x}} = \frac{\dot{\sigma}(t)}{\sigma(t)}\left(\mathbf{x}-D_\theta(\mathbf{x};\sigma(t))\right)

\hat{\mathbf{x}}_{i+1}=\mathbf{x}_i+(t_{i+1}-t_i)\,d_i, \qquad d_i=\frac{\mathbf{x}_i-D_\theta(\mathbf{x}_i;t_i)}{t_i},

d'_i=\frac{\hat{\mathbf{x}}_{i+1}-D_\theta(\hat{\mathbf{x}}_{i+1};t_{i+1})}{t_{i+1}}, \qquad \mathbf{x}_{i+1}=\mathbf{x}_i+(t_{i+1}-t_i)\left(\frac{d_i+d'_i}{2}\right)

EDM then adds optional stochastic "churn" between selected noise levels to improve low-NFE sampling.

Training Procedure

$\sigma_{\mathrm{data}} = 0.5$
$\sigma_{\min} = 0.002$ , $\sigma_{\max} = 80$
Noise schedule exponent $\rho = 7$
Log-normal noise sampling: $P_{\mathrm{mean}} = -1.2$ , $P_{\mathrm{std}} = 1.2$
CIFAR-10 training duration: 200 M images
CIFAR-10 batch size: 512
CIFAR-10 learning rate: 0.001
EMA half-life: 0.5 Mimages
Data augmentation probability: 0.12

Evaluation

Datasets

CIFAR-10, 32 $\times$ 32
FFHQ, 64 $\times$ 64
AFHQv2, 64 $\times$ 64
ImageNet-64, class-conditional

Metrics

FID
NFE for near-best FID
Qualitative sample quality

Headline results

CIFAR-10 class-conditional: FID 1.79.
CIFAR-10 unconditional: FID 1.97.
ImageNet-64 pre-trained model + EDM sampler: FID 1.55.
ImageNet-64 retrained with EDM: FID 1.36.
CIFAR-10/ImageNet-64 sampling: near-SOTA quality at 35 NFE.

Results plot: three FID-vs-NFE panels compare deterministic, stochastic, and prior samplers on unconditional CIFAR-10 and class-conditional ImageNet-64.

Ablations

Deterministic solver: Heun beats Euler and black-box RK45 at matched NFE.
Time-step schedule: larger $\rho$ concentrates steps near low noise and improves FID.
Stochasticity: tuned churn improves low-NFE FID over pure ODE sampling.
Training configuration: preconditioning, noise sampling, and augmentation each reduce FID.

Method Strengths and Weaknesses

Strengths

Separates diffusion design axes cleanly: solver, schedule, stochasticity, and preconditioning.
Heun sampler cuts NFE sharply without sacrificing FID.
Preconditioning yields stable loss scaling across noise levels.
Works as a drop-in sampler improvement for pre-trained models.

Weaknesses

Core backbones are inherited, not architecturally novel.
Strongest gains rely on careful hand-tuned sampler hyperparameters.
Main image results remain at low resolution: 32 $\times$ 32 and 64 $\times$ 64.
Design-space study is empirical; optimality is not theoretically settled.

Suggestions from the authors

Explore wider combinations of modular diffusion design choices.
Analyze why low-noise discretization errors dominate sample quality.
Study whether overshooting variants of Heun improve sampling more broadly.
Extend the framework beyond current score-network parameterizations.

Elucidating the Design Space of Diffusion-Based Generative Models

Elucidating the Design Space of Diffusion-Based Generative Models

Problem

Framing

Currently Used Methods

Foundational

Proposed Method

Architecture

Loss / Objective

Sampling Rule / Algorithm

Training Procedure

Evaluation

Datasets

Metrics

Headline results

Ablations

Method Strengths and Weaknesses

Strengths

Weaknesses

Suggestions from the authors

Links

Prior Papers

Further Papers