Sam Babichenko

Reading notes / Diffusion models

Sampling is as easy as learning the score

Chen et al., 2023

Context. Song et al. (2021) gave diffusion models a standard procedure: learn the score of the OU-noised data by L2L^2 regression, plug it into the reverse-time SDE, and discretize. Whether this procedure provably samples from the data distribution was open under realistic assumptions. Earlier analyses needed an L∞L^\infty-accurate score, which score matching does not deliver because the network sees no data where the density is small, or a log-Sobolev inequality for the data, which rules out well-separated modes, or gave bounds exponential in dimension. This paper shows that an L2L^2-accurate score suffices for essentially any data distribution, with step counts polynomial in dd, the smoothness and the accuracy. The reverse process never has to mix, because the forward OU process did all the mixing; it only has to follow the forward marginals backward. Girsanov bounds the error in following them as the time-integrated squared drift error under the true reverse process, whose marginals are the forward marginals, which is exactly where score matching controls the error, and the error of starting from noise instead of qTq_T can only shrink by data processing.

Setting. Data qq on Rd\mathbb R^d. The forward OU process dXˉt=−Xˉt dt+2 dBtd\bar X_t=-\bar X_t\,dt+\sqrt2\,dB_t, Xˉ0∼q\bar X_0\sim q, has marginals qt→γd=N(0,Id)q_t\to\gamma^d=N(0,I_d), and its reversal Xˉt←=XˉT−t\bar X^\leftarrow_t=\bar X_{T-t} solves dXˉt←={Xˉt←+2∇ln⁡qT−t(Xˉt←)} dt+2 dBtd\bar X^\leftarrow_t=\{\bar X^\leftarrow_t+2\nabla\ln q_{T-t}(\bar X^\leftarrow_t)\}\,dt+\sqrt2\,dB_t. The algorithm (DDPM with an exponential integrator) uses estimates skhs_{kh} on a grid h=T/Nh=T/N and freezes the score on each step, dXt←={Xt←+2 sT−kh(Xkh←)} dt+2 dBt,t∈[kh,(k+1)h],X0←∼γd,dX^\leftarrow_t=\{X^\leftarrow_t+2\,s_{T-kh}(X^\leftarrow_{kh})\}\,dt+\sqrt2\,dB_t,\qquad t\in[kh,(k+1)h],\qquad X^\leftarrow_0\sim\gamma^d, a linear SDE with Gaussian transitions; it outputs pT=law(XT←)p_T=\mathrm{law}(X^\leftarrow_T). Assumptions: (A1) ∇ln⁡qt\nabla\ln q_t is LL-Lipschitz for every t≥0t\ge0, with L≥1L\ge1; (A2) m22=Eq∥x∥2<∞m_2^2=\mathbb E_q\|x\|^2\lt\infty; (A3) Eqkh∥skh−∇ln⁡qkh∥2≤εscore2\mathbb E_{q_{kh}}\|s_{kh}-\nabla\ln q_{kh}\|^2\le\varepsilon_{\rm score}^2 at every grid time. Nothing like convexity or a functional inequality is assumed on qq, and ss need not be Lipschitz.

Theorem 2 (DDPM). Under (A1)-(A3) and h≲1/Lh\lesssim1/L, TV(pT,q)≲KL(q ∥ γd) e−T+(Ldh+Lm2h)T+εscoreT,\mathrm{TV}(p_T,q)\lesssim\sqrt{\mathrm{KL}(q\,\|\,\gamma^d)}\,e^{-T}+\big(L\sqrt{dh}+Lm_2h\big)\sqrt T+\varepsilon_{\rm score}\sqrt T, the three terms being convergence of the forward process, discretization, and score error. If KL(q∥γd)≤poly(d)\mathrm{KL}(q\|\gamma^d)\le\mathrm{poly}(d) and m2≤dm_2\le d, the choices T≍log⁡(KL(q∥γd)/ε)T\asymp\log(\mathrm{KL}(q\|\gamma^d)/\varepsilon) and h≍ε2/(L2d)h\asymp\varepsilon^2/(L^2d) give TV(pT,q)≤O~(ε+εscore)\mathrm{TV}(p_T,q)\le\tilde O(\varepsilon+\varepsilon_{\rm score}) after N=Θ~(L2d/ε2)N=\tilde\Theta(L^2d/\varepsilon^2) steps, so εscore=O~(ε)\varepsilon_{\rm score}=\tilde O(\varepsilon) suffices. This matches the best known complexity of Langevin Monte Carlo under a log-Sobolev inequality, here without one.

