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

面向扩散模型与对数凹分布的高精度采样

Fan Chen Affiliation:  MIT Email:  fanchen@mit.edu    Sinho Chewi Affiliation:  Yale University Email:  sinho.chewi@yale.edu    Constantinos Daskalakis Affiliation:  MIT Email:  costis@csail.mit.edu    Alexander Rakhlin Affiliation:  MIT Email:  rakhlin@mit.edu

Abstract

摘要

We present algorithms for diffusion model sampling which obtain $\delta$-error in $\mathrm{polylog}(1/\delta)$ steps, given access to $\widetilde{O}(\delta)$-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\,\mathrm{polylog}(1/\delta))$ where $d$ is the dimension of the data; under a non-uniform $L$-Lipschitz condition, the complexity is $\widetilde{O}(\sqrt{dL}\,\mathrm{polylog}(1/\delta))$; and if the data distribution has intrinsic dimension $d_{\star}$, then the complexity reduces to $\widetilde{O}(d_{\star}\,\mathrm{polylog}(1/\delta))$. Our approach also yields the first $\mathrm{polylog}(1/\delta)$ complexity sampler for general log-concave distributions using only gradient evaluations.

我们给出扩散模型采样算法：在仅能访问 $L^{2}$ 意义下 $\widetilde{O}(\delta)$ 精度的 score 估计时，可在 $\mathrm{polylog}(1/\delta)$ 步内达到 $\delta$ 误差。这相对此前所有结果是指数级改进。具体而言，在对数据几乎不做假设时，复杂度为 $\widetilde{O}(d\,\mathrm{polylog}(1/\delta))$（$d$ 为数据维数）；在非一致 $L$-Lipschitz 条件下为 $\widetilde{O}(\sqrt{dL}\,\mathrm{polylog}(1/\delta))$；若数据分布本征维为 $d_{\star}$，则可降到 $\widetilde{O}(d_{\star}\,\mathrm{polylog}(1/\delta))$。同一思路还给出首个仅用梯度求值、对一般对数凹分布达到 $\mathrm{polylog}(1/\delta)$ 复杂度的采样器。

1 Introduction

1 引言

What is the complexity of sampling from a continuous probability distribution, given access to evaluations of the gradient of the log-density? In particular, can one design algorithms whose iteration complexity scales as $\mathrm{polylog}(1/\delta)$, where $\delta$ is the target accuracy, or must they necessarily take $\mathrm{poly}(1/\delta)$ steps?

给定对数密度的梯度求值，从连续概率分布中采样的复杂度是什么？特别是，能否设计迭代复杂度随 $\mathrm{polylog}(1/\delta)$ 缩放的算法（$\delta$ 为目标精度），还是必须走 $\mathrm{poly}(1/\delta)$ 步？

Complexity bounds which scale as $\mathrm{polylog}(1/\delta)$, indicating that algorithms converge exponentially fast, are known as “high-accuracy” guarantees, as they ensure that one can draw an extremely accurate sample without too many iterations. When access to the log-density of the target distribution is available, and not just its gradient, then high-accuracy samplers abound, based on accept-reject mechanisms such as rejection sampling or the Metropolis–Hastings filter.

复杂度按 $\mathrm{polylog}(1/\delta)$ 缩放、意味着算法指数收敛的界，通常称为「高精度」保证：它们确保不必过多迭代就能抽出极高精度的样本。当目标分布的对数密度本身可求值（而不只是梯度）时，基于拒绝采样或 Metropolis–Hastings 滤波等接受–拒绝机制的高精度采样器已经很多。

Without density evaluations, the answer is less clear. In this setting, existing sampling methods are typically based on discretizations of stochastic differential equations, and the need to control the discretization error precludes high-accuracy guarantees. (Note that this is unlike the setting of optimization, in which gradient descent enjoys a high-accuracy guarantee under strong convexity—there, discretization does not bias the algorithm.) A notable exception is discretization of piecewise deterministic Markov processes (PDMPs), for which 38 show that $\mathrm{polylog}(1/\delta)$ evaluations of the gradient suffice from a warm-start.

若不能求密度，情况就不那么清楚。此时现有采样方法多基于随机微分方程的离散化，而控制离散化误差会妨碍高精度保证。（这与优化不同：在强凸设定下梯度下降本身就有高精度保证——那里离散化不会给算法带来偏置。）一个显著例外是分段确定性 Markov 过程（PDMP）的离散化：38 证明，从 warm-start 出发，$\mathrm{polylog}(1/\delta)$ 次梯度求值就够。

This question has become particularly interesting in light of recent developments on diffusion-based generative modeling. Such models are based on implementing a certain reverse Markov process, in which each iteration requires evaluation of a score function—the gradient of the log-density along a diffusion process. Since these models only learn the score function and not the density function itself (unlike earlier approaches such as energy-based models), it is natural to ask what the best achievable complexity is using score evaluations alone.

随着基于扩散的生成建模兴起，这一问题变得尤为有趣。这类模型实现某种反向 Markov 过程，每一步都需要求 score 函数——即沿扩散过程对数密度的梯度。由于它们只学习 score、不学习密度本身（不同于更早的能量模型等），自然要问：仅用 score 求值，能达到的最好复杂度是什么。

More precisely, we are interested in bounding the number of steps (and queries to the score function) required to sample from a distribution $\widehat{p}$ such that

更精确地说，我们关心从分布 $\widehat{p}$ 采样所需的步数（以及对 score 的查询次数），使得

$$\displaystyle D(p_{\mathsf{data}},\widehat{p})\leq\delta+C_{\texttt{apx}}\cdot\varepsilon_{\mathsf{score}}\,,$$

where $D$ is a measure of discrepancy, $\delta\in(0,1)$ is the target accuracy, $\varepsilon_{\mathsf{score}}$ measures the error in the score function estimates, and $C_{\texttt{apx}}$ is a target approximation factor.

