One-sentence Summary

Diffusion model sampling algorithms are proposed that achieve $\delta$δ-error in $\text{polylog}(1/\delta)$polylog(1/δ) steps given $\widetilde{O}(\delta)$O(δ)-accurate score estimates in $L^{2}$L2, yielding an exponential improvement over prior work with complexity $\widetilde{O}(d_{\star}\text{polylog}(1/\delta))$O(d⋆​polylog(1/δ)) under intrinsic dimension $d_{\star}$d⋆​ and providing the first $\text{polylog}(1/\delta)$polylog(1/δ) complexity sampler for general log-concave distributions using only gradient evaluations.

Key Contributions

• The work gives diffusion sampling algorithms that reach $\delta$δ error in $\operatorname{polylog}(1/\delta)$polylog(1/δ) steps when provided with $\widetilde{O}(\delta)$O(δ)-accurate score estimates in $L^{2}$L2, yielding $\widetilde{O}(d\operatorname{polylog}(1/\delta))$O(dpolylog(1/δ)) complexity under minimal data assumptions and an exponential improvement in dependence on target accuracy over previous results.

• The approach samples from the Bayes-rule diffusion transition rather than relying only on Gaussian DDPM updates, and the resulting guarantees remain robust to score estimation error while giving $\widetilde{O}(\sqrt{dL}\operatorname{polylog}(1/\delta))$O(dL​polylog(1/δ)) complexity under a non-uniform $L$L-Lipschitz condition and $\widetilde{O}(d_{\star}\operatorname{polylog}(1/\delta))$O(d⋆​polylog(1/δ)) complexity when the data have intrinsic dimension $d_{\star}$d⋆​.

• The framework also produces the first $\operatorname{polylog}(1/\delta)$polylog(1/δ)-complexity sampler for general log-concave distributions using only gradient evaluations.

Introduction

The authors study the complexity of sampling from continuous distributions when only evaluations of the gradient of the log-density are available, a setting central to diffusion-based generative models because these models learn score functions rather than unnormalized densities. This matters because fast convergence, ideally with iteration counts scaling polylog(1/delta) in the target accuracy, would allow high-accuracy generation without excessive score queries. Prior work faces a key limitation: score-only diffusion samplers typically rely on stochastic differential equation discretizations, whose discretization error prevents high-accuracy guarantees and leads to polynomial or only modestly subpolynomial dependence on 1/delta, while existing polylog-type results require density evaluations or stronger assumptions. The authors address this by introducing first-order rejection sampling, a meta-algorithm that simulates rejection sampling using only first-order queries, and show that it yields high-accuracy polylogarithmic guarantees for diffusion sampling under minimal data assumptions, with additional refinements for Lipschitz scores, low intrinsic dimension, and log-concave sampling.

Method

The method is organized around a reusable sampling primitive designed to sample from a tilted Gaussian target without direct access to the density value. The core target has the form

$$\nu(x)\propto \exp\left(-f(x)-\frac{\|x-x_0\|^2}{2\eta}\right),$$ν(x)∝exp(−f(x)−2η∥x−x0​∥2​),

where the algorithm is allowed to query only first order information about $f$f. In the diffusion sampling setting, one chooses $f=-\log p_t$f=−logpt​ and replaces $-\nabla f$−∇f by a score estimate. The main challenge is to perform rejection sampling against a Gaussian proposal while avoiding expensive integration or evaluation of $f$f itself.

The authors introduce a first order rejection sampling module. The abstract task is to sample from a target

$$\widehat p(x)\propto q(x)\exp(w(x)),$$p​(x)∝q(x)exp(w(x)),

where $q$q is a proposal distribution and $w(x)$w(x) is not directly evaluable. Instead, for each $x$x, the sampler can draw random variables $W_1,W_2,\ldots$W1​,W2​,… such that

$$\mathbb{E}[W_i\mid x]=w(x).$$E[Wi​∣x]=w(x).

The key observation is that an unbiased estimator of $w(x)$w(x) is enough to simulate the required Bernoulli acceptance probability. This is a Bernoulli factory construction. If the random variables are bounded after clipping, say in $[-B,B]$[−B,B], the acceptance step can be performed using a Poissonized product of scaled unbiased estimates. The resulting sampler outputs a point with law approximately proportional to

$$q(x)\exp(\mathbb{E}[W_i\mid x]),$$q(x)exp(E[Wi​∣x]),

and the expected number of drawn estimators grows on the order of

$$O\left(Be^{2B}(T+\log(1/\delta))\right)$$O(Be2B(T+log(1/δ)))

over $T$T calls, up to constant factors, with probability at least $1-\delta$1−δ.