Corollaries 3-5 (bounded support, early stopping). Let supp q⊆B(0,R)\mathrm{supp}\,q\subseteq B(0,R), R≥1R\ge1, with no density assumed, so manifold data are allowed, and stop the reverse process at time T−tT-t with t≍εW22/(d (R∨d))t\asymp\varepsilon_{W_2}^2/(\sqrt d\,(R\vee\sqrt d)). Then pT−tp_{T-t} is εTV\varepsilon_{\rm TV}-close in TV to qtq_t, which is εW2\varepsilon_{W_2}-close in W2W_2 to qq, provided εscore≤O~(εTV)\varepsilon_{\rm score}\le\tilde O(\varepsilon_{\rm TV}) and N=Θ~(d3R4(R∨d)4εTV2 εW28).N=\tilde\Theta\Big(\frac{d^3R^4(R\vee\sqrt d)^4}{\varepsilon_{\rm TV}^2\,\varepsilon_{W_2}^8}\Big). After projecting the output onto a ball of radius Θ~(R)\tilde\Theta(R), the W2W_2 distance to qq is at most ε\varepsilon with N=Θ~(d3R8(R∨d)4/ε12)N=\tilde\Theta(d^3R^8(R\vee\sqrt d)^4/\varepsilon^{12}), polynomial where the earlier manifold bound was exponential. A TV bound to qq itself is impossible here, since pTp_T has full support.

Theorems 6-7 (critically damped Langevin). The CLD forward process is dXˉ=Vˉ dtd\bar X=\bar V\,dt, dVˉ=−(Xˉ+2Vˉ) dt+2 dBd\bar V=-(\bar X+2\bar V)\,dt+2\,dB, and its reverse needs only the velocity score ∇vln⁡qt\nabla_v\ln\boldsymbol q_t. Under the analogous assumptions Theorem 6 gives the bound of Theorem 2, now on TV(pT,q⊗γd)\mathrm{TV}(p_T,q\otimes\gamma^d) and with first term KL(q∥γd)+FI(q∥γd) e−cT\sqrt{\mathrm{KL}(q\|\gamma^d)+\mathrm{FI}(q\|\gamma^d)}\,e^{-cT}. Theorem 7 shows the technique cannot do better: for q=γdq=\gamma^d, exact scores and h≤1/10h\le1/10, the path-space divergence satisfies KL(QT←∥PT)≥dhT\mathrm{KL}(\boldsymbol Q^\leftarrow_T\|\boldsymbol P_T)\ge dhT, so the Girsanov route needs h≲1/dh\lesssim1/d, the same dimension dependence as DDPM.

Step 1

The error splits into a part from starting at γd\gamma^d instead of qTq_T and a part from following the reverse path, both measured in information divergences. Compare the algorithm with the same algorithm started at qTq_T, whose path law is PTqTP^{q_T}_T, and that one with the true reverse path law QT←Q^\leftarrow_T. Data processing (the same Markov kernel acts on both starts, and marginals are images of paths) and Pinsker give TV(pT,q)≤TV(γd,qT)+TV(PTqT,QT←)≤KL(q∥γd) e−T+KL(QT←∥PTqT),\mathrm{TV}(p_T,q)\le\mathrm{TV}(\gamma^d,q_T)+\mathrm{TV}(P^{q_T}_T,Q^\leftarrow_T)\le\sqrt{\mathrm{KL}(q\|\gamma^d)}\,e^{-T}+\sqrt{\mathrm{KL}(Q^\leftarrow_T\|P^{q_T}_T)}, where the first bound is the exponential decay of KL along OU from the log-Sobolev inequality of γd\gamma^d. Information divergences are used because the reverse dynamics can be expansive: for q=N(0,σ2I)q=N(0,\sigma^2I) the reverse drift near the end is (1−2/σ2)x(1-2/\sigma^2)x, which pushes trajectories apart when σ2>2\sigma^2\gt2, so a W2W_2 coupling would need log-concavity or dissipativity and even then the error bound grows exponentially in time, while data processing in TV or KL holds for any dynamics.

