Source-linked AI summary
High-accuracy sampling for diffusion models and log-concave distributions
Fan Chen, Sinho Chewi, Constantinos Daskalakis, Alexander Rakhlin
TL;DR
Diffusion models raise the question of the best sampling complexity achievable using score evaluations alone, while prior methods remained short of polylogarithmic dependence on target accuracy. The paper develops a high-accuracy diffusion sampler from approximately correct first-order score evaluations and reports exponentially improved accuracy dependence, with implications for log-concave sampling using first-order queries.
Problem
Prior diffusion-sampling methods achieved sub-polynomial but not polylogarithmic dependence on target accuracy, motivating the search for the best complexity using score evaluations alone.
Method
The paper treats backward-kernel sampling with approximate score evaluations as a first-order sampling setup and derives a high-accuracy diffusion-sampling template.
Results
The algorithms improve exponentially over prior works in their dependence on target accuracy, with refined complexity under a non-uniform Lipschitz condition and implications for log-concave sampling using first-order queries.
Takeaways & Limitations
The framework provides high-accuracy diffusion sampling under minimal data assumptions and supports log-concave sampling using only first-order queries.
Takeaways & Limitations
Experimental implementation and evaluation are left for future work because the work is primarily theoretical.
Abstract
from arXiv · showhide
We present algorithms for diffusion model sampling which obtain $δ$-error in $\mathrm{polylog}(1/δ)$ steps, given access to $\widetilde O(δ)$-accurate score estimates in $L^2$. This is an exponential improvement over all previous results. Specifically, under minimal data assumptions, the complexity is $\widetilde O(d_\star \mathrm{polylog}(1/δ))$ where $d_\star$ is the intrinsic dimension of the data. Further, under a non-uniform $L$-Lipschitz condition, the complexity reduces to $\widetilde O(L \mathrm{polylog}(1/δ))$. Our approach also yields the first $\mathrm{polylog}(1/δ)$ complexity sampler for general log-concave distributions using only gradient evaluations.
1 Introduction
The paper asks whether diffusion sampling can achieve polylog(1/δ) complexity using only score evaluations under minimal assumptions. It answers affirmatively with first-order rejection sampling, extending high-accuracy guarantees to diffusion and log-concave sampling.
- Motivation: Polylog(1/δ) complexity is the target for high-accuracy sampling, whereas existing discretization methods struggle to control discretization error.Higher-order methods achieve sub-polynomial dependence on 1/δ but remain short of the desired polylogarithmic complexity.
- Motivation: The central question is whether score evaluations alone suffice for high-accuracy diffusion sampling under minimal assumptions on the data distribution and score error.Diffusion models learn score functions rather than density functions, motivating this score-only formulation.
- Method: FORS simulates rejection sampling using only first-order evaluations, providing the paper’s algorithmic route to high-accuracy guarantees.The method is designed for settings where only approximate first-order information is available.
- Diffusion sampling: Under a finite second-moment assumption, the paper obtains δ error with complexity scaling in the intrinsic dimension d⋆ rather than the embedding dimension d.The guarantee uses L2-accurate score estimates and improves on prior dimension-dependent bounds in this setting.
- Diffusion sampling: Under a non-uniform L-Lipschitz condition, the complexity becomes O(L log^3((d + M2^2)/δ)) steps with Capx = O(1).The result is described as almost dimension-free and recovers the minimal-assumption result because L ≤ eO(d⋆).
- Log-concave sampling: The framework also gives a polylog(1/δ)-complexity sampler for log-concave and isoperimetric densities using gradient evaluations without density evaluations.The paper states that this recovers state-of-the-art results while removing the density-evaluation requirement.
2 Background on diffusion models
DDPMs reverse a Gaussian-noising Markov process using score functions, but the paper focuses on sampling with approximate scores and early stopping. Its framework develops algorithms that directly sample the reverse transitions while remaining robust to score-estimation error.
- Forward and reverse processes: DDPMs transform pdata into noise through a forward Markov process and reverse this process using backward transition kernels.The reverse kernels are characterized using the densities pk and Bayes' rule.
- Sampling target: The paper targets sampling from the early-stopped distribution p1, corresponding to small 1 − α0 and σ1.This is the stated sampling target before converting guarantees to pdata itself.
- Score estimates: Score estimates sk approximate the true scores s⋆k = ∇log pk, with their quality measured by a controlled mean-squared error.The denoising function is related to the score through Tweedie's identity.
- Paper approach: The proposed algorithms learn approximate score functions and use them in DDPM while providing convergence guarantees robust to score-estimation error.The framework first studies exact scores and then instantiates the method for approximate diffusion scores.
3 Key subroutine: Gaussian tilts
The key subroutine, first-order rejection sampling, samples Gaussian tilts using only first-order information by replacing unavailable function values with unbiased estimators. Its guarantees yield high-accuracy sampling under Hölder-continuous gradients, with dimension-dependent step-size conditions.
- FORS construction: The Bernoulli-factory construction converts unbiased estimates of f(x) into samples from Ber(ce^−f(x)), avoiding direct integration of f′.This supplies the rejection step using derivative information alone.
- FORS construction: FORS samples from a tilted density q(x)e^w(x) without evaluating w(x), using bounded unbiased estimators of the tilt.The estimator samples Wj conditionally on x with E[W1 | x] = w(x).
- FORS guarantee: FORS outputs the target tilted density exactly, with at most 3Be^2B log(2/δ) sampled estimators with probability at least 1 − δ.Across T calls, the total estimator count is O(Be^2B(T + log(1/δ))).
- Gaussian tilts: The Gaussian-tilt method uses a Gaussian proposal based on a first-order expansion and constructs a bounded correction estimator from ∇f.The proposal is q = N(x0 − η∇f(x+), ηI).
- Generalization: The general path-integral construction extends the approach beyond gradient-Lipschitz functions and permits Hölder continuity for every s ∈ [0, 1].It also requires only βdη ≪ 1 in the gradient-Lipschitz case.
- Gaussian tilts: Under Hölder-continuous gradients, the method achieves high-accuracy sampling, interpolating between Lipschitz and smooth cases through the exponent s.The gradient condition is ∥∇f(x) − ∇f(y)∥ ≤ βs∥x − y∥^s.
4 Diffusion sampling
The paper applies FORS to reverse diffusion transitions, obtaining high-accuracy diffusion sampling whose complexity depends on intrinsic rather than embedding dimension under minimal assumptions. Additional non-uniform Lipschitz conditions yield refined bounds, including log-smooth settings.
- Algorithmic template: Reverse diffusion sampling is treated as a special case of Gaussian-tilt sampling with approximately correct first-order evaluations.Algorithm 2 applies FORS to backward diffusion transitions using score estimates.
- Algorithmic template: The proposal distribution is chosen as an approximate minimizer of the KL divergence to the true backward transition distribution.This choice motivates the exponential-integrator-style proposal.
- Intrinsic dimension: d⋆ is an intrinsic dimension based on covering numbers, is no larger than d, and satisfies d⋆ = eO(k) for distributions supported on compact k-dimensional manifolds.Finite-support distributions also admit logarithmic intrinsic dimension bounds.
- Main diffusion guarantees: The variance-preserving construction uses K reverse steps and makes O(K) queries with probability at least 1 − δ.The forward process is chosen sufficiently close to its terminal noise distribution before reverse sampling.
- Main diffusion guarantees: Under a finite second moment and minimal data assumptions, the diffusion sampler obtains δ error with complexity depending on d⋆ and logarithmic factors.The guarantee is stated in the bounded Lipschitz metric and uses L2 score-estimation error.
- Refined guarantees: For log-smooth pdata, the paper derives KL convergence to pdata, while the non-uniform Lipschitz analysis relates complexity to smoothness parameters.In this setting, d⋆ = d necessarily.
- Refined guarantees: Under a non-uniform Lipschitz condition, the framework yields refined complexity bounds that can be sublinear in the embedding dimension.The relevant smoothness assumptions can hold unconditionally with Lop,δ = O(d⋆ + log(1/δ)), and with Lop,δ ≡ 1 for log-concave pdata.
5 Log-concave sampling
The paper transfers its first-order rejection-sampling framework to proximal sampling for log-concave and isoperimetric distributions. This produces high-accuracy guarantees using gradient and proximal-oracle queries, including a log-concave guarantee based on first-order access.
- Connection to FORS: FORS implements the restricted Gaussian oracle because the oracle distribution is a Gaussian tilt already covered by the paper's framework.This removes the need for zeroth-order queries to f in the resulting sampler.
- Proximal sampling: The proximal sampler alternates Gaussian perturbations with restricted Gaussian oracle updates, using ∇f as input.Its iterations sample Yn ∼ N(Xn, ηI) and then update through RGOf,η,Yn.
- Guarantees: Under a log-Sobolev inequality, the sampler obtains χ2 error ε2 with query complexity eO(κ(d1/2 log3/2(R/ε2) + log2(R/ε2))).Here R = log(1 + Dχ2(µ0 ∥ µ)).
- Guarantees: Under a Poincaré inequality, the sampler obtains χ2 error ε2 with a logarithmic dependence on 1/ε and χ2.The stated complexity is eO(κ(d1/2 log1/2(1/ε) + log(1/ε)) log(χ2/ε2)).
- Guarantees: Under log-concavity, the sampler obtains KL error ε2 with complexity eO(β1d1/2W2_2(µ0, µ)/ε2), and the paper notes a high-accuracy guarantee through recent KLS progress.The KLS implication uses CPI(µ) ≤ O(log d)·∥Eµ[XX⊤]∥op.
- Guarantees: For Lipschitz f under log-concavity, the stated KL guarantee has complexity eO(β2_0W2_2(µ0, µ)/ε2).The result is the s = 0 counterpart to the smooth case.
6 Conclusion
The paper presents high-accuracy diffusion samplers under minimal data assumptions, with improved dimension dependence under a Lipschitz score condition and implications for first-order log-concave sampling. Experimental implementation and evaluation are left for future work.
- The algorithms provide high-accuracy diffusion sampling under minimal data assumptions.
- A Lipschitz score assumption improves the dimension dependence of the guarantees.
- The results improve exponentially over prior works in their dependence on target accuracy.
- The framework also has implications for log-concave sampling using only first-order queries.
- Implementation and experimental evaluation are deferred to future work.
A.1 Concurrent work
Concurrent work studies accelerated ODE-based diffusion sampling under structural assumptions, while this paper claims broader applicability to general first-order sampling methods. The comparison highlights differences in assumptions and data regimes.
- Concurrent work: Concurrent work obtains a high-accuracy guarantee for an accelerated ODE flow under structural assumptions.
- Comparison: The intrinsic dimension d⋆ can upper bound the comparison work’s dimension parameter, yielding dimension-free guarantees when R/σ⋆ is constant.
- Comparison: In the early-stopping regime, matching the real data distribution typically requires σ⋆ ≪ δ, suggesting poly(1/δ) rather than polylog(1/δ) dependence.
- Comparison: The concurrent work’s assumption on score-estimate tails is substantially stronger than this paper’s average-error condition.
- Comparison: This paper applies more broadly to general first-order sampling methods, including log-concave sampling.
C.1 Proof of Theorem 3.1
The proof analyzes the first-order rejection sampler through its per-iteration acceptance probability and controls the number of iterations and auxiliary draws with exponential-tail bounds.
- The algorithm’s termination probability on each iteration is analyzed conditional on the proposal draw.
- The acceptance probability is represented explicitly as A.
- The total number of auxiliary draws is controlled by an i.i.d. sequence of Poisson(2B) random variables.
- Choosing i0 = e^2B log(2/δ) and c = e^-2B bounds the relevant tail terms by δ/2.
- The total number of auxiliary draws also satisfies an exponential-moment bound after normalization by CBe^2B.
D Proofs for Section 3
The proofs establish coverage bounds connecting diffusion kernels and reverse processes, then control the resulting errors using intrinsic-dimension estimates and Lipschitz conditions. These estimates support the paper’s diffusion-sampling applications.
- Diffusion-process comparisons: The analysis compares the DDPM kernel with reverse SDE and ODE transitions over a short time interval.
- Coverage estimates: The coverage estimates bound divergence between these transitions by score-related error terms along the diffusion process.
- Intrinsic-dimension control: Intrinsic-dimension bounds control posterior means, covariances, and related quantities needed in the error analysis.
- Lipschitz conditions: Under operator-norm or Frobenius-norm Lipschitz conditions, the bounds depend on the probability of violating the corresponding gradient threshold.
- Posterior bounds: The technical development includes posterior-moment bounds based on covering arguments for the support of the data distribution.
E.5 Why Lipschitz condition under the Frobenius norm?
The analysis shows that the Frobenius-norm Lipschitz condition reflects an instance-specific barrier for Gaussian backward-step approximations, not merely an artifact of the proof. An exact one-step KL characterization separates score-estimation error from irreducible discretization error and identifies the optimal Gaussian mean.
- Motivation: An exact characterization of one-step Gaussian-approximation KL error provides the basis for analyzing the Lipschitz condition.The argument treats the condition as an instance-specific complexity measure for Gaussian approximation of the backward transition.
- One-step Gaussian approximation: The one-step divergence U_k(η, v) measures the performance of transitions X_k ∼ N(v(X_{k+1}), ηI), and rejection sampling requires U_k(η, v)=O(1).Here η is the step size and v is an arbitrary mean function.
- Error decomposition: The minimum one-step KL decomposes into score-estimation error and an irreducible discretization term governed by ∇m_τ(Y_τ).The discretization term is intrinsic to the data distribution and cannot be removed by choosing a different mean function.
- Step-size regimes: The optimality conditions distinguish step sizes below η_k, above η_k, and within the interval [η_k, η_k] through lower bounds involving Lip_k(0) and Lip_k(∞).For η<η_k, U_k exceeds Lip_k(0); for η>η_k, it exceeds Lip_k(∞).
- Optimal proposal: The optimal mean is v(X)=X+ηs_{k+1}(X), matching the DDPM proposal considered in the paper.This choice makes the first term in the decomposition vanish.
- Lipschitz characterization: Lip_k(λ) exactly characterizes the best possible one-step performance of a Gaussian transition scheme after optimizing the step size for fixed λ.The analysis further implies that rejection sampling is expected to succeed only when the corresponding Lip_k(λ)=O(1).
E.6 Proof of Theorem E.10
The proof develops the exact one-step KL decomposition and then connects it to Gaussian transition analysis, score estimation, and the paper’s sampling chain. It uses conditional Gaussian structure, concentration, change of measure, and related propositions to obtain the theorem’s bounds.
- Setup: The proof works with the conditional law of Y_β given Y_τ and rewrites the one-step KL quantity using differential entropy and mean-square estimation terms.It introduces M_t=E∥Y_0−m_t(Y_t)∥² and the corresponding decomposition of U(η).
- Case analysis: The analysis splits into three cases and compares the normalized error 2U(η) with quantities indexed by λ, including the limiting case λ→∞.The proof also identifies equality at the appropriate optimal step size.
- Optimality: The key identity yields lower bounds for η below or above the optimal step-size range and equality for an intermediate λ when η lies in that range.The proof concludes by combining all cases.
- Lipschitz specialization: Under an L-Lipschitz score assumption, conditional strong log-concavity and log-smoothness provide covariance bounds used to control the discretization contribution.The proof sends η to τ after deriving the relevant integral bound.
- Sampling chain: Algorithm 2 forms a Markov chain with approximate backward kernels, and KL data processing plus the chain rule transfer per-step control to the full chain.The step uses estimated scores s_k≈∇log p_k and s_{k+1}≈∇log p_{k+1}, with proposal mean g(x_+)=x_++ηs_{k+1}(x_+).
- Technical bounds: The remaining proof controls Gaussian perturbations, clipping, and change-of-measure terms before invoking the propositions that imply Theorems 4.3 and 4.9.The argument uses Gaussian independence and concentration to establish the required intermediate bounds.
F.5 Proof of Proposition F.4
The proof of Proposition F.4 bounds the discrepancy between true and approximate transition distributions through Gaussian perturbation analysis, concentration, and change of measure. It establishes moment and tail controls needed for the proposition’s final comparison.
- Distributional setup: The proof fixes an interpolation parameter r and studies the joint distribution of the exact transition, Gaussian proposal, perturbation variables, and conditioning state.This setup supports the subsequent comparison between ρ_k, its approximation, and the reference transition.
- Moment bounds: The derivative term has sub-exponential norm O(η^(1/3)d_⋆), supplying a dimension-dependent control used in the moment bounds.This estimate is one of the quantitative inputs to the proposition’s argument.
- Gaussian structure: The clipped estimator and auxiliary Gaussian variables are analyzed using independence properties that preserve Gaussian marginals under the interpolation.These properties are used to reduce the relevant bounds to Gaussian concentration and moment estimates.
- Tail control: The proof applies triangle inequalities, change of measure, and concentration bounds to control the two principal tail events.The resulting estimates are combined through union bounds and integration.
- Conclusion: The distribution induced by the Gaussian proposal is identified with the reference distribution used for the desired upper bound.This identification completes the comparison after the preceding probabilistic estimates.
G Log-concave sampling
The log-concave sampling result applies the paper’s RGO-based algorithm with first-order-query implementations to obtain divergence guarantees under log-Sobolev and Poincaré assumptions. The treatment assumes access to an exact proximal point for simplicity and notes several omitted extensions.
- Algorithmic setup: The analysis applies Theorem 3.3 at each step of Algorithm 3 to implement the RGO distribution using first-order queries.For simplicity, the proof assumes the required proximal point can be accessed exactly.
- Log-Sobolev case: Under a log-Sobolev inequality, R_λ(bµ∥µ)≤ε² is achieved using at most the theorem’s stated number of first-order queries in expectation.The condition is stated to require s=1, meaning f is smooth.
- Poincaré case: Under a Poincaré inequality, D_χ²(bµ∥µ)≤ε² is achieved using at most the theorem’s stated number of first-order queries in expectation.This is the second divergence guarantee in Theorem G.1.
- Log-concave sampling: For general log-concave µ, the theorem states a KL guarantee D_KL(bµ∥µ)≤ε² using at most the stated number of first-order queries in expectation.The result follows by tracking RGO implementation error and choosing δ appropriately.
- Scope and extensions: The theorem can be extended to Latała–Oleszkiewicz inequalities and more complicated composite settings, but those extensions are omitted.The paper also notes that proximal-map computation requires separate error analysis when exact access is unavailable.