High-accuracy sampling for diffusion models and log-concave distributions

Summary

This paper (Chen, Chewi, Daskalakis, Rakhlin) addresses a central open question in diffusion-sampling theory: can a sampler using only score (gradient) queries achieve high-accuracy convergence — polylog($1/\delta$) query complexity — under minimal data and score-error assumptions? The answer is affirmative, via a new meta-algorithm called first-order rejection sampling (FORS) that simulates rejection sampling against a Gaussian proposal using only an unbiased estimator of the log-tilt, built out of clipped path-integral gradient differences and a Bernoulli-factory / Poisson-truncated Taylor series. Deployed step-by-step in a backward diffusion chain, FORS delivers $\delta$-error in the bounded Lipschitz metric with $O(d \log^2(M/\delta) + \log^3(M/\delta))$ queries under only a finite-second-moment assumption; under a non-uniform $L$-Lipschitz score, the complexity improves to $O(\sqrt{dL}\,\mathrm{polylog})$; with intrinsic dimension $d_\star$, to $O(d_\star \mathrm{polylog})$. A parallel application to log-concave sampling recovers state-of-the-art proximal-sampler rates without density evaluations. The conceptual move is that the FORS Bernoulli-factory step decouples proposal accuracy from discretisation error, which had been the barrier preventing polylog rates for SDE-based samplers. This is a substantive theoretical contribution that dovetails with the prior high-accuracy literature that assumed either zeroth-order queries (Huang 2024a; Wainwright 2025) or Gaussian-convolution structure (concurrent Gatmiry et al. 2026). The main conceptual and practical caveats sit largely outside the proofs. The abstract's "exponential improvement over all previous results" ignores concurrent work whose rate is incomparable rather than dominated. The polylog claim is contingent on $\tilde O(\delta)$-accurate score estimates in $L^2$ — a training standard that must shrink as the caller demands smaller $\delta$. The FORS per-step overhead of $B e^{2B} \log(1/\delta)$ score queries, together with an unstated concrete choice of $B$, an overloaded symbol $G$ between main text and appendix, and an omitted proof of Theorem F.1 delegated to prior work, pushes the practical relevance further than the paper acknowledges. The theoretical machinery, however, is genuinely novel, and the framework's uniform reach across score-based diffusion and log-concave sampling is elegant.

Strengths

• Conceptual contribution. Reframing the RGO/backward-transition step as a Bernoulli-factory problem on an unbiased gradient-integral estimator is a genuinely new mechanism for closing the gap between SDE discretisation and polylog convergence. The mechanism is also unifying: the same FORS subroutine drives the diffusion, low-dimensional, and log-concave results.
• Assumption footprint. The minimal-data result requires only a second moment, and the intrinsic-dimension and Lipschitz results use notably weak notions (the "non-uniform" Lipschitz condition is a with-high-probability smoothness that holds unconditionally with $L_\delta = O(d + \log(1/\delta))$ and is much smaller for structured data). This is a genuine assumption weakening relative to prior high-accuracy work that required density oracles or sub-exponential score error.
• Log-concave transfer. Recovering Fan et al. (2023)'s proximal-sampler rates without density evaluations is a concrete deliverable outside diffusion, and the paper connects the two threads cleanly through Theorem 3.3.
• Honest positioning. The related-work discussion, including the appendix comparison with Gatmiry et al. (2026), is careful, and the paper explicitly acknowledges the practical price of FORS's per-step queries in the conclusion.
• Modular proof structure. Section 3 factors out an exact-oracle analysis, and Section 4 adds the score-error robustness on top; Appendix E cleanly separates convergence-under-true-score (Proposition E.3) from score-error (Proposition E.4). This makes the theorems easier to verify and generalise.

Weaknesses

