WP-2026-004 · ml · generative · sde

Score-Based Generative Models through SDEs - A Controlled Study

Ibrahim Laklaa · 2026

ODE score evaluations to 0.002 nats

SDEs (VE, VP, sub-VP)

bugs caught by theory tests

frameworks with parity

Abstract

A from-scratch reimplementation of Song et al. (ICLR 2021) in PyTorch and JAX. Using the analytic score of a Gaussian mixture isolates integrator error from model error: the probability-flow ODE is the most accurate integrator yet the least robust to a learned score, discretisation error shows up as a width error rather than mode collapse, and the Langevin corrector only pays in high dimension.

Question

Score-based diffusion models are usually judged on one FID number. What actually limits sample quality: the learned score, or the integrator that turns it into samples?

Forward and reverse SDEs

Data is diffused by a forward SDE and generated by integrating its time reversal (Anderson, 1982), which only needs the score $\nabla_{x} lo g p_{t} (x)$ :

d x = f (x, t) d t + g (t) d w, d x = [f (x, t) - g (t)^{2} \nabla_{x} lo g p_{t} (x)] d t + g (t) d w ˉ

The same marginals are reached by the deterministic probability-flow ODE, with a factor $\frac{1}{2}$ on the score term. VE, VP and sub-VP SDEs are implemented with their closed-form perturbation kernels, and the score network is trained by denoising score matching:

L (θ) = E_{t, x_{0}, x_{t}} [λ (t) s_{θ} (x_{t}, t) - \nabla_{x_{t}} lo g p_{0 t} (x_{t} ∣ x_{0})^{2}]

Design: make the CPU a feature

The study ran on an 8-core CPU. The headline sampler comparisons therefore use the analytic score of a Gaussian mixture: model error is exactly zero, so every difference is a property of the integrator, a cleaner measurement than the paper's.

Findings

The probability-flow ODE is the best integrator and the worst estimator. With the exact score it reaches the data log-density within 0.002 nats in 76 score evaluations, which no SDE sampler matches at 500. With a learned MNIST score its FID collapses to 828 against 7.4 for the SDE, and the failure is bimodal: 69 % of trajectories reach FID 13.4, better than the stochastic sampler at equal cost, while 31 % diverge.
Discretisation error is a width error, not mode collapse: too wide for Euler-Maruyama, too narrow for ancestral sampling. Metrics that only detect mode dropping call every run perfect.
The Langevin corrector only pays in high dimension: it loses 10 of 12 equal-budget comparisons on 2-D data and wins 3 of 4 on MNIST at 500 NFE.
Sliced Wasserstein is useless at these sample sizes: its noise floor exceeds the spread between converged samplers, so every distance is quoted against it.

Four bugs the loss never showed

A wrong sub-VP diffusion coefficient, stale discrete noise schedules, EMA weights without bias correction, and a VE network conditioned on raw $t$ instead of $lo g σ (t)$ . Each was caught by a test against theory, not by the training curve, and the write-up records the detection method.

Implementation

PyTorch and a JAX re-implementation (jit, lax.scan, exact ODE likelihood through jacfwd) checked for parity; FID written from scratch; every table generated from stored results.

Read the paper Code