一句话总结

扩散模型采样算法可在 $L^{2}$L2 中的 score 估计达到 $\widetilde{O}(\delta)$O(δ) 精度时，于 $\text{polylog}(1/\delta)$polylog(1/δ) 步内实现 $\delta$δ-error，相对于在内在维度 $d_{\star}$d⋆​ 下具有复杂度 $\widetilde{O}(d_{\star}\text{polylog}(1/\delta))$O(d⋆​polylog(1/δ)) 的先前工作带来指数级改进，并且提供了第一个仅使用梯度评估、针对一般 log-concave 分布的 $\text{polylog}(1/\delta)$polylog(1/δ) 复杂度采样器。

核心贡献

• 本工作给出扩散采样算法，在给定 $L^{2}$L2 中 $\widetilde{O}(\delta)$O(δ) 精度的 score 估计时，可在 $\operatorname{polylog}(1/\delta)$polylog(1/δ) 步内达到 $\delta$δ-error，在最弱数据假设下得到 $\widetilde{O}(d\operatorname{polylog}(1/\delta))$O(dpolylog(1/δ)) 复杂度，并且相对先前结果对目标精度依赖实现了指数级改进。

• 该方法从 Bayes-rule 扩散转移中采样，而不是仅依赖 Gaussian DDPM 更新。最终保证对 score 估计误差具有鲁棒性，并在非均匀 $L$L-Lipschitz 条件下给出 $\widetilde{O}(\sqrt{dL}\operatorname{polylog}(1/\delta))$O(dL​polylog(1/δ)) 复杂度；当数据具有内在维度 $d_{\star}$d⋆​ 时，给出 $\widetilde{O}(d_{\star}\operatorname{polylog}(1/\delta))$O(d⋆​polylog(1/δ)) 复杂度。

• 该框架还为一般 log-concave 分布给出第一个仅使用梯度评估的 $\operatorname{polylog}(1/\delta)$polylog(1/δ) 复杂度采样器。

引言

作者研究在只能获得 log-density 梯度评估的条件下从连续分布中采样的复杂度。这一设置对于基于扩散的生成模型至关重要，因为这类模型学习 score 函数而不是未归一化 density。这一点很重要，因为理想情况下，快速收敛意味着迭代次数随目标精度以 polylog(1/delta) 缩放，从而能够在不进行过多 score 查询的情况下实现高精度生成。先前工作面临一个关键限制：仅使用 score 的扩散采样器通常依赖随机微分方程离散化，而离散化误差会阻碍高精度保证，并导致对 1/delta 的多项式或仅适度亚多项式依赖；现有的 polylog 型结果则通常要求密度评估或更强的假设。作者通过引入 first-order rejection sampling 来解决这一问题，这是一种仅使用一阶查询来模拟拒绝采样的元算法，并证明该算法在最弱数据假设下为扩散采样提供高精度 polylog 级保证，同时针对 Lipschitz score、低内在维数和 log-concave 采样给出进一步改进。

方法

该方法围绕一个可复用的采样原语组织，该原语用于在无法直接访问 density 值的情况下从倾斜 Gaussian target 中采样。核心目标具有以下形式

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

其中算法只允许查询关于 $f$f 的一阶信息。在扩散采样场景中，可以选择 $f=-\log p_t$f=−logpt​，并用 score 估计替换 $-\nabla f$−∇f。主要挑战是针对 Gaussian 提议分布执行拒绝采样，同时避免对 $f$f 自身进行昂贵的积分或评估。

作者提出一个 first-order rejection sampling 模块。抽象任务是从一个目标中采样

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

其中 $q$q 是提议分布，$w(x)$w(x) 不能直接评估。取而代之，对于每个 $x$x，采样器可以抽取随机变量 $W_1,W_2,\ldots$W1​,W2​,…，使得

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

关键观察是，$w(x)$w(x) 的无偏估计就足以模拟所需的 Bernoulli 接受概率。这是一种 Bernoulli factory 构造。如果随机变量在截断后有界，例如在 $[-B,B]$[−B,B] 内，则接受步骤可以使用缩放后的无偏估计量的 Poissonized product 来执行。最终采样器输出一个分布近似正比于以下形式的点

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

并且所抽取估计器的期望数量按以下量级增长

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

在 $T$T 次调用内，在常数因子意义下，至少以 $1-\delta$1−δ 的概率成立。

对于 Gaussian tilt，该 first-order rejection sampling 模块可以通过从 $f$f 的局部线性近似构造 Gaussian 提议分布来具体化。选取点 $x_+$x+​ 并构造

$$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).

等价地，如果

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

则提议分布可取为

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

$\nu$ν 与 $q$q 之间的残差 log tilt 可以写成梯度差的 path integral。最简单的形式是

$$\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.

这提示无偏估计量为

$$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+​)⟩,

其中 $r\sim\mathrm{Unif}(0,1)$r∼Unif(0,1)。实践中，该估计量被截断到一个固定区间：

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

其中 $B=\Theta(1)$B=Θ(1)。如果 $\nabla f$∇f 为 $\beta$β-Lipschitz，则

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

因此，当 Gaussian 提议在锚点附近具有良好中心化且步长足够小时，截断后的估计量以高概率有界。