• Overclaim in the abstract. "Exponential improvement over all previous results" is only true against DDPM-style score-only samplers. Zeroth-order high-accuracy samplers (Huang 2024a; Wainwright 2025) are also polylog($1/\delta$), and the concurrent Gatmiry et al. (2026) rate is incomparable. The paper's body is careful about this distinction; the abstract should be too.
• Non-uniform Lipschitz result is conditional on structure. Theorem 4.4's $O(\sqrt{dL_\delta}\,\mathrm{polylog})$ collapses to essentially the same rate as Theorem 4.1 when $L_\delta = O(d + \log(1/\delta))$ (the unconditional case explicitly noted in the text). The abstract's $\sqrt{dL}$ bullet reads as a strict improvement over the minimal-data case, but the strict improvement only holds when the data admits a much smaller $L_\delta$ than the generic bound.
• Log-concave result quietly assumes an exact proximal oracle. Theorem F.1 is proved under $x_+ = \mathrm{prox}_{\eta f}(x_0)$; computing the prox requires additional gradient queries whose cost is not folded into the reported complexity. The "using only gradient evaluations" claim needs either a fully general prox-error analysis or a clear statement that the reported query counts are per-outer-iteration.
• Contingency on score-error scaling. Every headline complexity has the form $\delta^2 + (\cdot)\sum \eta_t \varepsilon^2_{t,\mathrm{score}}$, so the polylog $1/\delta$ story only kicks in once $\varepsilon_{\mathrm{score}}$ is also driven to $\tilde O(\delta)$. This is a nontrivial ask on the training side and deserves explicit acknowledgement outside the abstract's parenthetical.
• Symbol overloading obscures step-size regime. $G$ has one meaning in the main text ($C(d + \log(1/\delta))\log(d/\delta)$) and a materially different one throughout Appendix E.4–E.7 ($\eta + (\eta/\sigma^2)(d + \log(1/\delta))$). Because Proposition E.3's step-size hypothesis reads very differently under the two definitions, this is a substantive readability problem, not a cosmetic one.
• Local technical issues in the proofs. (i) Algorithm 1's per-iteration acceptance probability as printed does not obviously lie in $[0,1]$ — the motivating example uses $(1 + 2W_j)/2$ but the algorithm text reads $\prod (B + 2BW_j)$. (ii) The Poisson tail choice $c = e^{-2B}$ in the proof of Theorem 3.1 seems tuned for the first term of the concentration bound only; the second term (involving $c^2/(c+2)$) does not appear to reach $\delta/2$ with this $c$. (iii) The inequality $\mathbb{E}[Y\,\mathbf 1\{Y>M\}] \le M^{\lambda-1}\mathbb{E}[Y^\lambda]$ (used before Proposition E.1 (2)) has the exponent sign reversed; the intended bound is $M^{1-\lambda}\mathbb{E}[Y^\lambda]$.
• Reinterpreted target in concurrent-work comparison. The dismissal of Gatmiry et al. (2026)'s dimension-free rate as "we should expect poly($1/\delta$) dependence" swaps their target ($p_{\mathrm{data}} = p^\star * \mathcal N(0,\sigma^2 I)$) for the un-noised base $p^\star$ mid-argument. The rate is genuinely incomparable, but the comparison is not like-for-like as stated.
• Practical bar is understated. The per-step FORS cost of $3 B e^{2B} \log(2/\delta)$ score queries (Theorem 3.1) is a constant $\approx 22$ at $B = 1$ with rapid growth in $B$, and the paper offers no guidance on choosing $B$ given a raw un-clipped $W_{r,z,x}$. The conclusion nods at this, but does not quantify the concrete gap between the polylog scaling and DDPM's per-step-equals-one-query accounting.
• Terminology drift. The minimal-data bullet advertises "queries", the Lipschitz and intrinsic-dimension bullets call the analogous quantity "steps". Since a step is a FORS invocation and can burn many queries, the distinction is precisely what the paper's central claim rests on and should be made uniformly.
• Filing does not contain the current best results. A footnote directs the reader to an arxiv version with improved bounds, which is awkward for the archival submission on review.

Reproducibility & code

• No code accompanies the submission. The bundled codebase directory is empty. This is defensible for a theory paper, but nothing is provided to sanity-check the polylog scaling, the choice of $B$, or the per-step evaluation count on even a Gaussian target — so the practical-improvement claim in the conclusion has no empirical anchor.
• Theorem F.1 proof omitted. The Section 5 log-concave headline results (LSI, Poincaré, log-concavity) are deferred wholesale to Altschuler & Chewi (2024). A reader trying to verify the rates must recover the coupling between inner FORS accuracy and outer proximal-sampler error from a different paper.
• Intrinsic-dimension covering radius is $\delta$-dependent. Definition 4.6 uses $r = \sigma_0/\sqrt d$; since $\sigma_0 \asymp \delta/\sqrt{d + M_2^2}$ in Corollary 4.2, $d_\star$ actually depends on $\delta$, but this dependence is not carried through the $O(d_\star \mathrm{polylog})$ statement. Someone implementing an intrinsic-dimension experiment cannot compute $d_\star$ without pinning down this schedule.
• Absolute constants unspecified. Every headline theorem is stated with a $\gg$ step-size condition and a "sufficiently large" constant. Whether the polylog scaling is empirically visible in moderate $d$ depends on the size of these constants; the paper leaves them existential.
• FORS parameter $B$ has no concrete guidance. All bounds assume $B = \Theta(1)$, but for a given target the natural range of the un-clipped path-integral estimator can be much larger, and clipping introduces bias. A recommended default (or diagnostic recipe for $B$) is needed to make Algorithm 1 reproducible in practice.
• Assumption E.2 confines Section 4.2–4.3. All DDPM-like and low-dimensional analyses assume the current-time law is exactly a Gaussian convolution; readers reproducing the results should know this is not a corollary of the main-text framing but a genuine appendix assumption.