其中 $D$ 是某种差异度量，$\delta\in(0,1)$ 是目标精度，$\varepsilon_{\mathsf{score}}$ 度量 score 估计误差，$C_{\texttt{apx}}$ 是目标近似因子。

The initial work of 7; 27 showed that with $L^{2}$-accurate score estimates, Denoising Diffusion Probabilistic Models (DDPMs) achieve a query complexity of $1/\delta^{2}$ in total variation distance for sampling from an early stopped distribution, provided that $p_{\mathsf{data}}$ has bounded support, and with an overhead that is polynomial in the dimension and radius and with $C_{\texttt{apx}}=\widetilde{O}(1)$. By a standard argument, this can be converted to guarantees for sampling from $p_{\mathsf{data}}$ itself in a weaker metric. Importantly, these results impose almost no assumptions on the data distribution, providing strong theoretical justification for diffusion models.

早期工作 7；27 表明：在 $L^{2}$ 精确的 score 估计下，若 $p_{\mathsf{data}}$ 有界支撑，则去噪扩散概率模型（DDPM）对早停分布可在总变差距离下达到 $1/\delta^{2}$ 的查询复杂度，开销对维数与半径多项式，且 $C_{\texttt{apx}}=\widetilde{O}(1)$。用标准论证可换成在更弱度量下对 $p_{\mathsf{data}}$ 本身的保证。重要的是，这些结果几乎不对数据分布做假设，为扩散模型提供了有力的理论背书。

Since then, there has been an explosion of works aimed at extending and refining these guarantees (5; 32; 4; 15; 11; 22; 35; 6; 33; 16; 18, e.g.,), which we cannot fully survey here. However, it is known that the complexity of $1/\delta$ is unimprovable for DDPM (25), which motivates changing the algorithm. Under minimal assumptions, the algorithm of 30 achieves a query complexity of $1/\delta^{1/2}$. Other works achieved similar speed-ups assuming that the Jacobian of the score error is bounded (31; 36).

此后涌现大量扩展与精炼这些保证的工作（如 5；32；4；15；11；22；35；6；33；16；18），此处无法尽述。但已知对 DDPM，$1/\delta$ 的复杂度不可改进（25），这促使换算法。在最小假设下，30 的算法达到 $1/\delta^{1/2}$ 查询复杂度。另一些工作在 score 误差的 Jacobian 有界时得到类似加速（31；36）。

A recent line of works (18; 19; 36, e.g.,) studied higher-order discretization methods, with the resulting query complexity scaling roughly as $C_{\rm P}\,d^{1+1/p}/\delta^{1/p}$, where $p\geq 1$ is the chosen acceleration order, and $C_{\rm P}$ hides dependence on other problem parameters[^1] along with an implicit (often exponential) dependence on $p$. These results show that diffusion sampling can be performed with sub-polynomial dependence on $1/\delta$,[^2] but they still fall well short of the desired poly-logarithmic complexity.[^3]

近期一系工作（如 18；19；36）研究高阶离散化，查询复杂度大致按 $C_{\rm P}\,d^{1+1/p}/\delta^{1/p}$ 缩放，其中 $p\geq 1$ 为所选加速阶，$C_{\rm P}$ 掩盖其他问题参数[^1]，并对 $p$ 常有隐含（往往指数）依赖。这些结果表明扩散采样可以对 $1/\delta$ 达到亚多项式依赖[^2]，但仍远未达到所期望的多对数复杂度。[^3]

Meanwhile, the works 20; 42 propose to use density evaluations to achieve high-accuracy samplers. As discussed above, this is not compatible with the current practice of diffusion sampling, which only learns score estimates.

与此同时，20；42 提出用密度求值实现高精度采样器。如前所述，这与当前只学 score 估计的扩散采样实践并不兼容。

Therefore, we state our main question of interest:

因此，我们陈述关心的核心问题：

Is there a diffusion model sampler which achieves a high-accuracy guarantee using only score evaluations, under minimal assumptions on both the data distribution and the score error?

是否存在仅用 score 求值、在对数据分布与 score 误差几乎不做假设时，仍能给出高精度保证的扩散模型采样器？

1.1 Our contribution

1.1 我们的贡献

We answer this question affirmatively via a new meta-algorithm, which we call first-order rejection sampling (FORS); see Algorithm 1. This algorithm aims at simulating rejection sampling using only first-order (gradient) queries. We give consequences of our method for diffusion sampling and for log-concave sampling.

我们给出肯定答复：提出称为一阶拒绝采样（first-order rejection sampling，FORS）的元算法，见算法 1。该算法旨在仅用一阶（梯度）查询模拟拒绝采样。我们给出其对扩散采样与对数凹采样的推论。

Diffusion sampling

We show the following results, where the error is measured in the bounded Lipschitz metric and $\varepsilon_{\mathsf{score}}$ denotes the $L^{2}$-error of the score estimates (Definition 2.1).

扩散采样

我们证明如下结果：误差以有界 Lipschitz 度量衡量，$\varepsilon_{\mathsf{score}}$ 表示 score 估计的 $L^{2}$ 误差（定义 2.1）。

• • Under minimal data assumptions—namely, $p_{\mathsf{data}}$ has a finite second moment $\mathsf{M}_{2}^{2}$—we obtain $\delta$ error in $O(d\log^{2}(M/\delta)+\log^{3}(M/\delta))$ queries with $C_{\texttt{apx}}=O(1)$, where $M\coloneqq d+\mathsf{M}_{2}^{2}$.

• • 在最小数据假设下——即 $p_{\mathsf{data}}$ 有有限二阶矩 $\mathsf{M}_{2}^{2}$——我们以 $O(d\log^{2}(M/\delta)+\log^{3}(M/\delta))$ 次查询达到 $\delta$ 误差，且 $C_{\texttt{apx}}=O(1)$，其中 $M\coloneqq d+\mathsf{M}_{2}^{2}$。