For a Gaussian tilt, the first order rejection sampling module is instantiated by constructing a Gaussian proposal from a local linear approximation of $f$f. Choose a point $x_+$x+​ and form

$$q(x)\propto \exp\left(-\langle x-x_+,\nabla f(x_+)\rangle-\frac{1}{2\eta}\|x-x_0\|^2\right).$$q(x)∝exp(−⟨x−x+​,∇f(x+​)⟩−2η1​∥x−x0​∥2).

Equivalently, if

$$\widehat x=x_0-\eta\nabla f(x_+),$$x=x0​−η∇f(x+​),

then the proposal can be taken as

$$q=\mathsf N(\widehat x,\eta I).$$q=N(x,ηI).

The residual log tilt between $\nu$ν and $q$q can be written as a path integral of gradient differences. In the simplest form,

$$\log \nu(x)-\log q(x)-\mathrm{const} = \int_0^1 \langle x-x_+,\nabla f(x_+)-\nabla f(rx+(1-r)x_+)\rangle\,dr.$$logν(x)−logq(x)−const=∫01​⟨x−x+​,∇f(x+​)−∇f(rx+(1−r)x+​)⟩dr.

This suggests the unbiased estimator

$$W_{r,x} = \langle x-x_+,\nabla f(x_+)-\nabla f(rx+(1-r)x_+)\rangle,$$Wr,x​=⟨x−x+​,∇f(x+​)−∇f(rx+(1−r)x+​)⟩,

where $r\sim\mathrm{Unif}(0,1)$r∼Unif(0,1). In practice the estimator is clipped to a fixed interval:

$$\widehat W_{r,x}=\operatorname{Clip}_B(W_{r,x}),$$Wr,x​=ClipB​(Wr,x​),

with $B=\Theta(1)$B=Θ(1). If $\nabla f$∇f is $\beta$β Lipschitz, then

$$|W_{r,x}|\le \beta\|x-x_+\|^2.$$∣Wr,x​∣≤β∥x−x+​∥2.

Therefore, when the Gaussian proposal is well centered near the anchor and the step size is sufficiently small, the clipped estimator is bounded with high probability.

A more refined construction uses a general path integral. Fix a random shift distribution for $z$z and a path $\gamma_{z,r}(x)$γz,r​(x) satisfying $\gamma_{z,1}(x)=x$γz,1​(x)=x and $\gamma_{z,0}(x)$γz,0​(x) independent of $x$x. Then for a smooth function $h$h,

$$h(x)-\mathbb{E}_{z\sim P}[h(\gamma_{z,0})] = \mathbb{E}_{z\sim P,r\sim\mathrm{Unif}} \langle \dot\gamma_{z,r}(x),\nabla h(\gamma_{z,r}(x))\rangle,$$h(x)−Ez∼P​[h(γz,0​)]=Ez∼P,r∼Unif​⟨γ˙​z,r​(x),∇h(γz,r​(x))⟩,

where $\dot\gamma_{z,r}(x)=\frac{d}{dr}\gamma_{z,r}(x)$γ˙​z,r​(x)=drd​γz,r​(x). This leads to the path estimator

$$W_{r,z,x} = \langle \dot\gamma_{z,r}(x), \nabla f(x_+)-\nabla f(\gamma_{z,r}(x))\rangle.$$Wr,z,x​=⟨γ˙​z,r​(x),∇f(x+​)−∇f(γz,r​(x))⟩.

A useful choice is

$$\gamma_{z,r}(x) = a_r x+(1-a_r)\widehat x+b_r z,$$γz,r​(x)=ar​x+(1−ar​)x+br​z,

with

$$a_r=\sin(\pi r/2), \qquad b_r=\cos(\pi r/2),$$ar​=sin(πr/2),br​=cos(πr/2),

and $z$z sampled from a Gaussian noise distribution. This path interpolates smoothly between the proposal mean and the current candidate while adding randomness from $z$z. Under gradient Lipschitz or Hölder continuity assumptions on $\nabla f$∇f, this path estimator yields a Gaussian tilt sampler whose sample distribution is close in chi squared divergence to the target. In the smooth case, the step size must satisfy a condition of the form

$$\eta^{-1}\gg \beta_1 d^{1/2}\log(1/\delta),$$η−1≫β1​d1/2log(1/δ),

while in the merely Lipschitz case the condition becomes

$$\eta^{-1}\gg \beta_0^2\log(1/\delta).$$η−1≫β02​log(1/δ).

The same subroutine is then used as the engine for diffusion sampling. The key structural insight is that one reverse diffusion step can be viewed as sampling from a Gaussian tilted target, where the tilt is determined by the score function. At step $t$t, let the forward process have variance $\sigma_t^2$σt2​ and signal scaling $\alpha_t$αt​. Given a state $X_{t+1}$Xt+1​, define the deterministic anchor