更精细的构造使用一般的路径积分。固定 $z$z 的随机偏移分布和一条满足 $\gamma_{z,1}(x)=x$γz,1​(x)=x 且 $\gamma_{z,0}(x)$γz,0​(x) 不依赖于 $x$x 的路径 $\gamma_{z,r}(x)$γz,r​(x)。则对于光滑函数 $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))⟩,

其中 $\dot\gamma_{z,r}(x)=\frac{d}{dr}\gamma_{z,r}(x)$γ˙​z,r​(x)=drd​γz,r​(x)。由此得到路径估计量

$$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))⟩.

一个有用的选择是

$$\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,

其中

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

并且 $z$z 从 Gaussian 噪声分布中采样。该路径在提议均值与当前候选之间平滑插值，同时加入来自 $z$z 的随机性。在 $\nabla f$∇f 的梯度 Lipschitz 或 Hölder 连续性假设下，该路径估计量得到一个 Gaussian tilt 采样器，其样本分布在 chi squared divergence 意义下接近目标。在光滑情形下，步长必须满足形如以下条件的条件

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

而仅在 Lipschitz 情形下，条件变为

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

随后，同一子程序被用作扩散采样的引擎。关键结构洞察是，一个反向扩散步骤可被视为从 Gaussian tilted target 中采样，其中 tilt 由 score 函数确定。在步骤 $t$t，设前向过程的方差为 $\sigma_t^2$σt2​，信号缩放为 $\alpha_t$αt​。给定状态 $X_{t+1}$Xt+1​，定义确定性锚点

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

随后查询学习到的或估计的 score 函数 $\widehat{\mathsf s}_t$st​，以构造反向转移目标的估计器。采样器在每个扩散时间反复运行 Gaussian tilt 拒绝子程序，并从 $t=T$t=T 向 $t=1$t=1 推进过程中累积样本质量。

第一种变体是简单的 Gaussian 提议方法。它使用

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

以及截断路径估计量

$$\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​)⟩),

其中 $r\sim\mathrm{Unif}(0,1)$r∼Unif(0,1)。步长被选择为使得

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

最终采样器跟踪从先验分布朝早期停止分布 $p_1$p1​ 的反向转移核。其误差分解具有如下形式

$$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​,

其中 $\varepsilon_{t,\mathrm{score}}^2$εt,score2​ 度量局部 score 近似误差。除二阶矩界之外，该方法不需要对数据分布施加较强的形状假设。在适当 schedule 下，在 bounded Lipschitz metric 中达到小差异所需的步数按以下量级增长

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

其中

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

第二种具体化是类似 DDPM 的方法。这里提议分布包含一阶 score 修正：

$$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).

该提议更接近标准 ancestral 扩散采样，但目标修正仍通过拒绝估计器实现。估计器中使用的路径为

$$\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,

其中

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

且

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

估计量为

$$\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​)⟩).

该版本受益于真实 score 的非均匀 smoothness 条件。在该条件下，步长可选择为满足

$$\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/δ),

其中 $L_\delta$Lδ​ 是高概率 smoothness 参数。对于 Gaussian 数据，可以取 $L_\delta$Lδ​ 为常数；而对于 Gaussian 混合分布，它随环境维度和 mixture component 数量对数增长。最终总查询复杂度按以下量级增长

$$\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​​).

当数据分布具有低维结构时，该方法也可以适配。动机是 Gaussian 平滑后的 log density 的 Hessian 包含可能主导 smoothness 参数的各向同性项。如果

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

则

$$\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)​.

偏移后的 score 函数

$$\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​,

去除了 curvatures 的各向同性部分，并且在 $p_{\mathrm{data}}$pdata​ 的 support 本质低维时往往表现更好。

内在维度通过覆盖数定义。如果 $N_{\mathrm{data}}(r)$Ndata​(r) 是数据 support 在半径 $r$r 处的覆盖数，则

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

低维适配通过以下形式修改 Gaussian 提议方差

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

并在拒绝估计器中使用偏移后的 score 差异：

$$\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​)⟩).

相应的步长条件变为

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

最终复杂度依赖于 $\mathsf d_\star$d⋆​，而不是环境维度 $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​​).

同一 Gaussian tilt 采样器还为 log-concave 分布的 proximal sampler 中使用的 restricted Gaussian oracle 提供了仅依赖梯度的实现。对于目标 $\mu\propto e^{-f}$μ∝e−f，proximal sampler 使用形如以下形式的增广密度

$$\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).

一个 Gibbs 步骤需要从 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).

该 oracle 恰好是以 $y$y 为中心的 Gaussian tilt。通过用 first-order rejection sampling 模块实现该 oracle，proximal sampler 避免直接评估 $f$f，转而依赖梯度信息。在 smoothness、log-Sobolev、Poincaré 或 log-concavity 假设下，这得到具有 first-order query complexity 的高精度 proximal sampling，相比标准基于拒绝的 Gibbs 方案还改善了维度依赖。

实验

该研究将对 log-concave 和 isoperimetric 密度的 proximal sampling 应用到所提出方法，通过对 restricted Gaussian oracle 使用 Gaussian tilt 表示。结果表明，在 log-Sobolev、Poincare 和 log-concavity 假设下，该方法为 smooth 和 Lipschitz objectives 都提供高精度采样保证，同时仅依赖一阶和 proximal 查询。这些结果说明，该方法无需零阶信息，也能实现强收敛，并避免基于扩散的离散化方法中常见的精度退化。