This result strictly improves upon all prior results in the literature in this setting.

该结果严格改进了该设定下文献中此前所有结果。

Beyond this setting, many works in the literature aim to sample with a number of steps which is sublinear in the dimension (6; 24; 25; 43). In an orthogonal direction, other works aim to show that the number of steps depend on the intrinsic dimension of the data (2; 21; 34; 40; 29; 37; 41). We also incorporate these advances into our framework.

除此设定外，许多工作追求步数对维数亚线性（6；24；25；43）。另一正交方向则希望步数依赖数据的本征维（2；21；34；40；29；37；41）。我们也把这些进展纳入框架。

• • Assuming Lipschitz score functions—where we use the significantly weaker notion of a “non-uniform” Lipschitz constant $L$ (4.3)—we obtain $\delta$ error in $O(\sqrt{dL}\log^{3/2}(M/\delta)+L\log^{2}(M/\delta))$ steps, where $L$ is the Lipschitz constant, and with $C_{\texttt{apx}}=O(\sqrt{d/L})$.

• • 假定Lipschitz score 函数——我们使用弱得多的「非一致」Lipschitz 常数 $L$（4.3）——则以 $O(\sqrt{dL}\log^{3/2}(M/\delta)+L\log^{2}(M/\delta))$ 步达到 $\delta$ 误差，其中 $L$ 为 Lipschitz 常数，且 $C_{\texttt{apx}}=O(\sqrt{d/L})$。

• • For low-dimensional data distributions, we obtain $\delta$ error in $O(\mathsf{d}_{\star}\log^{2}(M/\delta)+\log^{3}(M/\delta))$ steps, where $\mathsf{d}_{\star}$ is the intrinsic dimension (Definition 4.5), and with $C_{\texttt{apx}}=O(d/\mathsf{d}_{\star})$ under a slightly stronger notion of error in the score estimates.

• • 对低维数据分布，以 $O(\mathsf{d}_{\star}\log^{2}(M/\delta)+\log^{3}(M/\delta))$ 步达到 $\delta$ 误差，其中 $\mathsf{d}_{\star}$ 为本征维（定义 4.5）；在略强的 score 误差概念下，$C_{\texttt{apx}}=O(d/\mathsf{d}_{\star})$。

Log-concave sampling

For sampling from log-concave (and isoperimetric) densities with gradient evaluations of the log-density, we recover state-of-the-art results from 13, except that we do not require density evaluations; see Section 5.

对数凹采样

对带对数密度梯度求值的对数凹（及等周）密度采样，我们恢复了 13 中的当前最优结果，但不需要密度求值；见第 5 节。

1.2 Related work

1.2 相关工作

As we do not have space to survey the vast literature on diffusion model guarantees, we focus on the ones most relevant to our work.

扩散模型保证的文献浩繁，篇幅所限，我们只聚焦与本文最相关者。

High-accuracy diffusion sampling

Prior to this work, 20; 42 studied diffusion sampling with additional zeroth-order queries. Specifically, given access to estimates of the log-densities with bounded error, query complexity bounds of $\widetilde{O}(d^{2}\log(1/\delta))$ (20) and $\widetilde{O}(\sqrt{d}\log^{3}(1/\delta))$ (42) under additional Lipschitz conditions on the score functions. Both works resort to standard high-accuracy sampling methods (e.g., Metropolis-adjusted Langevin). However, estimating the densities can be extremely challenging in practice.

高精度扩散采样

在本文之前，20；42 研究了带额外零阶查询的扩散采样。具体地，给定有界误差的对数密度估计，在 score 另有 Lipschitz 条件下，查询复杂度界分别为 $\widetilde{O}(d^{2}\log(1/\delta))$（20）与 $\widetilde{O}(\sqrt{d}\log^{3}(1/\delta))$（42）。两文都诉诸标准高精度采样方法（例如 Metropolis 调整 Langevin）。但在实践中估计密度可能极为困难。

Concurrent work

In concurrent work, 17 also obtains a high-accuracy guarantee for diffusion sampling. Assuming that $p_{\mathsf{data}}$ is a Gaussian convolution, i.e., $p_{\mathsf{data}}=p_{\star}*\mathsf{N}\left(0,\sigma^{2}I\right)$ where $p_{\star}$ is supported on a ball of radius $R$, and that the score errors are bounded in the sub-exponential norm, they obtain a query complexity of $\widetilde{O}((R/\sigma)^{2}\log^{2}(1/\delta))$. While their assumptions are stronger than ours, the final complexity bounds are incomparable. We defer a lengthier comparison to Section A.1.

同期工作

同期工作 17 也为扩散采样给出了高精度保证。假定 $p_{\mathsf{data}}$ 是高斯卷积，即 $p_{\mathsf{data}}=p_{\star}*\mathsf{N}\left(0,\sigma^{2}I\right)$，其中 $p_{\star}$ 支撑在半径 $R$ 的球上，且 score 误差在次指数范数下有界，他们得到查询复杂度 $\widetilde{O}((R/\sigma)^{2}\log^{2}(1/\delta))$。其假设强于我们，最终复杂度界则不可直接比较。更长的对照见附录 A.1。

Log-concave sampling

Our application to log-concave sampling is based on the proximal sampler (28; 8), and our results can be compared to 1; 13. See the book draft 10 for an overview of the subject.

对数凹采样

我们对对数凹采样的应用基于近端采样器（28；8），结果可与 1；13 比较。该主题概述见书稿 10。

1.3 Notation

1.3 记号

