Diffusion models

Deep dive

How it works

In: samples x₀ (training) or N(0, I) noise (sampling), optional condition · Out: a noise predictor; samples · State: weights, EMA copy, schedule · Cost: NFE × network cost per sample.

DDPM (ε-prediction). Schedule β₁..β_T, α_t = 1 − β_t, ᾱ_t = Π α_s. The forward process has a closed form:

x_t = √ᾱ_t · x₀ + √(1−ᾱ_t) · ε,      ε ~ N(0, I),     SNR(t) = ᾱ_t / (1−ᾱ_t)

1. Training: sample x₀, t ~ U{1..T}, ε; minimize ‖ε_θ(x_t, t) − ε‖². T = 1000, linear β from 1e-4 to 0.02 gives ᾱ_T ≈ 4.0e-5. The cosine schedule ᾱ(t) = cos²((t/T + s)/(1+s)·π/2)/ᾱ(0), s = 0.008, β clipped at 0.999, spends fewer steps at near-pure noise.
2. Ancestral sampling: x_{t−1} = (x_t − β_t/√(1−ᾱ_t)·ε_θ)/√α_t + σ_t·z, with σ_t² = β_t or β̃_t = β_t(1−ᾱ_{t−1})/(1−ᾱ_t). Needs about all T steps.
3. DDIM (η = 0): predict x̂₀ = (x_t − √(1−ᾱ_t)·ε_θ)/√ᾱ_t, then jump to any earlier level: x_s = √ᾱ_s·x̂₀ + √(1−ᾱ_s)·ε_θ. Deterministic, 25-50 steps, and the same trained network works unchanged.

Other parametrizations of the same model: x₀-prediction, and v-prediction v = √ᾱ·ε − √(1−ᾱ)·x₀ (stable at both ends of the schedule, needed for zero terminal SNR). Loss weighting matters: Min-SNR-γ weights the ε-loss by min(SNR, γ)/SNR with γ = 5.

Rectified flow / flow matching. Straight path x_t = (1−t)·x₀ + t·ε, t ∈ [0, 1]; target velocity v = ε − x₀; loss ‖v_θ(x_t, t) − v‖². Sample by integrating dx/dt = v_θ from t = 1 to 0 with Euler or Heun. SD3 samples t from a logit-normal (t = sigmoid(n), n ~ N(0, 1)) to concentrate training on mid noise levels, and shifts timesteps toward noise at high resolution: t' = α·t/(1 + (α−1)·t).

