EDM — Elucidating Diffusion Models

One axis for everything: the sigma-parameterization

EDM fixes the forward process as the simplest thing that could work: a clean sample y is corrupted by additive Gaussian noise of a chosen scale, x = y + n with n ∼ N(0, σ^2 I). That is it — there is no variance-preserving rescaling of y itself. The marginal at noise level σ is p(x; σ) = ∫ p_data(y) N(x; y, σ^2 I) dy, a Gaussian-blurred version of the data distribution whose blur radius is σ.

This single scalar replaces the timestep t, the β schedule, and α-bar all at once. When σ → 0 the distribution is the data; when σ is large (EDM uses σ_max = 80) it is indistinguishable from pure noise of that scale. Generation is then simply the task of walking σ down from large to near-zero. Because σ carries physical units of ‘standard deviations of added noise,’ every later choice — schedule, weighting, network scaling — can be reasoned about dimensionally instead of by convention.

The denoiser and the probability-flow ODE

The object EDM trains is a denoiser D_θ(x; σ) that, given a noisy x and its noise level, returns an estimate of the clean signal y. The score of the noised marginal follows directly from Tweedie’s formula: ∇_x log p(x; σ) = (D(x; σ) − x) / σ^2. A denoiser and a score are two views of the same function.

Sampling uses the probability-flow ODE, the deterministic trajectory whose marginals match the diffusion at every level. With σ playing the role of time it reduces to a strikingly plain form:

dx/dσ = (x − D(x; σ)) / σ
         = −σ · ∇_x log p(x; σ)

The right-hand side points from the noisy point toward the denoiser’s estimate of clean data, scaled by 1/σ. Integrating this field from σ_max down to σ_min transports a Gaussian sample onto the data manifold. No stochasticity is required — that becomes an optional design knob rather than a defining feature.

Preconditioning: why the network should not predict D directly

The subtle EDM contribution is that you should not ask a neural network F_θ to output D raw. At small σ the clean signal dominates and the network mostly needs to copy its input; at large σ the input is almost pure noise and the network must predict nearly the whole signal. Forcing one network to span that range with fixed input/output scales wrecks conditioning. EDM wraps F_θ in four σ-dependent coefficients:

D(x; σ) = c_skip(σ) · x  +  c_out(σ) · F_θ( c_in(σ) · x ; c_noise(σ) )

Read it left to right: c_in rescales the noisy input to unit variance before it ever reaches the network; c_noise maps σ to a well-behaved conditioning input; the network predicts a unit-scale residual F_θ; c_out rescales that prediction; and c_skip mixes in a straight copy of the input so the network only has to learn the correction. Every coefficient exists to keep the quantities the network sees and produces at order one.

Deriving c_in, c_skip, c_out, c_noise

The coefficients are not tuned — they fall out of two demands: the network’s input should have unit variance, and its training target should have unit variance, for every σ. Let σ_data be the standard deviation of the clean data (about 0.5 for images in [−1, 1]). Since x = y + n has variance σ_data^2 + σ^2, unit input variance forces

c_in(σ)   = 1 / √(σ^2 + σ_data^2)
c_skip(σ) = σ_data^2 / (σ^2 + σ_data^2)
c_out(σ)  = σ · σ_data / √(σ^2 + σ_data^2)
c_noise(σ) = ¼ · ln(σ)

c_skip is chosen to minimize the magnitude of what F_θ must predict, and c_out then normalizes that residual to unit variance. Check the limits: as σ → 0, c_skip → 1 and c_out → 0, so D → x — a clean input is returned untouched. As σ grows, c_skip → 0 and the skip path fades, handing the whole job to the network. The ln(σ)/4 for c_noise is empirical: a log scale spreads the enormous σ range into a tractable conditioning signal.

Loss weighting that makes every noise level count equally

Training minimizes a denoising objective over noise levels: E[ λ(σ) · || D(y + n; σ) − y ||^2 ]. Left alone, the raw loss at each σ has wildly different magnitude — huge where the target is large, tiny where it is small — so gradient signal is dominated by a narrow band of levels. EDM cancels this by choosing the weight to undo the preconditioning:

λ(σ) = 1 / c_out(σ)^2 = (σ^2 + σ_data^2) / (σ · σ_data)^2

With this weight the effective loss the network sees — the error on its own unit-scale output F_θ — has constant magnitude one across all σ. Every noise level then contributes comparable gradient, and no region silently starves. This is the same principle as the preconditioning, applied to the loss rather than the activations: keep the numbers the optimizer actually differentiates near unity so no part of the σ spectrum is over- or under-trained by accident.

The training noise distribution

A weight decides how much each level counts; a sampling distribution decides how often each level is seen. EDM draws training noise levels log-normally: ln(σ) ∼ N(P_mean, P_std^2) with P_mean = −1.2 and P_std = 1.2. In linear terms the median σ is e^(−1.2) ≈ 0.30, with a long tail toward larger values.

The reasoning is that both extremes are nearly useless to train on. At very small σ the denoising task is trivial (copy the input); at very large σ it is hopeless (the signal is gone) and the target is roughly the data mean regardless. The useful learning happens in the middle band where structure is partly visible, and the log-normal concentrates training exactly there while still occasionally sampling the tails so the model stays calibrated everywhere the sampler will visit. This is one of the largest practical levers in the paper: it moves sample quality more than most architecture changes.