We make use of the following divergences between probability distributions: the total variation (TV) distance $D_{\mathsf{TV}}\left(\mu,\nu\right)\coloneqq\frac{1}{2}\int|\mu-\nu|$; the Hellinger distance $D^{2}_{\mathsf{H}}\left(\mu,\nu\right)\coloneqq\frac{1}{2}\int(\sqrt{\mu}-\sqrt{\nu})^{2}$; the KL divergence $D_{\mathsf{KL}}\left(\mu\,\|\,\nu\right)\coloneqq\En_{\mu}\log(\mu/\nu)$; the chi-squared divergence $D_{\chi^{2}}\left(\mu\;\|\;{}\nu\right)\coloneqq\mathbb{E}_{\mu}(\mu/\nu)^{2}-1$; and the Wasserstein distance $W_{2}^{2}(\mu,\nu)\coloneqq\inf_{\gamma~\text{coupling of }\mu,\nu}\En_{(X,Y)\sim\gamma}\|X-Y\|^{2}$. In addition, we consider the bounded Lipschitz metric which metrizes weak convergence:

我们使用下列概率分布之间的散度：总变差（TV）距离 $D_{\mathsf{TV}}\left(\mu,\nu\right)\coloneqq\frac{1}{2}\int|\mu-\nu|$；Hellinger 距离 $D^{2}_{\mathsf{H}}\left(\mu,\nu\right)\coloneqq\frac{1}{2}\int(\sqrt{\mu}-\sqrt{\nu})^{2}$；KL 散度 $D_{\mathsf{KL}}\left(\mu\,\|\,\nu\right)\coloneqq\En_{\mu}\log(\mu/\nu)$；卡方散度 $D_{\chi^{2}}\left(\mu\;\|\;{}\nu\right)\coloneqq\mathbb{E}_{\mu}(\mu/\nu)^{2}-1$；以及 Wasserstein 距离 $W_{2}^{2}(\mu,\nu)\coloneqq\inf_{\gamma~\text{coupling of }\mu,\nu}\En_{(X,Y)\sim\gamma}\|X-Y\|^{2}$。此外，我们考虑度量弱收敛的有界 Lipschitz 度量：

$$\displaystyle D_{\mathsf{BL}}\left(\mu,\nu\right)\coloneqq\sup\{\mathbb{E}_{\mu}f-\mathbb{E}_{\nu}f:|f|\leq 1,\,\|f\|_{\rm Lip}\leq 1\}.$$

We define

我们定义

$$\displaystyle\mathsf{Clip}_{B}(x)\coloneqq\max\{-B,\min\{B,x\}\}\,,\qquad\forall x\in\mathbb{R}\,.$$

We use $\lesssim$ and $O(\cdot)$ to hide absolute constant, i.e., $f\lesssim g$ (and $f=O(g)$) if there is an absolute constant such that $f\leq Cg$. The notation $\widetilde{O}(\cdot)$ hides logarithmic factors.

我们用 $\lesssim$ 与 $O(\cdot)$ 隐藏绝对常数，即若存在绝对常数使 $f\leq Cg$，则 $f\lesssim g$（以及 $f=O(g)$）。记号 $\widetilde{O}(\cdot)$ 隐藏对数因子。

2 Background on diffusion models

2 扩散模型背景

Recall that Denoising Diffusion Probabilistic Models (DDPMs) are based on the following forward process that transforms a sample $X_{0}\sim p_{\mathsf{data}}$ to noise:

回顾去噪扩散概率模型（DDPM）基于如下前向过程，把样本 $X_{0}\sim p_{\mathsf{data}}$ 变成噪声：

$$\displaystyle X_{0}\sim p_{\mathsf{data}}\,,\quad X_{t+1}\sim\mathsf{N}\left(\alpha_{t}X_{t},\alpha_{t}^{2}\eta_{t}I\right)\,,~~t\in[T]\,. \tag{1}$$

Note that this is a Markov chain, and it is easy to see that $X_{t}\mid X_{0}\sim\mathsf{N}\left(\bar{\alpha}_{t}X_{0},\sigma_{t}^{2}I\right)$, where

这是一条 Markov 链，易见 $X_{t}\mid X_{0}\sim\mathsf{N}\left(\bar{\alpha}_{t}X_{0},\sigma_{t}^{2}I\right)$，其中

$$\displaystyle\bar{\alpha}_{t}=\prod_{s=0}^{t-1}\alpha_{s}\,,\qquad\sigma_{t}^{2}=\alpha_{t-1}^{2}\,(\sigma_{t-1}^{2}+\eta_{t-1})\,.$$

Conversely, given the parameter sequence $(\bar{\alpha}_{t},\sigma_{t})_{t\in[T]}$, the corresponding $(\alpha_{t},\eta_{t})_{t\in[T]}$ is given by

反之，给定参数序列 $(\bar{\alpha}_{t},\sigma_{t})_{t\in[T]}$，对应的 $(\alpha_{t},\eta_{t})_{t\in[T]}$ 由下式给出

$$\displaystyle\alpha_{t}=\frac{\bar{\alpha}_{t+1}}{\bar{\alpha}_{t}}\,,\qquad\eta_{t}=\frac{\sigma_{t+1}^{2}}{\alpha_{t}^{2}}-\sigma_{t}^{2}=\frac{\sigma_{t+1}^{2}}{\bar{\alpha}_{t+1}^{2}}\,\bar{\alpha}_{t}^{2}-\sigma_{t}^{2}\,.$$

We say that the DDPM is variance-preserving if $\bar{\alpha}_{t}^{2}+\sigma_{t}^{2}=1$, and hence $\alpha_{t}^{2}=\frac{1-\sigma_{t+1}^{2}}{1-\sigma_{t}^{2}}$, $\sigma_{t+1}^{2}=\frac{1-\sigma_{t+1}^{2}}{1-\sigma_{t}^{2}}\,(\sigma_{t}^{2}+\eta_{t})$.