Conditioning and guidance. Timestep → sinusoidal features → MLP. Text → frozen encoder (CLIP, T5) → cross-attention or joint attention. Classifier-free guidance evaluates conditional and unconditional outputs (one batch of 2) and extrapolates with scale w (diffusers' SD 1.x default: 7.5 at 50 steps).

Backbones.

• U-Net (DDPM, ADM, SD 1.x/2.x/SDXL): ResBlocks with GroupNorm, timestep embedding added or applied as scale/shift, downsample ×2 per stage, self-attention plus text cross-attention at the low-resolution stages, skip connections concatenated into the decoder. SD 1.x U-Net: about 860M params.
• DiT (Peebles & Xie): patchify the latent (p = 2), plain transformer blocks, condition via adaLN-Zero (timestep/class embedding regresses per-block scale, shift and a gate initialized to zero, so each block starts as identity). DiT-XL/2: 28 blocks, width 1152, 675M params, 118.6 GFLOPs per forward at 256², FID 2.27 on ImageNet 256 with guidance. SD3 and FLUX use transformer backbones of this family with joint text-image attention.

Latent space. A VAE encoder compresses images 8× per side (SD 1.x: 4 channels, latents scaled by 0.18215 to roughly unit variance; SD3: 16 channels). Diffusion runs on latents; the decoder runs once at the end.

Code

Schedules, the DDPM training loss (with optional Min-SNR), DDIM sampling, the rectified-flow loss and shifted Euler sampler, and CFG. The test uses a one-point data distribution whose optimal predictors are known in closed form, so every sampler must return that point.

// DDPM (epsilon-prediction) and rectified-flow objectives and samplers on flat f32 slices.
// `model` stands for the network: (x_t, timestep) -> prediction with x_t's shape.
pub struct Rng(pub u64);
impl Rng {
    pub fn uniform(&mut self) -> f32 { // xorshift64*, 24-bit mantissa in (0, 1)
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        ((self.0.wrapping_mul(0x2545F4914F6CDD1D) >> 40) as f32 + 0.5) / (1u64 << 24) as f32
    }
    pub fn normal(&mut self) -> f32 { // Box-Muller
        let (u, v) = (self.uniform(), self.uniform());
        (-2.0 * u.ln()).sqrt() * (std::f32::consts::TAU * v).cos()
    }
}

pub struct Schedule { pub beta: Vec<f32>, pub abar: Vec<f32> } // abar[i] = prod_{j<=i} (1 - beta[j])
impl Schedule {
    pub fn linear(n: usize, b0: f32, b1: f32) -> Self { // DDPM: n = 1000, 1e-4 .. 0.02
        Self::from_beta((0..n).map(|i| b0 + (b1 - b0) * i as f32 / (n - 1) as f32).collect())
    }
    pub fn cosine(n: usize) -> Self { // Nichol & Dhariwal: s = 0.008, beta clipped at 0.999
        let f = |t: f64| (((t / n as f64 + 0.008) / 1.008) * std::f64::consts::FRAC_PI_2).cos().powi(2);
        Self::from_beta((0..n).map(|i| (1.0 - f(i as f64 + 1.0) / f(i as f64)).min(0.999) as f32).collect())
    }
    fn from_beta(beta: Vec<f32>) -> Self {
        let mut a = 1.0f64; // accumulate the product in f64
        let abar = beta.iter().map(|b| { a *= 1.0 - *b as f64; a as f32 }).collect();
        Schedule { beta, abar }
    }
}

/// One training sample: x_t = sqrt(abar) x0 + sqrt(1-abar) eps, loss = w * mean (eps_hat - eps)^2.
/// Min-SNR-gamma weighting w = min(SNR, gamma) / SNR when gamma is set.
pub fn ddpm_loss(s: &Schedule, x0: &[f32], rng: &mut Rng, gamma: Option<f32>,
                 model: &dyn Fn(&[f32], usize) -> Vec<f32>) -> f32 {
    let i = ((rng.uniform() * s.abar.len() as f32) as usize).min(s.abar.len() - 1);
    let (a, b) = (s.abar[i].sqrt(), (1.0 - s.abar[i]).sqrt());
    let eps: Vec<f32> = x0.iter().map(|_| rng.normal()).collect();
    let xt: Vec<f32> = x0.iter().zip(&eps).map(|(x, e)| a * x + b * e).collect();
    let pred = model(&xt, i);
    let snr = s.abar[i] / (1.0 - s.abar[i]);
    let w = gamma.map_or(1.0, |g| snr.min(g) / snr);
    w * pred.iter().zip(&eps).map(|(p, e)| (p - e) * (p - e)).sum::<f32>() / x0.len() as f32
}

/// Deterministic DDIM (eta = 0) over a subsequence of timesteps, from noise x to a sample.
pub fn ddim_sample(s: &Schedule, steps: usize, mut x: Vec<f32>, model: &dyn Fn(&[f32], usize) -> Vec<f32>) -> Vec<f32> {
    let n = s.abar.len();
    let ts: Vec<usize> = (0..steps).rev().map(|k| k * n / steps).collect(); // e.g. 980, 960, ..., 0
    for (k, &i) in ts.iter().enumerate() {
        let a_prev = if k + 1 < ts.len() { s.abar[ts[k + 1]] } else { 1.0 }; // final step lands on x0
        let eps = model(&x, i);
        let (sa, sb) = (s.abar[i].sqrt(), (1.0 - s.abar[i]).sqrt());
        for (xv, e) in x.iter_mut().zip(&eps) {
            let x0 = (*xv - sb * e) / sa;
            *xv = a_prev.sqrt() * x0 + (1.0 - a_prev).sqrt() * e;
        }
    }
    x
}

/// Rectified flow: x_t = (1-t) x0 + t eps, target velocity v = eps - x0, t ~ logit-normal(0, 1).
pub fn flow_loss(x0: &[f32], rng: &mut Rng, model: &dyn Fn(&[f32], f32) -> Vec<f32>) -> f32 {
    let t = 1.0 / (1.0 + (-rng.normal()).exp());
    let eps: Vec<f32> = x0.iter().map(|_| rng.normal()).collect();
    let xt: Vec<f32> = x0.iter().zip(&eps).map(|(x, e)| (1.0 - t) * x + t * e).collect();
    let v = model(&xt, t);
    v.iter().zip(x0.iter().zip(&eps)).map(|(v, (x, e))| (v - (e - x)).powi(2)).sum::<f32>() / x0.len() as f32
}

/// Euler ODE solve from t = 1 (noise) to t = 0, with the SD3 resolution shift t' = a t / (1 + (a-1) t).
pub fn flow_sample(steps: usize, shift: f32, mut x: Vec<f32>, model: &dyn Fn(&[f32], f32) -> Vec<f32>) -> Vec<f32> {
    let tt = |k: usize| { let t = k as f32 / steps as f32; shift * t / (1.0 + (shift - 1.0) * t) };
    for k in (1..=steps).rev() {
        let (t, t_next) = (tt(k), tt(k - 1));
        let v = model(&x, t);
        for (xv, vv) in x.iter_mut().zip(&v) { *xv -= (t - t_next) * vv; }
    }
    x
}

/// Classifier-free guidance on two model outputs (epsilon or velocity): u + w (c - u).
pub fn cfg(cond: &[f32], uncond: &[f32], w: f32) -> Vec<f32> {
    cond.iter().zip(uncond).map(|(c, u)| u + w * (c - u)).collect()
}

fn main() {
    // With data = one point c, the optimal predictors are known in closed form: samplers must return c.
    let c = vec![0.5f32, -1.0, 2.0];
    let s = Schedule::linear(1000, 1e-4, 0.02);
    assert!((s.abar[999] - 4.04e-5).abs() < 1e-6);
    let eps_star = |x: &[f32], i: usize| -> Vec<f32> {
        x.iter().zip(&c).map(|(x, c)| (x - s.abar[i].sqrt() * c) / (1.0 - s.abar[i]).sqrt()).collect()
    };
    let mut rng = Rng(42);
    assert!(ddpm_loss(&s, &c, &mut rng, Some(5.0), &eps_star) < 1e-6);
    let noise: Vec<f32> = (0..3).map(|_| rng.normal()).collect();
    let cs = Schedule::cosine(1000);
    let out = ddim_sample(&cs, 50, noise.clone(), &|x, i| {
        let a = cs.abar[i];
        x.iter().zip(&c).map(|(x, c)| (x - a.sqrt() * c) / (1.0 - a).sqrt()).collect()
    });
    assert!(out.iter().zip(&c).all(|(o, c)| (o - c).abs() < 1e-3));
    let v_star = |x: &[f32], t: f32| -> Vec<f32> { x.iter().zip(&c).map(|(x, c)| (x - c) / t).collect() };
    assert!(flow_loss(&c, &mut rng, &v_star) < 1e-6);
    let out = flow_sample(28, 3.0, noise, &v_star);
    assert!(out.iter().zip(&c).all(|(o, c)| (o - c).abs() < 1e-4));
    assert_eq!(cfg(&[2.0], &[1.0], 7.5), vec![8.5]);
    println!("ok");
}

Complexity and constants

• Training: one forward + backward per sample at one random t; order of 10⁵-10⁶ steps at batch 256-2048 (DiT-XL/2: 7M steps at batch 256).
• Sampling: NFE × forward cost, ×2 with CFG (numbers under Efficiency).
• DiT cost scales with token count: patch size 2 on a 32×32 latent gives 256 tokens; halving the patch quadruples tokens and roughly quadruples FLOPs.

Efficiency

• Cost: DiT-XL/2 takes 118.6 GFLOPs per forward at 256²; a sample costs NFE × that, ×2 with CFG. Bound: GPU compute.
• Levers:
  • latent space — 48× fewer values than pixels (Rombach et al.) — images, video
  • better solvers — 15-25 NFE instead of 1000 (DPM-Solver++, DDIM) — any trained model
  • distillation — 1-4 steps (progressive distillation) — real-time use
  • Min-SNR loss weighting — 3.4x faster convergence (Hang et al.) — training
• Naive → standard → state of the art:
  • ancestral DDPM · about 1000 calls
  • DDIM, DPM-Solver++ · 20-50 NFE
  • EDM Heun · 35 NFE; distilled · 1-4 NFE

Implementation notes

• Compute ᾱ in f64; with linear β the last values are around 4e-5 and 1 − ᾱ near t = 0 is about 1e-4.
• Keep an EMA of weights (decay 0.9999 in DDPM/ADM) and sample from the EMA copy.
• Timestep input: DDPM-style integer index or continuous t (flow, EDM uses log σ / 4). Match the convention the sampler uses exactly.
• Pre-encode VAE latents and text embeddings offline; online they often cost more than the diffusion step.
• EDM (Karras et al.) σ-parametrization defaults: σ_min = 0.002, σ_max = 80, ρ = 7 time discretization, σ_data = 0.5.
• adaLN-Zero / zero-initialized output convs make the untrained model predict zero, which stabilizes early training.
• Run attention with FlashAttention and bf16; keep the loss and the schedule arithmetic in fp32.

Pitfalls

• Non-zero terminal SNR: SD's scaled-linear schedule (β 0.00085 → 0.012) ends at ᾱ_T ≈ 0.0047, so training never sees pure noise but sampling starts from it; images drift to mid brightness. Fix: rescale to zero terminal SNR, use v-prediction, start sampling at t = T (Lin et al.).
• Mixing timestep conventions (0- vs 1-based, t vs 1 − t in flow papers): samples come out as noise or blur.
• CFG scale too high: oversaturated, burnt images. Values 3-8 are typical; use guidance rescale or interval guidance at high w.
• Uniform t sampling for flow models: wastes capacity at the easy ends; use logit-normal.
• Evaluating with the non-EMA weights, or FID computed on too few samples (use 50k for comparisons).

Variants and alternatives

In production

Stable Diffusion 1.x/2.x/XL (latent U-Net), Stable Diffusion 3 and FLUX (rectified-flow transformers), DALL-E 2/3, Imagen; Hugging Face diffusers schedulers (DDPM, DDIM, DPM-Solver++, FlowMatchEuler). Also RFdiffusion (protein design) and Diffusion Policy (robot control).

Implementations

Sources

• Ho, Jain, Abbeel, Denoising Diffusion Probabilistic Models, NeurIPS 2020. paper
• Song, Meng, Ermon, Denoising Diffusion Implicit Models, ICLR 2021. paper
• Nichol and Dhariwal, Improved Denoising Diffusion Probabilistic Models, ICML 2021. paper
• Dhariwal and Nichol, Diffusion Models Beat GANs on Image Synthesis, NeurIPS 2021. paper
• Ho and Salimans, Classifier-Free Diffusion Guidance, 2022. paper
• Rombach et al., High-Resolution Image Synthesis with Latent Diffusion Models, CVPR 2022. paper · code
• Karras, Aittala, Aila, Laine, Elucidating the Design Space of Diffusion-Based Generative Models, NeurIPS 2022. paper
• Salimans and Ho, Progressive Distillation for Fast Sampling of Diffusion Models, ICLR 2022. paper
• Lipman et al., Flow Matching for Generative Modeling, ICLR 2023. paper
• Liu, Gong, Liu, Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, ICLR 2023. paper
• Peebles and Xie, Scalable Diffusion Models with Transformers, ICCV 2023. paper · code
• Hang et al., Efficient Diffusion Training via Min-SNR Weighting Strategy, ICCV 2023. paper
• Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, ICML 2024. paper
• Lin et al., Common Diffusion Noise Schedules and Sample Steps are Flawed, WACV 2024. paper
• Lu et al., DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, 2022. paper
• Watson et al., De novo design of protein structure and function with RFdiffusion, Nature 2023. paper · code
• Chi et al., Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, RSS 2023. paper · code