Step 2

Girsanov turns the tracking error into the score error along the true process. The two path laws share the noise 2 dB\sqrt2\,dB and the start qTq_T, so Girsanov's formula involves only the drift gap bt=2{sT−kh(Xkh)−∇ln⁡qT−t(Xt)}b_t=\sqrt2\{s_{T-kh}(X_{kh})-\nabla\ln q_{T-t}(X_t)\}, and since the stochastic integral has mean zero under QT←Q^\leftarrow_T, KL(QT← ∥ PTqT)≤∑k=0N−1∫kh(k+1)hEQT←∥sT−kh(Xkh)−∇ln⁡qT−t(Xt)∥2 dt.\mathrm{KL}(Q^\leftarrow_T\,\|\,P^{q_T}_T)\le\sum_{k=0}^{N-1}\int_{kh}^{(k+1)h}\mathbb E_{Q^\leftarrow_T}\big\|s_{T-kh}(X_{kh})-\nabla\ln q_{T-t}(X_t)\big\|^2\,dt . The direction of the divergence matters here. The expectation is under the true reverse process, whose time-tt marginal is qT−tq_{T-t}, so at grid points the score error is literally (A3), and the error under the algorithm's own law never has to be bounded. The earlier interpolation method differentiates the divergence along the algorithm, so it needs the score error under the algorithm's marginals, which forced a log-Sobolev assumption.

Step 3

Each step contributes εscore2+L2dh+L2m22h2\varepsilon_{\rm score}^2+L^2dh+L^2m_2^2h^2, and the dominant term is the Brownian movement over one step, of size dhdh, times L2L^2 from the Lipschitz score. On [kh,(k+1)h][kh,(k+1)h] split the integrand into estimation error, the change of the score in time, and the change of its argument: E∥sT−kh(Xkh)−∇ln⁡qT−t(Xt)∥2≲εscore2+E∥∇ln⁡qT−khqT−t(Xkh)∥2+L2 E∥Xkh−Xt∥2.\mathbb E\|s_{T-kh}(X_{kh})-\nabla\ln q_{T-t}(X_t)\|^2\lesssim\varepsilon_{\rm score}^2+\mathbb E\Big\|\nabla\ln\frac{q_{T-kh}}{q_{T-t}}(X_{kh})\Big\|^2+L^2\,\mathbb E\|X_{kh}-X_t\|^2 . The first term is (A3), and the last is L2L^2 times the path's squared movement over a time δ=t−kh≤h\delta=t-kh\le h, which is ≲δ2m22+δd\lesssim\delta^2m_2^2+\delta d. The middle term compares two forward marginals δ\delta apart: qT−khq_{T-kh} is qT−tq_{T-t} contracted by e−δe^{-\delta}, which contributes an Lδ∥θ∥L\delta\|\theta\| term, and smoothed by N(0,(1−e−2δ)I)N(0,(1-e^{-2\delta})I). The score perturbation lemma handles the smoothing by writing the change of score as the gap between ∇H\nabla H at a point and its average over a small Gaussian tilt there: with q=e−Hq=e^{-H} and qθ∝q⋅N(θ,M)q^\theta\propto q\cdot N(\theta,M), ∇ln⁡q∗N(0,M)q(θ)=∇H(θ)−Eqθ∇H.\nabla\ln\frac{q*N(0,M)}{q}(\theta)=\nabla H(\theta)-\mathbb E_{q^\theta}\nabla H . For L≤1/(2∥M∥)L\le1/(2\|M\|) the tilt is strongly log-concave, so it concentrates within ∥M∥d\sqrt{\|M\|d} of a mode that lies within 2∥M∥∥∇H(θ)∥2\|M\|\|\nabla H(\theta)\| of θ\theta, and with ∥M∥=O(h)\|M\|=O(h) the middle term is ≲L2dh+L2h2(∥Xkh∥2+∥∇ln⁡qT−t(Xkh)∥2)\lesssim L^2dh+L^2h^2(\|X_{kh}\|^2+\|\nabla\ln q_{T-t}(X_{kh})\|^2). Bounding the moments (E∥∇ln⁡qt∥2≤Ld\mathbb E\|\nabla\ln q_t\|^2\le Ld by integration by parts) and summing gives KL(QT←∥PTqT)≲(εscore2+L2dh+L2m22h2)T\mathrm{KL}(Q^\leftarrow_T\|P^{q_T}_T)\lesssim(\varepsilon_{\rm score}^2+L^2dh+L^2m_2^2h^2)T, and Step 1 gives Theorem 2. The L2dhL^2dh term is the dhdh of Brownian movement over one step passed through the Lipschitz score, which is why h≍ε2/(L2d)h\asymp\varepsilon^2/(L^2d) and NN is linear in dd.