若 $\bar{\alpha}_{t}^{2}+\sigma_{t}^{2}=1$，则称该 DDPM 为方差保持的，从而 $\alpha_{t}^{2}=\frac{1-\sigma_{t+1}^{2}}{1-\sigma_{t}^{2}}$，$\sigma_{t+1}^{2}=\frac{1-\sigma_{t+1}^{2}}{1-\sigma_{t}^{2}}\,(\sigma_{t}^{2}+\eta_{t})$。

We let $p_{t}$ be the probability density function of $X_{t}$, and let $\rho_{t}(\cdot\mid x^{\prime})$ be the probability density function of the backward transition kernel $\mathbb{P}(X_{t}=\cdot\mid X_{t+1}=x^{\prime})$. Then, by Bayes rule, it holds that

令 $p_{t}$ 为 $X_{t}$ 的概率密度，令 $\rho_{t}(\cdot\mid x^{\prime})$ 为反向转移核 $\mathbb{P}(X_{t}=\cdot\mid X_{t+1}=x^{\prime})$ 的概率密度。则由 Bayes 法则有

$$\displaystyle\rho_{t}(x\mid x^{\prime})\propto_{x}p_{t}(x)\exp\Big(-\frac{\|x-\alpha_{t}^{-1}x^{\prime}\|^{2}}{2\eta_{t}}\Big)\,. \tag{2}$$

In this work, we focus on the task of sampling from $p_{1}$ with extremely small values of $1-\alpha_{0}$ and $\sigma_{1}$. This corresponds to early stopping; see Section 4.1 for further discussion.

本文聚焦在 $1-\alpha_{0}$ 与 $\sigma_{1}$ 极小的情形下从 $p_{1}$ 采样。这对应于早停；进一步讨论见第 4.1 节。

Score function and estimates

As above, we assume that we have access to approximate score functions $(\mathsf{s}_{t})_{t\in[T]}$ such that $\mathsf{s}_{t}\approx\mathsf{s}^{\star}_{t}\vcentcolon=\nabla\log p_{t}$ with controlled mean-squared error. Recall that by Tweedie’s identity,

Score 函数与估计

如上，我们假定可访问近似 score 函数 $(\mathsf{s}_{t})_{t\in[T]}$，使得 $\mathsf{s}_{t}\approx\mathsf{s}^{\star}_{t}\vcentcolon=\nabla\log p_{t}$，且均方误差受控。回顾由 Tweedie 恒等式，

$$\displaystyle\mathsf{s}^{\star}_{t}(x)=\nabla\log p_{t}=\frac{1}{\sigma_{t}^{2}}\En[\bar{\alpha}_{t}X_{0}-X_{t}\mid X_{t}]\,, \tag{3}$$

where the conditional expectation is taken over $X_{0}\sim p_{\mathsf{data}}$, $X_{t}\sim\mathsf{N}\left(\bar{\alpha}_{t}X_{0},\sigma_{t}^{2}\right)$.

其中条件期望对 $X_{0}\sim p_{\mathsf{data}}$、$X_{t}\sim\mathsf{N}\left(\bar{\alpha}_{t}X_{0},\sigma_{t}^{2}\right)$ 取。

Definition 2.1(Score estimation error).

Given the score estimates $(\mathsf{s}_{t})_{t\in[T]}$, we define

定义 2.1（Score 估计误差）。

给定 score 估计 $(\mathsf{s}_{t})_{t\in[T]}$，我们定义

$$\displaystyle\varepsilon_{t,\mathsf{score}}^{2}\vcentcolon=\En_{X_{t}\sim p_{t}}\|\mathsf{s}_{t}(X_{t})-\mathsf{s}^{\star}_{t}(X_{t})\|^{2}\,,\qquad t\in[T]\,.$$

Roadmap

To approximately generate a sample from $p_{\mathsf{data}}$, we learn approximate score functions $\mathsf{s}_{t}\approx\nabla\log p_{t}$. DDPM uses the approximation $\log p_{t}(x)\approx\log p_{t}(x^{\prime})+\langle\mathsf{s}_{t}(x^{\prime}),x-x^{\prime}\rangle$, yielding

路线图

为近似生成 $p_{\mathsf{data}}$ 的样本，我们学习近似 score 函数 $\mathsf{s}_{t}\approx\nabla\log p_{t}$。DDPM 使用近似 $\log p_{t}(x)\approx\log p_{t}(x^{\prime})+\langle\mathsf{s}_{t}(x^{\prime}),x-x^{\prime}\rangle$，从而得到

$$\displaystyle\rho_{t}(\cdot\mid x^{\prime})\approx\mathsf{N}\left(\alpha_{t}^{-1}x^{\prime}+\eta_{t}\mathsf{s}_{t}(x^{\prime}),\eta_{t}I\right). \tag{4}$$

In our work, we develop algorithms for sampling from (2) directly. In Section 3, we motivate our approach by assuming that $\mathsf{s}_{t}=\nabla\log p_{t}$ is exact, which has implications for log-concave sampling described in Section 5. In Section 4, we instantiate our methods on diffusion sampling and show that the convergence guarantees are robust with respect to the error in the score function estimates.

本文中，我们直接为从 (2) 采样设计算法。第 3 节先假定 $\mathsf{s}_{t}=\nabla\log p_{t}$ 精确，以引出方法，其对对数凹采样的含义见第 5 节。第 4 节把方法实例化到扩散采样，并证明收敛保证对 score 函数估计误差是稳健的。

3 Key subroutine: Gaussian tilts

3 关键子程序：高斯倾斜

In this section, we focus on the problem of sampling from a Gaussian tilt of the form

本节聚焦从如下形式的高斯倾斜中采样

$$\displaystyle\nu(x)\propto\exp\Big(-f(x)-\frac{\|x-x_{0}\|^{2}}{2\eta}\Big)\,, \tag{5}$$

