Source-linked AI summary
Nearly $d$-Linear Convergence Bounds for Diffusion Models via Stochastic Localization
Joe Benton, Valentin De Bortoli, Arnaud Doucet, George Deligiannidis
TL;DR
The paper addresses the lack of diffusion-model convergence bounds linear in data dimension without smoothness assumptions. It analyzes reverse-SDE discretization and obtains iteration complexity linear in dimension up to logarithmic factors, while identifying a trade-off with step-size dependence.
Problem
Prior convergence bounds lacked linear dependence on data dimension without smoothness assumptions.
Method
The analysis uses posterior-mean/score relations, initializes the reverse process in a standard Gaussian, and controls reverse-SDE discretization through expected drift bounds.
Results
The diffusion model requires steps to approximate qδ within ε^2 in KL divergence, with iteration complexity linear in data dimension up to logarithmic factors.
Takeaways & Limitations
The discretization error is ˜O(d/N) or equivalently ˜O(dη), and its dependence on d is stated to be optimal up to logarithmic factors.
Takeaways & Limitations
Improving the step-size dependence beyond ˜O(dη) appears to require stronger drift control, which brings worse dimension dependence and additional smoothness assumptions.
Abstract
from arXiv · showhide
Denoising diffusions are a powerful method to generate approximate samples from high-dimensional data distributions. Recent results provide polynomial bounds on their convergence rate, assuming $L^2$-accurate scores. Until now, the tightest bounds were either superlinear in the data dimension or required strong smoothness assumptions. We provide the first convergence bounds which are linear in the data dimension (up to logarithmic factors) assuming only finite second moments of the data distribution. We show that diffusion models require at most $\tilde O(\frac{d \log^2(1/δ)}{\varepsilon^2})$ steps to approximate an arbitrary distribution on $\mathbb{R}^d$ corrupted with Gaussian noise of variance $δ$ to within $\varepsilon^2$ in KL divergence. Our proof extends the Girsanov-based methods of previous works. We introduce a refined treatment of the error from discretizing the reverse SDE inspired by stochastic localization.
1 INTRODUCTION
The paper studies diffusion models for sampling arbitrary high-dimensional distributions under weak assumptions, using stochastic localization to sharpen reverse-SDE discretization analysis. It develops a simulation framework involving score approximation, Gaussian initialization, time discretization, and early stopping.
- Diffusion models: Diffusion models corrupt data through an SDE and learn a reverse process whose score determines the reverse dynamics.The forward process uses an Ornstein–Uhlenbeck SDE, while the reverse process can generate samples by starting from q_T and integrating backward.
- Motivation: The paper targets convergence bounds linear in data dimension up to logarithmic factors without smoothness assumptions.Earlier bounds were polynomial under finite-second-moment assumptions, but the dimension dependence was not known to be linear.
- Proof strategy: The proof extends Girsanov-based analyses by controlling reverse-path discrepancies through a differential inequality for time-varying drift errors.This refined treatment bounds the difference between true and approximate path measures more tightly than previous discretization analyses.
- Stochastic localization: Stochastic localization supplies identities that connect posterior covariance control to more precise control of diffusion time-discretization error.The paper uses the equivalence between stochastic localization sampling and diffusion models under a time change.
- Assumptions and score learning: The analysis assumes identity covariance for convenience, focuses on OU noising, and can be extended to linear SDEs including the VE SDE.Score learning can use tractable denoising or implicit score-matching objectives estimated from analytically simulated forward-process samples.
- Diffusion-model approximations: The practical approximation pipeline learns an L2 score estimator, initializes from the standard Gaussian, discretizes time, and stops at T−δ.The standard Gaussian approximates q_T because the OU process converges exponentially quickly, while early stopping avoids score blow-up near time zero.
2 MAIN RESULTS
The paper establishes diffusion-model convergence bounds linear in data dimension up to logarithmic factors, without smoothness assumptions, under finite-second-moment conditions. Its main bound separates score-estimation, reverse-SDE discretization, and forward-process convergence errors, while a tailored time-step schedule yields the iteration-complexity result.
- Main assumptions: The main result assumes finite second moments and a controlled score approximation, without requiring smoothness of the data distribution.The covariance normalization is included for convenience; the analysis can extend to general covariance matrices.
- Main bound: Theorem 1 bounds KL error through score-estimation error, reverse-SDE discretization terms κ^2dN + κdT, and forward-process error de^-2T.The three contributions respectively reflect learned-score error, discretization, and convergence of the forward SDE.
- Time-step schedule: For any N, the schedule achieves κ = ˜O(1/N) up to factors linear in T and logarithmic in 1/δ, making discretization scale as ˜O(d/N).The resulting maximum-step constraint is constant before T − 1 and decays exponentially afterward.
- Comparison with prior bounds: The bound improves prior results that either required Lipschitz scores or had at least quadratic dependence on data dimension.Removing Lipschitz assumptions avoids excluding distributions such as those supported on a submanifold.
- Iteration complexity: The resulting iteration complexity is ˜O(d log^2(1/δ) / ε^2) for approximating qδ within ε^2 in KL divergence.This is linear in data dimension up to logarithmic factors under minimal smoothness assumptions.
- Time-step schedule: A suitable schedule uses linearly spaced steps on [0, T − 1] and exponentially spaced steps on [T − 1, T − δ].The schedule satisfies γk ≤ κ min{1, T − tk+1}, forcing exponentially decaying steps near the terminal time.
3 PROOF OF THEOREM 1
The proof has three steps: bound reverse-SDE discretization error, convert it to a path-measure KL bound with Girsanov’s theorem, and transfer this bound to terminal marginals using initialization and data processing.
- Step 1: Step 1 bounds reverse-SDE discretization error through a new Itô-calculus differential inequality for the expected drift difference.The coefficients are related to conditional means and covariances, then bounded using properties of the OU process and a key lemma.
- Step 2: Step 2 uses a Girsanov-based argument to bound KL divergence between the true and approximate reverse path measures.The proof represents both processes using different probability measures on a common path space.
- Step 3: Step 3 bounds terminal KL divergence by path-measure distance and the initialization error between qT and πd.The initialization term is controlled by OU convergence, with KL(qT ||πd) ≲ de^-2T for T ≥1.
- Step 1: Lemma 2 establishes the discretization-error bound needed for the theorem.Its proof combines the differential-inequality estimates across the selected time intervals.
- Step 3: Data processing transfers the path-measure bound to KL(qδ||ptN), completing the proof of Theorem 1.The terminal distributions are pushforwards of the corresponding path measures.
4 DISCUSSION
The discussion separates score, initialization, and discretization errors, showing optimality in the dimension dependence while identifying an unresolved step-size trade-off.
- Error terms: The initialization error from approximating qT by πd decays exponentially in T and is typically negligible.This follows from convergence of the OU process to the standard Gaussian.
- Error terms: An approximate score contributes an unavoidable KL-error term of order ε2score under the stated approximation setting.The discussion says this term cannot be improved because it comes from using an approximate score.
- Discretization: With perfect scores and initialization in qT, discretization induces KL error of order κ2dN +κdT, interpreted as ˜O(d/N) or ˜O(dη).The linear dependence on d, up to logarithmic factors, is stated to be optimal.
- Open questions: The optimality of linear dependence on η remains unclear because increasing noise schedules may permit faster convergence than the worst-case Girsanov rate.Under additional smoothness assumptions, related Euler–Maruyama results give O(η2) reverse-KL error.
- Open questions: The paper identifies a trade-off between dimension dependence and step-size dependence, leaving their simultaneous improvement to future work.Better η dependence requires stronger drift control and additional smoothness assumptions, which worsen dimension dependence.
A EQUIVALENCE OF DIFFUSION MODELS AND STOCHASTIC LOCALIZATION
The appendix establishes an equivalence between the OU forward diffusion and stochastic localization after a time change and rescaling, including corresponding conditional quantities.
- Process equivalence: The OU process is transformed using a time change τ(s) and an exponential rescaling of Xτ(s).The transformed process is expressed through the initial variable and a time-changed Brownian motion.
- Process equivalence: The transformed process Us has the same law as the stochastic localization process after matching the time variables.The equivalence uses the fact that the rescaled time-changed Brownian motion has the law of standard Brownian motion.
- Conditional quantities: Conditioning on Us is equivalent to conditioning on Xt(s), so the associated conditional means and covariances have the same laws under the correspondence.The appendix identifies as(Us) with mt(Xt) and As(Us) with Σt(Xt) when t=t(s).
B PROOFS OF STOCHASTIC LOCALIZATION RESULTS
The appendix reproduces stochastic-localization results by applying Itô calculus, integration by parts, and martingale identities to the localization process.
- Proposition proofs: The propositions are standard stochastic-localization results, reproduced for convenience and based on earlier arguments in the literature.The appendix attributes Proposition 1 and the proof strategy for Proposition 2 to prior work.
- Proposition proofs: The proofs begin from the stochastic-localization definition and derive differential identities for its likelihood-related quantities.Integration by parts for continuous semimartingales is used in the derivation.
- Proposition proofs: Itô’s lemma is applied to exponential and logarithmic transforms, producing identities involving the localization drift and normalization process.The derivation includes quadratic-variation terms and stochastic integrals.
- Proposition proofs: The proof introduces a Brownian-motion representation and uses Itô’s isometry to control the resulting stochastic terms.The final expectation identity follows from the relationship between the localization process and its conditional distribution.
C ADAPTATIONS REQUIRED TO HANDLE A GENERAL COVARIANCE OF pdata
The analysis extends the convergence theorem to data distributions with general covariance and finite second moments, including singular or unknown covariance. The resulting bounds replace dimension-only terms with dependence on d and the second-moment parameter M2.
- Proof adaptations: The proofs of Lemmas 3, 4, and 5 remain unchanged in the general-covariance setting.The adaptations are concentrated in Lemmas 6 and 7 and Proposition 4.
- Proof adaptations: In Lemma 6, substituting the finite-second-moment bounds causes all instances of M2 to cancel, leaving equation (12) unchanged.The proof of equation (13) also requires no alteration.
- Generalized theorem: Under Assumption 3, Proposition 5 bounds KL(qT ||πd) by O((d + M2)e^-2T) for T ≥ 1.This proposition replaces Proposition 4 in the generalized-covariance analysis.
- Generalized theorem: Theorem 2 applies even when the data covariance is singular or unknown to the diffusion-model algorithm.The result is considered essentially comparable to Theorem 1 when M2 scales linearly with d.
D PROOF OF COROLLARY 1
The corollary constructs a time discretization that is linear near the terminal interval and exponentially refined near the noise scale δ. This schedule satisfies the theorem’s step-size condition and yields logarithmic dependence of T on d and εscore.
- Time discretization: The schedule sets t0 = 0, tN/2 = T − 1, and tN = T − δ, with linear spacing followed by exponential spacing.The exponential segment decays from 1 to δ.
- Time discretization: The step-size condition is equivalent to κ ≥ (T − 1)/(N/2) and κ ≥ (1/δ)1/N − 1.These constraints correspond to the linear and exponentially decaying portions of the schedule.
- Time discretization: For a suitable κ = Θ(...), the schedule satisfies γk ≤ κ min{1, T − tk+1} for every k.The construction proves the first part of Corollary 1.
- Convergence guarantee: Choosing T = 1/2 log(... ) yields KL(qδ||ptN) = O(εscore^2).The displayed passage gives the terminal-time choice and the resulting KL guarantee.
- Convergence guarantee: T depends only logarithmically on d and εscore, producing the stated iteration-complexity bound.The iteration bound follows from the logarithmic dependence of the terminal time.
E OMITTED PROOFS FROM SECTION 3
The omitted proofs establish the auxiliary identities and error estimates used in the convergence analysis. They combine Tweedie’s formula, conditional-expectation calculations, trace bounds, and Itô/Fokker–Planck manipulations.
- Reverse-process calculations: The reverse-process calculations use Itô’s lemma and the Fokker–Planck equation to manipulate derivatives of log qt.These identities yield the desired intermediate differential relation.
- Auxiliary lemmas: The omitted algebraic steps expand earlier identities and rearrange them to obtain equations (10) and the associated intermediate results.These steps are presented as direct consequences of equations (8) and (9).
- Auxiliary lemmas: Lemma 5 uses Tweedie’s formula for its first part and previously established results for its second part.The paper supplies proofs of both components for completeness.
- Auxiliary lemmas: Lemma 6 derives moment identities by combining conditional expectations with the forward process relation E[Xt|X0] = e^-tX0.The proof also uses the second part of Lemma 5 and trace calculations.
- Error bounds: Lemma 7 controls discretization error terms separately on the ranges k = 0, ..., M − 1 and k = M, ..., N − 1.The bounds use the assumptions γk ≤ κ and γk ≤ κ(T − tk+1), together with monotonicity properties of T − t and E[Tr(ΣT−t)].
F APPLICATION OF GIRSANOV’S THEOREM
The proof applies Girsanov’s theorem to compare the discretized reverse process with the exact reverse diffusion. Stopping times, martingale changes of measure, coupling, weak convergence, and data processing transfer the path-space comparison to the desired KL bound.
- Measure change: Girsanov’s theorem is applied on [0, tN] with the exact reverse diffusion as the reference path measure.The drift difference is represented through an adapted process satisfying the proposition’s assumptions.
- Measure change: Stopping times Tn localize the stochastic exponential so it becomes a true martingale before the change of measure.The stopped exponentials have expectation one, allowing construction of the measures Pn.
- Measure change: Under the new measure Pn, the shifted Brownian motion is a Brownian motion, and the resulting process has the discretized reverse drift up to the stopping time.After the stopping interval, the process follows the exact reverse drift.
- Coupling and limits: The processes are coupled using a common Brownian motion, so the law of Yn is Pn while the law of Y is PqT.This coupling supports the subsequent path-space convergence argument.
- Coupling and limits: Projection onto paths stopped at tN − ε, weak convergence, lower semicontinuity, and data processing yield the KL comparison before taking ε → 0.Uniform convergence of the projected paths is used to pass to the exact path-space limit.
G CONVERGENCE OF FORWARD PROCESS
Under Assumption 2, the forward OU process converges exponentially in KL divergence.
- The forward OU process converges exponentially in KL divergence under Assumption 2.
H LINEAR DEPENDENCE ON DATA DIMENSION IS OPTIMAL
For product distributions, the exact-score reverse process factorizes across copies, making total KL error scale linearly with the number of copies and showing the linear dependence on data dimension is optimal.
- The reverse process factorizes across m copies of R^d when using exact scores, q_T^⊗m initialization, and shared discretization times.
- The total KL error becomes mε^2 by tensorization of KL divergence.
- Linear dependence on the data dimension is optimal for these KL bounds.