$$\bar x_t=\alpha_t^{-1}X_{t+1}.$$xˉt​=αt−1​Xt+1​.

A learned or estimated score function $\widehat{\mathsf s}_t$st​ is then queried to construct an estimator for the reverse transition target. The sampler repeatedly runs the Gaussian tilt rejection subroutine at each diffusion time, accumulating sample quality as it progresses from $t=T$t=T toward $t=1$t=1.

The first variant is a simple Gaussian proposal method. It uses

$$q_t=\mathsf N(\bar x_t,\eta_t I)$$qt​=N(xˉt​,ηt​I)

and the clipped path estimator

$$\widehat W_{r,x}^t = \operatorname{Clip}_B \left( \langle x-\bar x_t, \widehat{\mathsf s}_t(rx+(1-r)\bar x_t) \rangle \right),$$Wr,xt​=ClipB​(⟨x−xˉt​,st​(rx+(1−r)xˉt​)⟩),

where $r\sim\mathrm{Unif}(0,1)$r∼Unif(0,1). The step size is chosen so that

$$\frac{\sigma_t^2}{\eta_t} \gg d\log(1/\delta)+\log^2(1/\delta).$$ηt​σt2​​≫dlog(1/δ)+log2(1/δ).

The resulting sampler tracks the reverse transition kernels from the prior distribution toward the early stopped distribution $p_1$p1​. Its error decomposition has the form

$$D_{\mathsf{KL}}(p_1\|\widehat p_1) \lesssim D_{\mathsf{KL}}(p_T\|\widehat p_T) + T\delta + \sum_{t=1}^{T-1} \eta_t\varepsilon_{t,\mathrm{score}}^2,$$DKL​(p1​∥p​1​)≲DKL​(pT​∥p​T​)+Tδ+t=1∑T−1​ηt​εt,score2​,

where $\varepsilon_{t,\mathrm{score}}^2$εt,score2​ measures the local score approximation error. The method does not require strong shape assumptions on the data distribution beyond a second moment bound. With an appropriate schedule, the number of steps needed to reach small discrepancy in the bounded Lipschitz metric is on the order of

$$\max\{d,\log(1/\delta)\} \log^2\left(\frac{d+M_2^2}{\delta^2}\right),$$max{d,log(1/δ)}log2(δ2d+M22​​),

where

$$M_2^2=\mathbb{E}_{X_0\sim p_{\mathrm{data}}}\|X_0\|^2.$$M22​=EX0​∼pdata​​∥X0​∥2.

A second instantiation is a DDPM like method. Here the proposal includes a first order score correction:

$$q_t = \mathsf N\left( \bar x_t+\eta_t\widehat{\mathsf s}_t(\bar x_t), \eta_t I \right).$$qt​=N(xˉt​+ηt​st​(xˉt​),ηt​I).

This proposal is closer to standard ancestral diffusion sampling, but the target correction is still implemented using the rejection estimator. The path used in the estimator is

$$\gamma_{z,r,\bar x_t}(x) = a_r x+(1-a_r)\bar x_t+b_r z,$$γz,r,xˉt​​(x)=ar​x+(1−ar​)xˉt​+br​z,

with

$$a_r=\sin(\pi r/2), \qquad b_r=\cos(\pi r/2),$$ar​=sin(πr/2),br​=cos(πr/2),

and

$$z\sim\mathsf N(0,\eta_t I).$$z∼N(0,ηt​I).

The estimator is

$$\widehat W_{r,z,x}^t = \operatorname{Clip}_B \left( \langle \dot\gamma_{z,r,\bar x_t}(x), \widehat{\mathsf s}_t(\gamma_{z,r,\bar x_t}(x)) - \widehat{\mathsf s}_t(\bar x_t) \rangle \right).$$Wr,z,xt​=ClipB​(⟨γ˙​z,r,xˉt​​(x),st​(γz,r,xˉt​​(x))−st​(xˉt​)⟩).

This version benefits from a nonuniform smoothness condition on the true score. Under this condition, the step size can be chosen to satisfy

$$\frac{\sigma_t^2}{\eta_t} \gg \sqrt{dL_\delta\log(d/\delta)} + L_\delta\log(d/\delta),$$ηt​σt2​​≫dLδ​log(d/δ)​+Lδ​log(d/δ),

where $L_\delta$Lδ​ is a high probability smoothness parameter. For Gaussian data one may take $L_\delta$Lδ​ constant, while for mixtures of Gaussians it grows logarithmically in the ambient dimension and the number of mixture components. The resulting total query complexity is on the order of