where we assume access to first-order queries for $f$. For our eventual application to diffusion models, we will choose $f=\log p_{t}$ in view of Eq. 2 and approximate $\nabla f\approx\mathsf{s}_{t}$.

其中假定可对 $f$ 做一阶查询。最终用于扩散模型时，将取 $f=\log p_{t}$（见式 2），并用 $\nabla f\approx\mathsf{s}_{t}$ 近似。

3.1 First-order rejection sampling (FORS)

3.1 一阶拒绝采样（FORS）

We motivate our approach by considering the simple problem of sampling from a density $p\propto e^{-f}$, where $f:[0,1]\to\mathbb{R}$, $f(0)=0$, and $-1\leq f^{\prime}\leq 1$. Our goal is to develop a high-accuracy sampler—a sampler whose sampled distribution has $\delta$ error in total variation distance to $p$—in $\mathrm{polylog}(1/\delta)$ steps, using only queries to $f^{\prime}$.

我们用一个简单问题引出方法：从密度 $p\propto e^{-f}$ 采样，其中 $f:[0,1]\to\mathbb{R}$，$f(0)=0$，且 $-1\leq f^{\prime}\leq 1$。目标是仅用对 $f^{\prime}$ 的查询，在 $\mathrm{polylog}(1/\delta)$ 步内得到高精度采样器——其采样分布相对 $p$ 的总变差误差为 $\delta$。

Consider performing rejection sampling with the base measure $\mathsf{Unif}([0,1])$. To do so, we must generate randomness $b\sim\mathsf{Ber}(ce^{-f(x)})$ for any given $x\in[0,1]$. However, we only have access to $f^{\prime}$. The identity $f(x)=\int_{0}^{x}f^{\prime}(y)\,dy$, written as

考虑以基测度 $\mathsf{Unif}([0,1])$ 做拒绝采样。为此，对任意给定 $x\in[0,1]$，需生成随机性 $b\sim\mathsf{Ber}(ce^{-f(x)})$。但我们只有 $f^{\prime}$。恒等式 $f(x)=\int_{0}^{x}f^{\prime}(y)\,dy$ 可写成

$$\displaystyle f(x)=\En_{y\sim\mathsf{Unif}([0,x])}[xf^{\prime}(y)]\,, \tag{6}$$

suggests an unbiased estimate of $f(x)$. Is it possible to sample from $\mathsf{Ber}(ce^{-f(x)})$ if we have access to an unbiased estimate of $f(x)$? A general version of this question is known as the “Bernoulli factory” problem (26; 39). For our purposes, this can be stated as the following abstract task:

这提示了 $f(x)$ 的无偏估计。若有 $f(x)$ 的无偏估计，能否从 $\mathsf{Ber}(ce^{-f(x)})$ 采样？该问题的一般形式即「Bernoulli 工厂」问题（26；39）。对我们而言，可陈述为如下抽象任务：

Task: Given i.i.d. random variables $W_{1},W_{2},W_{3},\dotsc$ in $[-1,1]$, generate a sample $b\sim\mathsf{Ber}(ce^{\En W_{1}})$.

任务： 给定取值于 $[-1,1]$ 的 i.i.d. 随机变量 $W_{1},W_{2},W_{3},\dotsc$，生成样本 $b\sim\mathsf{Ber}(ce^{\En W_{1}})$。

To solve this, we write the Taylor series as

为此，把 Taylor 级数写成

$$\displaystyle e^{\En W_{1}}=e^{-1}\cdot e^{\En[1+W_{1}]}=\sum_{j\geq 0}\frac{e^{-1}}{j!}\,\bigl(\En[1+W_{1}]\bigr)^{j}\,.$$

Suppose that $J\sim\mathsf{Poisson}(2)$ is independent of the i.i.d. sequence $W_{1},W_{2},W_{3},\dotsc$. Then we notice that

设 $J\sim\mathsf{Poisson}(2)$ 独立于 i.i.d. 序列 $W_{1},W_{2},W_{3},\dotsc$。则注意到

$$\displaystyle e^{\En W_{1}}=e\En\Big[\prod_{j=1}^{J}\bigl(\frac{1+W_{j}}{2}\bigr)\Big]\,,$$

and simply set $b\sim\mathsf{Ber}\big(\prod_{j=1}^{J}\bigl(\frac{1+W_{j}}{2}\bigr)\big)$. Indeed, $\mathbb{P}(b=1)=\mathbb{E}\prod_{j=1}^{J}\bigl(\frac{1+W_{j}}{2}\bigr)=e^{-1+\En W_{1}}$.

并直接取 $b\sim\mathsf{Ber}\big(\prod_{j=1}^{J}\bigl(\frac{1+W_{j}}{2}\bigr)\big)$。的确，$\mathbb{P}(b=1)=\mathbb{E}\prod_{j=1}^{J}\bigl(\frac{1+W_{j}}{2}\bigr)=e^{-1+\En W_{1}}$。

In summary, we can generate a sample $b\sim\mathsf{Ber}(ce^{-f(x)})$ without having to compute or accurately approximate $f(x)$ (this requires the integration step $f(x)=\int_{0}^{x}f^{\prime}(y)\,dy$ and can be expensive). Instead, it is sufficient to have access to (a random number of) unbiased estimates of $f(x)$. The latter can be achieved with derivative information, in view of (6).

总之，我们可以生成样本 $b\sim\mathsf{Ber}(ce^{-f(x)})$，不必计算或高精度近似 $f(x)$（那需要积分步 $f(x)=\int_{0}^{x}f^{\prime}(y)\,dy$，可能很昂贵）。只需（随机次数的）$f(x)$ 无偏估计即可；由 (6)，这可用导数信息做到。