Recommended Changes

Essential

• Qualify the "exponential improvement" headline. Rewrite the abstract sentence "This is an exponential improvement over all previous results" to make explicit that the improvement is over score-only samplers, and that zeroth-order and concurrent Gaussian-convolution samplers achieve incomparable polylog rates. (Ties to Weaknesses: overclaim in the abstract.)
• Fix or restate the Algorithm 1 acceptance probability. Print the per-iteration acceptance probability in a form ($(B + W_j)/(2B)$, or an equivalent explicitly normalised expression) that is manifestly in $[0,1]$ for $W_j \in [-B, B]$, and check that the guarantee in Theorem 3.1 reads off directly.
• Rework the Poisson tail step in Theorem 3.1. Show explicitly that the choice $c = e^{-2B}$, $i_0 = e^{2B}\log(2/\delta)$ bounds both terms of the concentration inequality by $\delta/2$, or amend the constants so that the $3 B e^{2B}\log(2/\delta)$ conclusion is correct.
• Rename the appendix $G$. Use a distinct symbol for $\eta + (\eta/\sigma^2)(d + \log(1/\delta))$ throughout Appendix E.4–E.7 so it does not collide with the main-text complexity constant $G = C(d + \log(1/\delta))\log(d/\delta)$.
• Correct the tail inequality in the proof of Theorem 4.1. The inequality preceding Proposition E.1 (2) should read $\mathbb{E}[Y\,\mathbf 1\{Y>M\}] \le M^{1-\lambda}\mathbb{E}[Y^\lambda]$ (not $M^{\lambda-1}$).
• State when Theorem 4.4 strictly beats Theorem 4.1. Explicitly quantify the regime where $\sqrt{dL_\delta}$ is materially smaller than $d$, and mark the Section 1.1 bullet as conditional on $L_\delta \ll d$ rather than an unconditional improvement.
• Fold the prox-map cost into the log-concave complexity. Either present Theorem F.1 with the prox-oracle assumption boxed as an "under an exact proximal oracle" caveat in the abstract and Section 1, or supply the query-complexity accounting when the prox is realised by inner gradient descent as in Altschuler & Chewi (2024).
• Include the Theorem F.1 proof (or a real sketch). At minimum, show the coupling between inner FORS accuracy $\delta_{\mathrm{inner}}$ and the outer proximal-sampler target so the log-concave rates are verifiable from this manuscript alone.

Suggested

• Fix the "steps" vs "queries" wording in the contribution bullets so all three lines report the same accounting.
• Add a one-line explanation of $C_{\mathrm{apx}} = O(\sqrt{d/L})$. Note that $C_{\mathrm{apx}}$ enters the score-error weighting rather than a per-step penalty, so the inverse dependence on $L$ is not a typo.
• Reconcile the concurrent-work comparison. Compare Gatmiry et al. (2026) at their own target ($p^\star * \mathcal N(0,\sigma^2 I)$), then separately note the additional cost of going to $p^\star$; the current text switches targets mid-argument.
• State the score-error scaling in the abstract. One extra clause: "$\tilde O(\delta)$-accurate in $L^2$ for the same $\delta$" — so the training-side requirement is visible from the headline.
• Recommend a default for FORS's parameter $B$. Even a diagnostic ("choose $B$ so that unclipped $|W|$ exceeds $B$ with probability $\le \delta$") would make Algorithm 1 reproducible.
• Pin the covering radius in Definition 4.6 to a $\delta$-independent scale, or write $d_\star(\delta)$ throughout the intrinsic-dimension theorem statement.
• Add a small illustrative experiment. A one-page sanity check on a Gaussian mixture (where the exact score is available and $L_\delta = O(\log(dH/\delta))$) would substantiate the "opens a path toward practical improvements" comment; the paper already identifies the natural test case.
• Resolve the arxiv-version footnote. Either add a one-sentence delta describing what the improved arxiv results give beyond this filing, or omit the footnote.
• Quantify the KLS-based log-concave bullet. State the warm-start dependence explicitly (the operator-norm second-moment prefactor and the $\log(\chi^2/\varepsilon^2)$ factor) so this bullet is not misread as unconditional.
• Flag Assumption E.2 in the main text. A single sentence at the top of Section 4.2 stating that all DDPM-like analyses work under the exact-Gaussian-convolution structure would prevent readers from over-generalising.
• Rewrite the impact statement as a paragraph that lists concrete downstream applications rather than declining to comment.