$$\max\left\{ \sqrt{dL_\delta\log(d/\delta)}, L_\delta\log(d/\delta) \right\} \log\left(\frac{d+M_2^2}{\delta^2}\right).$$max{dLδ​log(d/δ)​,Lδ​log(d/δ)}log(δ2d+M22​​).

The method can also be adapted when the data distribution has low dimensional structure. The motivation is that the Hessian of a Gaussian smoothed log density contains an isotropic term that can dominate the smoothness parameter. If

$$p=p_{\mathrm{data}}*\mathsf N(0,\sigma^2I),$$p=pdata​∗N(0,σ2I),

then

$$\nabla^2\log p(x) = -\frac{I}{\sigma^2} + \frac{\operatorname{Cov}(X_0\mid X=x)}{\sigma^4}.$$∇2logp(x)=−σ2I​+σ4Cov(X0​∣X=x)​.

The shifted score function

$$\widetilde{\mathsf s}_t(x) = \widehat{\mathsf s}_t(x)+\lambda_t x, \qquad \lambda_t=\frac{1}{\sigma_t^2},$$st​(x)=st​(x)+λt​x,λt​=σt2​1​,

removes the isotropic part of the curvature and often yields better behavior when the support of $p_{\mathrm{data}}$pdata​ is intrinsically low dimensional.

The intrinsic dimension is defined using covering numbers. If $N_{\mathrm{data}}(r)$Ndata​(r) is the covering number of the data support at radius $r$r, then

$$\mathsf d_\star = 1+\log N_{\mathrm{data}}(\sigma_0/\sqrt d).$$d⋆​=1+logNdata​(σ0​/d​).

The low dimensional adaptation modifies the Gaussian proposal variance through

$$\frac{1}{\bar\eta_t} = \frac{1}{\eta_t}+\lambda_t,$$ηˉ​t​1​=ηt​1​+λt​,

and uses the shifted score difference in the rejection estimator:

$$\widehat W_{r,z,x}^t = \operatorname{Clip}_B \left( \left\langle \dot\gamma_{z,r,\bar x_t}(x), \widehat{\mathsf s}_t(\gamma_{z,r,\bar x_t}(x)) - \widehat{\mathsf s}_t(\bar x_t) + \lambda_t(\gamma_{z,r,\bar x_t}(x)-\bar x_t) \right\rangle \right).$$Wr,z,xt​=ClipB​(⟨γ˙​z,r,xˉt​​(x),st​(γz,r,xˉt​​(x))−st​(xˉt​)+λt​(γz,r,xˉt​​(x)−xˉt​)⟩).

The corresponding step size condition becomes

$$\frac{\sigma_t^2}{\eta_t} \gg \mathsf d_\star\log(d/\delta) + \log^2(d/\delta),$$ηt​σt2​​≫d⋆​log(d/δ)+log2(d/δ),

and the resulting complexity depends on $\mathsf d_\star$d⋆​ rather than the ambient dimension $d$d:

$$\max\{\mathsf d_\star,\log(d/\delta)\} \log^2\left(\frac{d+M_2^2}{\delta^2}\right).$$max{d⋆​,log(d/δ)}log2(δ2d+M22​​).

The same Gaussian tilt sampler also gives a gradient only implementation of the restricted Gaussian oracle used in the proximal sampler for log concave distributions. For a target $\mu\propto e^{-f}$μ∝e−f, the proximal sampler uses an augmented density of the form

$$\widetilde\mu(x,y) \propto \exp\left( -f(x)-\frac{1}{2\eta}\|y-x\|^2 \right).$$μ​(x,y)∝exp(−f(x)−2η1​∥y−x∥2).

A Gibbs step requires sampling from the restricted Gaussian oracle

$$\operatorname{RGO}_{f,\eta,y}(x) \propto \exp\left( -f(x)-\frac{1}{2\eta}\|y-x\|^2 \right).$$RGOf,η,y​(x)∝exp(−f(x)−2η1​∥y−x∥2).

This oracle is exactly a Gaussian tilt with center $y$y. By implementing the oracle with the first order rejection sampling module, the proximal sampler avoids direct evaluations of $f$f and relies on gradient information instead. Under smoothness, log Sobolev, Poincaré, or log concavity assumptions, this yields high accuracy proximal sampling with first order query complexity, including improved dimension dependence compared with standard rejection based Gibbs schemes.

Experiment

The study applies the proposed method to proximal sampling from log-concave and isoperimetric densities by using a Gaussian tilt representation for the restricted Gaussian oracle. The findings show that this approach yields high-accuracy sampling guarantees for both smooth and Lipschitz objectives under log-Sobolev, Poincare, and log-concavity assumptions, while relying only on first-order and proximal queries. These results indicate that the method can achieve strong convergence without zeroth-order information or the accuracy degradation typical of diffusion-based discretization methods.