We now generalize this setup via the following meta-algorithm, called first-order rejection sampling (FORS). Given a proposal distribution $q$ and a tilt function $w$, the goal of Algorithm 1 is to produce a sample from $\widehat{p}(x)\propto q(x)\,e^{w(x)}$ without having access to the value $w(x)$. Instead, for each $x\in\mathbb{R}^{d}$, we can generate i.i.d. samples $W_{1},W_{2},W_{3}\dotsc$ such that $\En[W_{1}\mid x]=w(x)$. Let $\mathcal{W}_{x}$ denote the conditional distribution of $W_{1}$ given $x$.

我们把上述设定推广为如下元算法，称为一阶拒绝采样（FORS）。给定提议分布 $q$ 与倾斜函数 $w$，算法 1 的目标是在无法访问数值 $w(x)$ 的情况下，从 $\widehat{p}(x)\propto q(x)\,e^{w(x)}$ 产生样本。取而代之，对每个 $x\in\mathbb{R}^{d}$，我们可生成 i.i.d. 样本 $W_{1},W_{2},W_{3}\dotsc$，满足 $\En[W_{1}\mid x]=w(x)$。记 $\mathcal{W}_{x}$ 为给定 $x$ 时 $W_{1}$ 的条件分布。

Input: Parameter $B>0$, proposal distribution $q$ over $\mathbb{R}^{d}$, estimator distributions $(\mathcal{W}_{x})_{x\in\mathbb{R}^{d}}$ supported on $[-B,B]$
for $i=1,2,3,\dotsc$ do
Sample $x\sim q$.
Sample $J\sim\mathsf{Poisson}(2B)$.
Sample i.i.d. $W_{1},\dotsc,W_{J}\sim\mathcal{W}_{x}$.
Output $x$ with probability $\prod_{j=1}^{J}\frac{B+W_{j}}{2B}$.
end for

Algorithm 1 First-order rejection sampling (FORS)

输入： 参数 $B>0$，提议分布 $q$（定义在 $\mathbb{R}^{d}$ 上），估计量分布 $(\mathcal{W}_{x})_{x\in\mathbb{R}^{d}}$（支撑在 $[-B,B]$）
for $i=1,2,3,\dotsc$ do
从 $q$ 采样 $x$。
采样 $J\sim\mathsf{Poisson}(2B)$。
独立同分布采样 $W_{1},\dotsc,W_{J}\sim\mathcal{W}_{x}$。
以概率 $\prod_{j=1}^{J}\frac{B+W_{j}}{2B}$ 输出 $x$。
end for

算法 1 一阶拒绝采样（FORS）

Theorem 3.1(FORS guarantee).

Algorithm 1**outputs a random point with density $\widehat{p}(x)\propto q(x)\,e^{\En[W_{1}\mid x]}$. The number of sampled $W_{j}$’s is bounded, with probability at least $1-\delta$, by $3Be^{2B}\log(2/\delta)$.

定理 3.1（FORS 保证）。

算法 1**输出的随机点密度为 $\widehat{p}(x)\propto q(x)\,e^{\En[W_{1}\mid x]}$。以至少 $1-\delta$ 的概率，所采样的 $W_{j}$ 个数至多为 $3Be^{2B}\log(2/\delta)$。

Moreover, if Algorithm 1 is called $T$ times, then with probability at least $1-\delta$, the total number of sampled $W_{j}$’s is $O(Be^{2B}\,(T+\log(1/\delta)))$.

此外，若算法 1 被调用 $T$ 次，则以至少 $1-\delta$ 的概率，采样的 $W_{j}$ 总数为 $O(Be^{2B}\,(T+\log(1/\delta)))$。

3.2 Sampling from Gaussian tilts with an exact oracle

3.2 用精确 oracle 从高斯倾斜采样

To perform rejection sampling, we have to appropriately choose the proposal distribution $q$. By Taylor’s expansion, for a fixed $x_{+}\in\mathbb{R}^{d}$, we have

要做拒绝采样，需恰当选择提议分布 $q$。由 Taylor 展开，对固定的 $x_{+}\in\mathbb{R}^{d}$，有

$$\displaystyle f(x)\approx f(x_{+})+\langle x-x_{+},\nabla f(x_{+})\rangle\,,$$

and hence it is natural to choose $q=\mathsf{N}\left(x_{0}-\eta\nabla f(x_{+}),\eta I\right)$ so that

因此自然取 $q=\mathsf{N}\left(x_{0}-\eta\nabla f(x_{+}),\eta I\right)$，从而

$$\displaystyle q(x)\propto\exp\Bigl(-\langle x-x_{+},\nabla f(x_{+})\rangle-\frac{1}{2\eta}\left\lVert x-x_{0}\right\rVert^{2}\Bigr)\,.$$

Then, we can express

于是可写

$$\displaystyle\log\nu(x)-\log q(x)-\mathrm{const}= \\ \displaystyle~\langle x-x_{+},\nabla f(x_{+})\rangle-f(x)+f(x_{+}) \\ \displaystyle= \\ \displaystyle~\int_{0}^{1}\langle x-x_{+},\nabla f(x_{+})-\nabla f(rx+(1-r)x_{+})\rangle\,dr\,.$$

In particular, we can set

特别地，可取

$$\displaystyle W_{r,x}\vcentcolon=\langle x-x_{+},\nabla f(x_{+})-\nabla f(rx+(1-r)x_{+})\rangle\,,$$

so that

从而

$$\displaystyle\nu(x)\propto q(x)\exp\left(\En_{r}W_{r,x}\right)\,,$$

where the expectation $\En_{r}[\cdot]$ is taken over $r\sim\mathsf{Unif}([0,1])$.

其中期望 $\En_{r}[\cdot]$ 对 $r\sim\mathsf{Unif}([0,1])$ 取。

