One-sentence Summary

Researchers from MIT and Yale present a diffusion sampling method that obtains $\delta$δ-error in $\mathrm { p o l y l o g } ( 1 / \delta )$polylog(1/δ) steps using $\widetilde O ( \delta )$O(δ)-accurate score estimates in $L ^ { 2 }$L2, an exponential speedup over all previous results, with complexity $\widetilde { O } ( d _ { \star } \mathrm { p o l y l o g } ( 1 / \delta ) )$O(d⋆​polylog(1/δ)) or $\widetilde { \cal O } ( L \mathrm { p o l y l o g } ( 1 / \delta ) )$O(Lpolylog(1/δ)) under a non-uniform $L$L-Lipschitz condition, and further achieve the first $\mathrm { p o l y l o g } ( 1 / \delta )$polylog(1/δ) sampler for log-concave distributions via gradient evaluations.

Key Contributions

• Given $\widetilde{O}(\delta)$O(δ)-accurate score estimates in $L^2$L2, the sampling algorithm achieves complexity $\widetilde{O}(d_\star \operatorname{polylog}(1/\delta))$O(d⋆​polylog(1/δ)) under minimal data assumptions, where $d_\star$d⋆​ is the intrinsic dimension, an exponential improvement in accuracy dependence over prior methods.
• Under a non-uniform $L$L-Lipschitz score condition, the complexity reduces to $\widetilde{O}(L \operatorname{polylog}(1/\delta))$O(Lpolylog(1/δ)), replacing the dimension dependence with the Lipschitz constant.
• The framework yields the first $\operatorname{polylog}(1/\delta)$polylog(1/δ)-complexity sampler for general log-concave distributions using only gradient evaluations.

Introduction

The authors tackle a central question in diffusion-based generative modeling: can sampling be performed with high accuracy—meaning only polylogarithmic steps in the inverse target error—using only gradient (score) evaluations, without access to the density? Standard high-accuracy samplers rely on accept/reject methods that require density evaluations, while score-only discretizations of stochastic processes suffer from a bias that forces a polynomial dependence on 1/δ. Existing diffusion samplers, even with higher-order discretizations, achieve at best sub-polynomial but still polynomially bounded query complexity. The authors overcome this barrier by introducing first-order rejection sampling (FORS), a meta-algorithm that simulates rejection sampling using only gradient queries. Applied to diffusion models, FORS yields a query complexity of $O(\mathsf{d}_\star \log^3((d + \mathsf{M}_2^2)/\delta))$O(d⋆​log3((d+M22​)/δ)) under minimal data assumptions, and near-dimension-free guarantees under a Lipschitz score condition, representing an exponential improvement over prior work in the dependence on the accuracy parameter.

Method

The authors develop a unified high-accuracy sampling framework that operates solely with first-order (gradient) queries. The central building block is a novel subroutine called first-order rejection sampling (FORS), which allows exact sampling from a tilted distribution $\widehat{p}(x) \propto q(x) e^{w(x)}$p​(x)∝q(x)ew(x) without ever evaluating the log‑density $w(x)$w(x). Instead, for each candidate $x$x, FORS draws i.i.d. unbiased estimates $W_1, W_2, \dots$W1​,W2​,… such that $\mathbb{E}[W_1 \mid x] = w(x)$E[W1​∣x]=w(x), and uses a Bernoulli‑factory technique to accept or reject $x$x with the correct probability. Concretely, writing $e^{w(x)} = e^{-1} e^{\mathbb{E}[1 + W_1]}$ew(x)=e−1eE[1+W1​] and sampling a Poisson random variable $J$J, the probability of acceptance is recovered as $\mathbb{E}\prod_{j=1}^J \frac{1+W_j}{2}$E∏j=1J​21+Wj​​. The number of gradient queries per acceptance is bounded with high probability, making the procedure efficient when the estimates are suitably bounded.

This FORS subroutine is then specialised to Gaussian tilt distributions of the form

which arise naturally in both diffusion model sampling and log‑concave sampling. The key idea is to choose a Gaussian proposal $q = \mathcal{N}(x_0 - \eta \nabla f(x_+), \eta \mathbf{I})$q=N(x0​−η∇f(x+​),ηI) that approximates $\nu$ν via a local linearisation of $f$f. A path‑integral representation expresses the log‑ratio $\log \nu(x) - \log q(x)$logν(x)−logq(x) as an expectation:

