[ICML 2026] High-accuracy sampling for diffusion models and log-concave distributions
Authors: Fan Chen, Sinho Chewi, Constantinos Daskalakis, Alexander Rakhlin
Paper: https://arxiv.org/abs/2602.01338, ICML Submission
Code: N/A
Model: N/A
This is another paper with the ICML 2026 Outstanding Paper Award.
TL;DR
WHAT was done? The authors introduce First-Order Rejection Sampling (FORS), a training-free meta-algorithm that achieves high-accuracy sampling (δ-error in polylog(1/δ) steps) for generative diffusion models and log-concave distributions using only first-order gradient (score) evaluations.
WHY it matters? Traditional stochastic differential equation (SDE) samplers are limited by polynomial discretization bias, scaling as poly(1/δ). By reframing accept-reject decisions as a “Bernoulli factory” solved via line-path integrals, this work provides the first high-accuracy score-based sampler that completely avoids expensive and intractable density evaluations, proving that high-accuracy sampling is possible with standard L2 score estimates.
Executive summary: For machine learning researchers and research leaders, this work represents a major theoretical milestone that won an ICML Outstanding Paper Award. It bridges the gap between classic Markov Chain Monte Carlo (MCMC) error correction and modern diffusion models. Since its initial peer-review submission, the authors have expanded their results on ArXiv, showing that we can bypass the discretization barrier of SDEs to draw highly accurate samples exponentially faster than previously thought possible, requiring zero architectural modifications or density queries.
Details
The Discretization Barrier in Score-Based Generative Models
Traditional sampling methods from continuous probability distributions are fundamentally limited by discretization bias. In generative diffusion models, sampling is typically executed by discretizing a reverse-time SDE or ODE. To control the resulting discretization error and reach a target accuracy δ, algorithms such as Denoising Diffusion Probabilistic Models (DDPMs) must scale their number of steps polynomially, yielding query complexities on the order of O(1/δ2) or O(1/δ1/2) under refined assumptions. In classic optimization, gradient descent achieves high-accuracy exponential convergence (polylog(1/δ)) under strong convexity because discretization does not bias the optimal solution. In sampling, however, discretization directly biases the stationary distribution.
To achieve high-accuracy convergence in sampling, classic MCMC resorts to Metropolis-Hastings correction. Yet, Metropolis-Hastings requires evaluating unnormalized probability densities (zeroth-order queries), which are notoriously intractable in diffusion setups where we only learn score functions (first-order gradients of the log-density). Prior attempts to bypass this either required learning a quantized score or introducing additional zeroth-order estimators, which degrade practical viability. This paper breaks this bottleneck by establishing that high-accuracy sampling can be performed using only standard first-order score estimates, representing a major advancement over early guarantees established by Chen et al. (2023c) and Lee et al. (2023).
First-Order Rejection Sampling First Principles: The Bernoulli Factory
The theoretical foundation of this work rests on reformulating rejection sampling as an instance of the “Bernoulli factory” problem. Consider the objective of sampling from a target density proportional to p(x)∝q(x)ew(x), where q(x) is a tractable proposal distribution and w(x) is a tilt function that we cannot evaluate directly. If we can generate an unbiased estimator W such that E[W∣x]=w(x) and W is almost surely bounded within [−B,B] for a parameter B>0, the Bernoulli factory enables us to simulate a Bernoulli trial with success probability proportional to ew(x) without ever evaluating w(x) itself.
The mathematical mechanics leverage the Taylor series expansion of the exponential function to construct this randomized acceptance test. Let J∼Poisson(2B) be a random variable independent of the candidate sample x∼q. The algorithm draws J independent samples W1,…,WJ from the estimator distribution and accepts the candidate x with the product probability: ∏j=1:J(B+Wj)/2B By taking the expectation over both the Poisson variable J and the estimators, the acceptance probability matches exp(E[W1∣x]−B), which exactly mirrors the target tilted density up to a constant factor. Theorem 3.1 guarantees that this process correctly samples from the target distribution, bounding the query complexity with high probability by O(Be2B(T+log(1/δ))) for T total algorithm calls.
The Path Integral Mechanism
To construct the unbiased estimators Wj using only first-order queries, the authors introduce a general path integral framework. For a target distribution ν(x)∝exp(−f(x)−∥x−x0∥2/2η) and a proposal distribution q=N(x0−η∇f(x+),ηI) defined via a first-order expansion around a proximal step x+, the log-ratio of the target and proposal densities can be represented as an integral along a continuous trajectory.
Specifically, they define a path function γz,r(x) that interpolates between x and a base point x̂=x0−η∇f(x+): γz,r(x)=arx+(1−ar)x̂+brz where ar=sin(πr/2), br=cos(πr/2), r∼Unif([0,1]), and z∼N(0,ηI). This setup yields the path derivative γ̇z,r(x)=a′r(x−x̂)+b’rz. Applying the fundamental theorem of calculus to line-path integrals allows them to define the unbiased estimator: Wr,z,x:=⟨γ̇z,r(x),∇f(x+)−∇f(γz,r(x))⟩ In a single step of the algorithm (detailed in Algorithm 1), a candidate x is drawn from q, and J∼Poisson(2B) independent paths are simulated. The score differences evaluated along these continuous trajectories provide the estimators W1,…,WJ. The algorithm clips these values using a tuning parameter B=Θ(1) via the function ClipB(Wr,z,x) to control the variance, entirely bypassing the need for explicit density evaluations.
Optimization and Step-Size Schedules
When integrating FORS into the backward step of a diffusion model (Algorithm 2), the algorithm is designed to correct the discretization error of the reverse-time transition kernel.
Given a sample Xk+1 at backward step k+1, the proposal distribution ρ̅k is defined as: ρ̅k(⋅∣Xk+1)=N(αk−1Xk+1+αkηksk+1(Xk+1),η̅kI) This proposal corresponds to applying an exponential integrator step utilizing the learned score function sk+1, where αk and ηk are coefficients of the backward SDE, and the adjusted proposal covariance η̅k satisfies 1/η̅k=1/ηk+1/σk2 with σk denoting the noise scale. Conflating the SDE coefficient ηk with this adjusted covariance η̅k must be avoided, as the step-size adjustments directly dictate the sampling fidelity. The score function is trained to minimize the standard L2 score estimation error defined in Definition 2.1:
where sk⋆=∇logpk is the ground-truth score. To guarantee convergence to the target accuracy δ, the algorithm coordinates step sizes according to the schedule in Theorem 4.3:
where d⋆ represents the intrinsic dimension of the data distribution. By terminating the reverse process early at p1 and choosing the truncation parameter B=Θ(1), the total query complexity scales polylogarithmically with the target accuracy while maintaining robustness to score estimation errors.
Dimensionality-Agnostic Complexity and the Source of Performance
The primary source of the paper’s efficiency is its ability to adapt to the low-dimensional structure of real-world data, replacing the ambient dimension d with the intrinsic dimension d⋆. Under minimal assumptions (simply a finite second moment M22), Corollary 4.4 proves that the algorithm reaches a target discrepancy δ in the bounded Lipschitz metric using only:
queries to the score estimator. When the data distribution satisfies a non-uniform L-Lipschitz condition (Assumption 4.6) under the Frobenius norm, the complexity reduces further to O(Llog3((d+M22)/δ)) steps, which is virtually dimension-free.
To explain the mechanism behind this performance, the authors analyze the one-step KL divergence discretization error Uk(η,v) in Theorem E.10. They prove that the discretization error cannot be entirely eliminated by optimizing the mean function v of the transition kernel, but is fundamentally bounded by an irreducible term governed by the Frobenius norm of the gradient of the conditional mean function:
This exact mathematical characterization explains why the Frobenius norm of the Jacobian of the conditional mean controls the step-size requirement, proving that controlling this geometric property is both necessary and sufficient for any Gaussian-approximation-based sampler.
Placing FORS in the Broader Sampling Ecosystem
This paper significantly advances the state-of-the-art in generative modeling theory. Standard diffusion samplers are limited to polynomial convergence rates because SDE discretizations bias the trajectory. While concurrent work by Gatmiry et al. (2026) also attempts to resolve this and obtains high-accuracy guarantees, their framework relies on a highly restrictive sub-exponential score error assumption that is difficult to enforce or verify during neural network training. In contrast, FORS operates under standard, average L2 score errors.
Additionally, prior high-accuracy samplers like those proposed by Wainwright (2025) or Huang et al. (2024) rely on unnormalized density evaluations or specialized quantized score functions, which are incompatible with standard score-matching workflows. FORS resolves this, and additionally improves log-concave proximal sampling by implementing the Restricted Gaussian Oracle (RGO) steps without requiring any density queries. Since its initial award-winning submission to ICML, the authors have updated their analysis on ArXiv, deriving tighter complexity bounds and providing a more streamlined exposition of the non-uniform Lipschitz conditions.
Open Challenges: Variance and Practical Overhead
Despite its elegant theory, FORS faces several practical challenges. First, rejection sampling algorithms are susceptible to high rejection rates in high dimensions. If the proposal distribution q is not positioned extremely close to the target distribution, the acceptance probability can decay rapidly, leading to high variance in execution times.
Furthermore, the Bernoulli factory steps rely on the truncation parameter B=Θ(1) and clipping functions to keep the estimator bounded. If the target score function exhibits high local fluctuations, the clipping error will introduce bias, degrading the accuracy of the generated samples in practice. Lastly, the theoretical analysis assumes we have access to L2-accurate score estimators across the entire support, whereas practical neural score estimators often suffer from high estimation errors in low-density or out-of-distribution regions.
The High-Accuracy Frontier
This work marks a paradigm shift in our understanding of generative diffusion model complexity. By demonstrating that high-accuracy sampling is achievable using only score evaluations, the authors have bridged the gap between classic MCMC error-correction and modern score-based generative architectures. The implications of this research extend far beyond image generation, offering powerful new tools for scientific AI applications—such as molecular dynamics, structural biology, and high-precision physical simulations—where drawing exact, high-fidelity samples is critical. We highly recommend this paper to theoretical deep learning researchers and practitioners working on advanced sampling algorithms, as it establishes a foundational framework that will likely inspire the next generation of high-precision generative models.