Further, assuming that $\nabla f$ is $\beta$-Lipschitz, we can bound $\lvert W_{r,x}\rvert\leq\beta\,\|x-x_{+}\|^{2}$. Therefore, as long as $x_{+}\approx x_{0}-\eta\nabla f(x_{+})$ and $\beta d\eta\ll 1$, we can guarantee that $\sup_{r}{\lvert W_{r,x}\rvert}\leq 1$ with high probability over $x\sim q$. This implies the following distribution is a good approximation of $\nu$:

进一步，若 $\nabla f$ 是 $\beta$-Lipschitz，则可界住 $\lvert W_{r,x}\rvert\leq\beta\,\|x-x_{+}\|^{2}$。因此，只要 $x_{+}\approx x_{0}-\eta\nabla f(x_{+})$ 且 $\beta d\eta\ll 1$，就能保证对 $x\sim q$ 以高概率有 $\sup_{r}{\lvert W_{r,x}\rvert}\leq 1$。这意味着下列分布是 $\nu$ 的良好近似：

$$\displaystyle\widehat{\nu}(x)\propto q(x)\exp(\En_{r}\widehat{W}_{r,x})\,,$$

where $\widehat{W}_{r,x}=\mathsf{Clip}_{B}(W_{r,x})$ and $B=\Theta(1)$ is a tuneable parameter.

其中 $\widehat{W}_{r,x}=\mathsf{Clip}_{B}(W_{r,x})$，$B=\Theta(1)$ 为可调参数。

General path integral

We can generalize the above pipeline with any path integral and demonstrate (a) that in the gradient Lipschitz case, we in fact only need to ensure $\beta\sqrt{d}\eta\ll 1$, and (b) the Lipschitz continuity of $\nabla f$ can be weakened to Hölder continuity for any exponent $s\in[0,1]$.

一般路径积分

可用任意路径积分推广上述流程，并证明：(a) 在梯度 Lipschitz 情形，实际上只需保证 $\beta\sqrt{d}\eta\ll 1$；(b) $\nabla f$ 的 Lipschitz 连续性可弱化为任意指数 $s\in[0,1]$ 的 Hölder 连续性。

Fix a distribution $P$ (to be determined later) over $\mathbb{R}^{d}$ and a path function $\gamma_{z,r}(x)\vcentcolon=\gamma(x;z,r)$ such that $\gamma_{z,1}(x)=x$ and $\gamma_{z,0}(x)=\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}(z)$ is independent of $x$. Then we can express any smooth function $h:\mathbb{R}^{d}\to\mathbb{R}$ as

固定分布 $P$（稍后确定）以及路径函数 $\gamma_{z,r}(x)\vcentcolon=\gamma(x;z,r)$，满足 $\gamma_{z,1}(x)=x$ 且 $\gamma_{z,0}(x)$ 与 $x$ 无关。则任意光滑函数 $h:\mathbb{R}^{d}\to\mathbb{R}$ 可写成

$$\displaystyle\begin{aligned} h(x)-\En_{z\sim P}[h(\macc@depth\char 1\relax\macc@set@skewchar\macc@nested@a 111{}(z))]=&\int_{0}^{1}\En_{z\sim P}\langle\dot{\gamma}{z,r}(x),\nabla h(\gamma{z,r}(x))\rangle,dr\
=&\En_{z\sim P,,r\sim\mathsf{Unif}([0,1])}\langle\dot{\gamma}{z,r}(x),\nabla h(\gamma{z,r}(x))\rangle,,\end{aligned} \tag{7}$$

where $\dot{\gamma}_{z,r}(x)=\frac{d}{dr}\gamma_{z,r}(x)$ is the path derivative with respect to $r$. This inspires us to consider

其中 $\dot{\gamma}_{z,r}(x)=\frac{d}{dr}\gamma_{z,r}(x)$ 为相对 $r$ 的路径导数。这启发我们考虑

$$\displaystyle W_{r,z,x}\vcentcolon=\langle\dot{\gamma}_{z,r}(x),\nabla f(x_{+})-\nabla f(\gamma_{z,r}(x))\rangle\,, \tag{8}$$

so that

从而

$$\displaystyle\nu(x)\propto q(x)\exp\left(\En_{r,z}W_{r,z,x}\right)\,,$$

where the expectation $\En_{r,z}[\cdot]$ is taken over $r\sim\mathsf{Unif}([0,1])$ and $z\sim P$. Again, we will truncate $W_{r,z,x}$ to ensure that the estimator lies in $[-B,B]$.

其中期望 $\En_{r,z}[\cdot]$ 对 $r\sim\mathsf{Unif}([0,1])$ 与 $z\sim P$ 取。同样，我们将截断 $W_{r,z,x}$，以保证估计量落在 $[-B,B]$。

For better dimension dependence, we consider the following path function ($\widehat{x}=x_{0}-\eta\nabla f(x_{+})$):

为获得更好的维数依赖，我们考虑如下路径函数（$\widehat{x}=x_{0}-\eta\nabla f(x_{+})$）：

$$\displaystyle\begin{aligned} &~\gamma_{z,r}(x)=a_{r}x+(1-a_{r})\widehat{x}+b_{r}z,,\
&~a_{r}=\sin(\pi r/2),,~~b_{r}=\cos(\pi r/2),,\end{aligned} \tag{9}$$

so that $\dot{\gamma}_{z,r}(x)=a_{r}^{\prime}(x-x_{0})+b_{r}^{\prime}z$.

从而 $\dot{\gamma}_{z,r}(x)=a_{r}^{\prime}(x-x_{0})+b_{r}^{\prime}z$。

Assumption 3.2.

There exists $s\in[0,1]$ and $\beta_{s}\geq 0$ such that $\left\lVert\nabla f(x)-\nabla f(y)\right\rVert\leq\beta_{s}\left\lVert x-y\right\rVert^{s}$ for all $x,y\in\mathbb{R}^{d}$.