where $\gamma_{z,r}(x)$γz,r​(x) is a carefully designed curve with $\gamma_{z,1}(x)=x$γz,1​(x)=x and $\gamma_{z,0}(x)=\bar{\gamma}(z)$γz,0​(x)=γˉ​(z) independent of $x$x. Under a Hölder‑smoothness assumption $\|\nabla f(x)-\nabla f(y)\| \le \beta_s \|x-y\|^s$∥∇f(x)−∇f(y)∥≤βs​∥x−y∥s, the estimators are clipped to a constant interval $[-B,B]$[−B,B], and Theorem 3.3 guarantees that the resulting law $\widehat{\nu}$ν is very close to the true tilt in $\chi^2$χ2 divergence provided the step size $\eta$η is chosen appropriately (e.g., $\eta^{-1} \gg \beta_1 \sqrt{d} \log(1/\delta)$η−1≫β1​d​log(1/δ) in the smooth case).

The framework is then lifted to diffusion sampling. The backward transition kernel of the diffusion can be written as a Gaussian tilt, so each reverse step becomes an application of FORS with a score estimate $\mathsf{s}_k$sk​. The proposal distribution is chosen as

which corresponds to an exponential integrator approximation and is almost the KL‑minimiser to the true transition. The “corrector” distribution that supplies the unbiased estimates is built from a path function $\gamma_{z,r,\widehat{x}}(x)$γz,r,x​(x) with coefficients $a_r,b_r$ar​,br​ that satisfy a norm preservation identity, ensuring dimension‑free variance bounds.

A central concept is the intrinsic dimension of the data distribution:

and $\mathsf{d}_\star = \dim_{\sigma_0^2/\alpha_0^2}(p_{\mathsf{data}})$d⋆​=dimσ02​/α02​​(pdata​). This quantity captures low‑dimensional structure, small support size, or bounded radius, and is never larger than the ambient dimension $d$d. Using $d_\star$d⋆​, the step size condition becomes $\sigma_k^2/\eta_k \gg \mathsf{d}_\star \log(1/\delta) + \log^2(1/\delta)$σk2​/ηk​≫d⋆​log(1/δ)+log2(1/δ), leading to a diffusion sampling complexity of order

for the bounded‑Lipschitz metric. When the data distribution is log‑smooth the same scheme yields a KL guarantee with $d$d replacing $\mathsf{d}_\star$d⋆​.

A further refinement uses non‑uniform Lipschitz conditions. Under the assumption that the score Jacobian $\nabla m_\tau(Y_\tau)$∇mτ​(Yτ​) has bounded operator norm with high probability (which holds unconditionally with $L_{\mathrm{op},\delta} = O(\mathsf{d}_\star + \log(1/\delta))$Lop,δ​=O(d⋆​+log(1/δ)) and becomes constant for log‑concave or mixture models), the effective Frobenius norm bound $L_{\mathrm{F},\delta}$LF,δ​ can be substantially smaller than $\sqrt{d}$d​. This reduces the required step size condition to $\sigma_k^2/\eta_k \gg L_{\mathrm{F},\delta} \log(\mathsf{d}_\star/\delta) + \log^2(1/\delta)$σk2​/ηk​≫LF,δ​log(d⋆​/δ)+log2(1/δ), resulting in a complexity as low as $\min\{\sqrt{d L_{\mathrm{op}}},\, \mathsf{d}_\star^{2/3} L_{\mathrm{op}}^{1/3}\}$min{dLop​​,d⋆2/3​Lop1/3​} times polylogarithmic factors. The analysis also verifies that the same Lipschitz parameter controls the smoothness of the score under the one‑step DDPM distribution, closing the loop of the induction.

Finally, the authors apply the same Gaussian tilt technique to log‑concave sampling via the proximal sampler. The restricted Gaussian oracle (RGO) that appears in the Gibbs step is exactly a tilt distribution, so FORS replaces zeroth‑order queries with first‑order ones. Under log‑concavity, strong log‑concavity, or isoperimetric assumptions, the resulting outer loop of the proximal sampler combined with FORS yields high‑accuracy guarantees whose complexities scale with dimension only through $\sqrt{d}$d​ (for smooth potentials) or are dimension‑free in the Lipschitz case. This unifies the treatment of diffusion models and classical log‑concave sampling within a single first‑order query framework.

Experiment

The evaluation applies the proximal sampler, where the restricted Gaussian oracle is realized via FORS, to sampling from log-concave and isoperimetric targets using only first-order queries. For smooth potentials (s=1), it attains high-accuracy guarantees in chi-squared or KL divergence under log-Sobolev, Poincaré, or log-concavity, with query complexity that scales polynomially with the condition number, dimension (often d^{1/2}), and the initial distributional distance. For Lipschitz potentials (s=0), analogous guarantees are derived under Poincaré and log-concavity, requiring no zeroth-order queries and avoiding discretization-error degradation. Overall, the results demonstrate that first-order sampling can achieve high precision without the limitations of diffusion-based discretization or zeroth-order oracles.