Step 4

For bounded data, stop early and treat the slightly noised data as the target. Without a smooth density, (A1) fails at t=0t=0 and KL(q∥γd)\mathrm{KL}(q\|\gamma^d) may be infinite, but the OU marginal qtq_t at a small time satisfies every hypothesis of Theorem 2, though with a Lipschitz constant that grows like R2/t2R^2/t^2. The explicit OU solution gives W22(qt,q)≲R2t2+dtW_2^2(q_t,q)\lesssim R^2t^2+dt, at most ε2\varepsilon^2 for the stated tt; Otto-Villani short-time regularization makes the start-up term finite, KL(qt∥γd)≲W22(q,γd)/t≲d(R∨d)3/ε2\mathrm{KL}(q_t\|\gamma^d)\lesssim W_2^2(q,\gamma^d)/t\lesssim\sqrt d(R\vee\sqrt d)^3/\varepsilon^2; and the second-order Tweedie identity bounds the Hessian through the posterior covariance of Xˉ0\bar X_0 given Xˉt\bar X_t, which is at most R2IR^2I: 11−e−2tI−e−2tR2(1−e−2t)2I ⪯ −∇2ln⁡qt ⪯ 11−e−2tI,\frac{1}{1-e^{-2t}}I-\frac{e^{-2t}R^2}{(1-e^{-2t})^2}I\ \preceq\ -\nabla^2\ln q_t\ \preceq\ \frac{1}{1-e^{-2t}}I, so every qt′q_{t'}, t′≥tt'\ge t, has L≲R2/t2≲dR2(R∨d)2/ε4L\lesssim R^2/t^2\lesssim dR^2(R\vee\sqrt d)^2/\varepsilon^4. Theorem 2 with target qtq_t gives Corollary 3, and N∝L2dN\propto L^2d is where d3R4(R∨d)4/ε8d^3R^4(R\vee\sqrt d)^4/\varepsilon^8 comes from. For W2W_2, project onto a ball: there W2W_2 is at most the radius times TV\sqrt{\mathrm{TV}} plus an exponentially small tail term, and the projection is 1-Lipschitz.

Step 5

CLD gains nothing, because its score is taken in the velocity, which moves like Brownian motion. Theorem 6 repeats Steps 2-3 with ∇vln⁡qt\nabla_v\ln\boldsymbol q_t, replacing OU's KL decay by entropic hypocoercivity. For Theorem 7 take q=γdq=\gamma^d: the joint law stays γ2d\gamma^{2d}, the exact velocity score is −v-v, Novikov holds directly, and Girsanov is an identity whose integrand is the velocity's own increment, KL(QT←∥PT)=2∑k∫kh(k+1)hE∥Vkh−Vt∥2dt ≥ 2∑k∫kh(k+1)hd (t−kh) dt=dhT,\mathrm{KL}(\boldsymbol Q^\leftarrow_T\|\boldsymbol P_T)=2\sum_k\int_{kh}^{(k+1)h}\mathbb E\|V_{kh}-V_t\|^2dt\ \ge\ 2\sum_k\int_{kh}^{(k+1)h}d\,(t-kh)\,dt=dhT, using E∥Vkh−Vt∥2≥2d(t−kh)−10d(t−kh)2≥d(t−kh)\mathbb E\|V_{kh}-V_t\|^2\ge2d(t-kh)-10d(t-kh)^2\ge d(t-kh) for h≤1/10h\le1/10. The velocity carries the Brownian increment. Underdamped Langevin gets its d\sqrt d speedup in log-concave sampling because its gradient depends only on the smooth position, whose increment is O(dh2)O(dh^2); the reverse CLD drift depends on VV, so the speedup is lost. This is a lower bound for the technique only; the authors conjecture, without proof, that CLD has no better dimension dependence in general.

Consequences and limits. Theorem 2 is a black-box reduction: any score estimator with small L2L^2 error under the forward marginals yields a sampler at polynomial cost, so all the difficulty of diffusion models is in learning the score, and known hardness of learning Gaussian mixtures or shallow-ReLU pushforwards transfers to L2L^2 score estimation, the first information-computation gaps for that task. A network that fits the empirical objective exactly learns the score of the Gaussian-smoothed training set, and Corollary 3 applied to the empirical measure (an application the paper does not make) says the sampler then reproduces training points; whether the network instead approximates the population score is the generalization question (Kadkhodaie et al.), which (A3) assumes away. The limits are that the score must be uniformly LL-Lipschitz at all times, which manifold data violate near t=0t=0, and the early-stopping fix puts large powers of d/εd/\varepsilon into the bound; that (A3) is assumed at every grid time with one εscore\varepsilon_{\rm score}; and that the score error enters as εscoreT\varepsilon_{\rm score}\sqrt T, so it must already be at the target accuracy. Later work removed the Lipschitz assumption and reached step counts nearly linear in dd (Benton et al., 2024).

Terms

dXˉt=−Xˉt dt+2 dBtd\bar X_t=-\bar X_t\,dt+\sqrt2\,dB_t
forward OU; qt→γdq_t\to\gamma^d with KL\mathrm{KL} decaying like e−2te^{-2t}
Reverse SDE
drift x+2∇ln⁡qT−t(x)x+2\nabla\ln q_{T-t}(x), noise 2 dB\sqrt2\,dB, started at qTq_T
Exponential integrator
score frozen at XkhX_{kh} on each step; linear SDE with Gaussian transitions
(A1)-(A3)
LL-Lipschitz scores, finite second moment, L2(qkh)L^2(q_{kh}) score error ≤εscore\le\varepsilon_{\rm score}
Girsanov bound
KL(Q←∥P)≤∑k∫EQ←∥sT−kh(Xkh)−∇ln⁡qT−t(Xt)∥2dt\mathrm{KL}(Q^\leftarrow\|P)\le\sum_k\int\mathbb E_{Q^\leftarrow}\|s_{T-kh}(X_{kh})-\nabla\ln q_{T-t}(X_t)\|^2dt
Theorem 2
TV≲KL(q∥γd)e−T+(Ldh+Lm2h)T+εscoreT\mathrm{TV}\lesssim\sqrt{\mathrm{KL}(q\|\gamma^d)}e^{-T}+(L\sqrt{dh}+Lm_2h)\sqrt T+\varepsilon_{\rm score}\sqrt T
N=Θ~(L2d/ε2)N=\tilde\Theta(L^2d/\varepsilon^2)
steps for TV ε\varepsilon, with h≍ε2/(L2d)h\asymp\varepsilon^2/(L^2d) and TT logarithmic
Score perturbation lemma
∥∇ln⁡q∗N(0,M)q∥≲L∥M∥d+L∥M∥∥∇H∥\|\nabla\ln\frac{q*N(0,M)}{q}\|\lesssim L\sqrt{\|M\|d}+L\|M\|\|\nabla H\| via a strongly log-concave tilt
Localization
stopped exponential martingales plus lower semicontinuity of KL replace Novikov
Early stopping
target qtq_t with t≍ε2/(d(R∨d))t\asymp\varepsilon^2/(\sqrt d(R\vee\sqrt d)); L≲R2/t2L\lesssim R^2/t^2, W2(qt,q)≤εW_2(q_t,q)\le\varepsilon
CLD lower bound
KL≥dhT\mathrm{KL}\ge dhT for Gaussian data and exact score; the velocity moves by order dhdh per step

Quiz

  1. Why does the analysis avoid Wasserstein coupling of the reverse process?
  2. Why does the total time TT need to be only logarithmic in the target accuracy?
  3. The data are a mixture of well-separated modes, a case where Langevin Monte Carlo mixes exponentially slowly. Why does Theorem 2 still give a polynomial number of steps?
  4. A network fits the empirical denoising objective exactly on nn training points. What does Theorem 2 (through Corollary 3) say about its samples?

Sources