The sampling schedule: a curved march in sigma

Sampling discretizes the ODE into N steps, and where you place the σ_i matters as much as how many you use. EDM does not space them uniformly; it warps them with a power law:

σ_i = ( σ_max^(1/&rho) + (i / (N−1)) · ( σ_min^(1/&rho) − σ_max^(1/&rho) ) )^&rho
for i = 0 .. N−1,   with σ_min = 0.002,  σ_max = 80,  &rho = 7

Uniform spacing in σ^(1/&rho) with &rho = 7 produces steps that are coarse at high noise and fine near σ = 0. The intuition is that the ODE trajectory is nearly straight when the noise is large — a big Euler step there costs little — but curves sharply as it approaches the data manifold, where discretization error is most damaging. Concentrating steps at small σ spends the step budget where curvature actually lives. The final step lands at σ = 0 exactly, denoising to a clean sample.

Sampling curvature and truncation error

Why does the schedule shape move quality at all? Because a numerical ODE solver accumulates local truncation error at every step, and that error scales with the curvature of the trajectory times the step size. EDM analyzes the probability-flow field and observes that with the σ-time parameterization the paths are unusually close to straight lines over most of their length — the tangent (x − D)/σ changes slowly while σ is large.

That near-linearity is exactly why an explicit solver can take a handful of big steps up top and still track the true path. It also motivates the &rho = 7 curved schedule as the discretization that equalizes error per step rather than error per unit σ. The payoff is concrete: EDM reaches state-of-the-art FID on CIFAR-10 in the neighborhood of 35 network evaluations, where earlier DDPM samplers needed hundreds or thousands. The structure was always in the ODE; EDM just chose coordinates that expose it.

Heun's second-order sampler

The final accelerator is the integrator itself. A plain Euler step uses the derivative at the current point only, so its error per step is O(Δσ^2). EDM instead uses Heun’s method, a second-order predictor-corrector that evaluates the field twice and averages, cutting error to O(Δσ^3) per step for one extra denoiser call:

d_i    = (x_i − D(x_i; σ_i)) / σ_i          # slope at current level
x'     = x_i + (σ_{i+1} − σ_i) · d_i        # Euler predictor
d_i'   = (x' − D(x'; σ_{i+1})) / σ_{i+1}     # slope at predicted point
x_{i+1} = x_i + (σ_{i+1} − σ_i) · ½ (d_i + d_i')   # corrected step

The correction cancels the first-order term that a single Euler step leaves behind, so each step tracks the curved trajectory far better. EDM applies the correction on every step except the last one into σ = 0, where d_i' would divide by zero. Two evaluations per step at ~18 steps beats one evaluation per step at ~50 for equal quality — second-order accuracy is cheaper than more first-order steps.

Optional stochasticity: churn as error correction

The deterministic ODE is the default, but EDM also gives a principled stochastic sampler and, crucially, explains what the randomness buys. Each step optionally injects a little fresh noise — nudging σ back up by a factor 1 + γ before stepping down — controlled by parameters S_churn, S_min, S_max, and S_noise.

The insight is that this churn acts as an error-correcting mechanism: adding noise and re-denoising pulls the trajectory back toward the true marginal, counteracting the drift that a biased or under-trained denoiser introduces. So stochasticity is not a source of diversity per se — the ODE already reproduces the full distribution — it is a numerical correction that helps some models but can also blur fine detail if overdone. EDM’s framing lets you dial it deliberately: enable churn only in a mid-σ band, keep S_noise slightly above one, and disable it entirely for a well-trained network where the deterministic path is already faithful.

Why this matters for small models on modest hardware

For the CPU-SLM reader the relevant lesson is that EDM buys quality with math, not parameters. The preconditioning and loss weighting cost nothing at inference — they are closed-form scalars — yet they make a small network train stably across the whole noise range, which is precisely where an under-parameterized model would otherwise fall apart. A compact denoiser paired with EDM’s recipe punches well above a larger model trained with a naive objective.

The sampler side is even more directly relevant when compute is scarce. Sample cost is dominated by the number of denoiser evaluations, and EDM’s curved schedule plus Heun steps cut that count by an order of magnitude versus vanilla DDPM. Fewer forward passes means a diffusion model becomes viable on a CPU or a phone, where hundreds of steps would be a non-starter. When you must ship a small generative model on modest hardware, the EDM design space is where the affordable wins live.

Common pitfalls when applying EDM

The first trap is mismatching σ_data. The coefficients assume your data really has that standard deviation; if you change normalization — latents from a VAE, audio, a different pixel range — and leave σ_data = 0.5, the unit-variance guarantees break and training conditioning degrades quietly. Measure it on your actual data.

The second is treating the constants as universal. P_mean, P_std, &rho, and the σ bounds were tuned for natural images; a very different modality often wants a different median noise level and schedule curvature. The third is porting only the sampler. People grab Heun and the curved schedule but keep a DDPM-style ε-prediction network without the preconditioning — and then wonder why quality stalls. EDM’s pieces are a coherent system: the network parameterization, the loss weighting, the training distribution, and the sampler were co-designed, and they deliver most of their gain together.