HyperAIHyperAI

Command Palette

Search for a command to run...

1 天前
扩散模型

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

Fan Chen Sinho Chewi Constantinos Daskalakis Alexander Rakhlin

摘要

本文提出一种扩散模型采样算法,在获得 O~(δ)\widetilde O(\delta)O(δ) 精度的 L2L^2L2 分数估计的条件下,可在 polylog(1/δ)\mathrm{polylog}(1/\delta)polylog(1/δ) 步内达到 δ\deltaδ 误差,这是相对此前所有结果的指数级改进。具体而言,在最小数据假设下,复杂度为 O~(dpolylog(1/δ))\widetilde{O}(d_{\star} \mathrm{polylog}(1/\delta))O(dpolylog(1/δ)),其中 dd_{\star}d 是数据的内禀维度。进一步地,在非均匀 L-Lipschitz 条件下,复杂度降低至 O~(Lpolylog(1/δ))\widetilde{\cal O}(L \mathrm{polylog}(1/\delta))O(Lpolylog(1/δ))。我们的方法还首次实现了仅用梯度评估即可达到 polylog(1/δ) 复杂度的通用对数凹分布采样器。

一句话总结

麻省理工学院与耶鲁大学的研究人员提出了一种扩散采样方法,该方法在 O~(δ)\widetilde O ( \delta )O(δ) 精度的 L2L ^ { 2 }L2 分数估计下,仅需 polylog(1/δ)\mathrm { p o l y l o g } ( 1 / \delta )polylog(1/δ) 步即可达到 δ\deltaδ-误差,相较于所有先前结果实现了指数级加速;在非均匀 LLL-Lipschitz 条件下,其复杂度为 O~(dpolylog(1/δ))\widetilde { O } ( d _ { \star } \mathrm { p o l y l o g } ( 1 / \delta ) )O(dpolylog(1/δ))O~(Lpolylog(1/δ))\widetilde { \cal O } ( L \mathrm { p o l y l o g } ( 1 / \delta ) )O(Lpolylog(1/δ)),并进一步通过梯度评估首次实现了针对对数凹分布的 polylog(1/δ)\mathrm { p o l y l o g } ( 1 / \delta )polylog(1/δ) 采样器。

核心贡献

  • 给出在 L2L^2L2 下精度为 O~(δ)\widetilde{O}(\delta)O(δ) 的分数估计时,采样算法在最小数据假设下达到 O~(dpolylog(1/δ))\widetilde{O}(d_\star \operatorname{polylog}(1/\delta))O(dpolylog(1/δ)) 的复杂度,其中 dd_\stard 为本征维度;这在与准确率相关的依赖上相较先前方法实现了指数级改进。
  • 在非均匀 LLL-Lipschitz 分数条件下,复杂度降至 O~(Lpolylog(1/δ))\widetilde{O}(L \operatorname{polylog}(1/\delta))O(Lpolylog(1/δ)),用 Lipschitz 常数替代了维度依赖。
  • 该框架仅通过梯度评估,便首次实现了针对一般对数凹分布、复杂度为 polylog(1/δ)\operatorname{polylog}(1/\delta)polylog(1/δ) 的采样器。

引言

作者们处理了基于扩散的生成建模中的一个核心问题:能否仅通过梯度(分数)评估,在无法获取密度的情况下,以高精度——即仅需与目标误差倒数成多对数关系的步数——完成采样?标准的高精度采样器依赖需要密度评估的接受/拒绝方法,而仅基于分数的随机过程离散化会因偏差而迫使复杂度与 1/δ1/\delta1/δ 成多项式关系。现有的扩散采样器即使采用高阶离散化,最多也只能达到次多项式但仍有界于多项式的查询复杂度。作者通过引入一阶拒绝采样(FORS)突破了这一障碍,FORS 是一种仅通过梯度查询来模拟拒绝采样的元算法。应用于扩散模型时,FORS 在最小数据假设下实现 O(dlog3((d+M22)/δ))O(\mathsf{d}_\star \log^3((d + \mathsf{M}_2^2)/\delta))O(dlog3((d+M22)/δ)) 的查询复杂度,并在 Lipschitz 分数条件下给出近乎无维度的保证,这在与准确性参数的依赖上相较先前工作实现了指数级改进。

方法

作者们开发了一个统一的高精度采样框架,该框架仅使用一阶(梯度)查询。其核心构件是一个名为 一阶拒绝采样(FORS) 的新颖子程序,它允许从倾斜分布 p^(x)q(x)ew(x)\widehat{p}(x) \propto q(x) e^{w(x)}p(x)q(x)ew(x) 中精确采样,而无需计算对数密度 w(x)w(x)w(x)。相反,对于每个候选 xxx,FORS 抽取独立同分布的无偏估计 W1,W2,W_1, W_2, \dotsW1,W2, 使得 E[W1x]=w(x)\mathbb{E}[W_1 \mid x] = w(x)E[W1x]=w(x),并利用伯努利工厂技术以正确概率接受或拒绝 xxx。具体地,记 ew(x)=e1eE[1+W1]e^{w(x)} = e^{-1} e^{\mathbb{E}[1 + W_1]}ew(x)=e1eE[1+W1] 并采样泊松随机变量 JJJ,接受概率即可恢复为 Ej=1J1+Wj2\mathbb{E}\prod_{j=1}^J \frac{1+W_j}{2}Ej=1J21+Wj。每次接受的梯度查询次数以高概率有界,从而在估计值适当有界时使该过程高效。

该 FORS 子程序随后被特化到形如

ν(x)exp(f(x)xx022η),\nu(x) \propto \exp\Bigl(-f(x) - \frac{\|x - x_0\|^2}{2\eta}\Bigr),ν(x)exp(f(x)2ηxx02),

高斯倾斜分布,这类分布在扩散模型采样和对数凹采样中自然出现。关键思路是选取高斯提议分布 q=N(x0ηf(x+),ηI)q = \mathcal{N}(x_0 - \eta \nabla f(x_+), \eta \mathbf{I})q=N(x0ηf(x+),ηI),通过 fff 的局部线性化来近似 ν\nuν。路径积分表示将对数比 logν(x)logq(x)\log \nu(x) - \log q(x)logν(x)logq(x) 表达为一个期望:

w(x)=ErUnif,zPγ˙z,r(x),f(x+)f(γz,r(x)),w(x) = \mathbb{E}_{r \sim \text{Unif}, z \sim P} \bigl\langle \dot{\gamma}_{z,r}(x),\, \nabla f(x_+) - \nabla f(\gamma_{z,r}(x)) \bigr\rangle,w(x)=ErUnif,zPγ˙z,r(x),f(x+)f(γz,r(x)),

其中 γz,r(x)\gamma_{z,r}(x)γz,r(x) 是精心设计的曲线,满足 γz,1(x)=x\gamma_{z,1}(x)=xγz,1(x)=xγz,0(x)=γˉ(z)\gamma_{z,0}(x)=\bar{\gamma}(z)γz,0(x)=γˉ(z)xxx 无关。在 Hölder 光滑性假设 f(x)f(y)βsxys\|\nabla f(x)-\nabla f(y)\| \le \beta_s \|x-y\|^s∥∇f(x)f(y)βsxys 下,估计量被截断到常值区间 [B,B][-B,B][B,B],定理 3.3 保证只要步长 η\etaη 选择得当(例如在光滑情形下 η1β1dlog(1/δ)\eta^{-1} \gg \beta_1 \sqrt{d} \log(1/\delta)η1β1dlog(1/δ)),所得分布 ν^\widehat{\nu}νχ2\chi^2χ2 散度下与真实倾斜分布十分接近。

该框架随后被提升至 扩散采样。扩散过程的反向转移核可以写为高斯倾斜形式,因此每个逆向步骤成为利用分数估计 sk\mathsf{s}_ksk 的 FORS 应用。提议分布选为

ρˉk(Xk+1)=N(αk1Xk+1+αkηksk+1(Xk+1),ηˉkI),\bar{\rho}_k(\cdot \mid X_{k+1}) = \mathcal{N}\bigl(\alpha_k^{-1} X_{k+1} + \alpha_k \eta_k \mathsf{s}_{k+1}(X_{k+1}),\, \bar{\eta}_k \mathbf{I}\bigr),ρˉk(Xk+1)=N(αk1Xk+1+αkηksk+1(Xk+1),ηˉkI),

它对应指数积分近似,且几乎是真实转移的 KL 最小化子。提供无偏估计的“校正”分布由一个路径函数 γz,r,x^(x)\gamma_{z,r,\widehat{x}}(x)γz,r,x(x) 构建,其系数 ar,bra_r,b_rar,br 满足范数保持恒等式,从而保证了无维度依赖的方差界。

一个核心概念是数据分布的 本征维度

dimσ2(p)=1infr0(logN(p;r)+r2σ2)d,\dim_{\sigma^2}(p) = 1 \vee \inf_{r\ge 0} \Bigl(\log N(p;r) + \frac{r^2}{\sigma^2}\Bigr) \wedge d,dimσ2(p)=1r0inf(logN(p;r)+σ2r2)d,

并定义 d=dimσ02/α02(pdata)\mathsf{d}_\star = \dim_{\sigma_0^2/\alpha_0^2}(p_{\mathsf{data}})d=dimσ02/α02(pdata)。该量捕捉低维结构、小支撑集或有界半径等特性,且永不大于环境维度 ddd。借助 dd_\stard,步长条件变为 σk2/ηkdlog(1/δ)+log2(1/δ)\sigma_k^2/\eta_k \gg \mathsf{d}_\star \log(1/\delta) + \log^2(1/\delta)σk2/ηkdlog(1/δ)+log2(1/δ),从而在有界 Lipschitz 度量下得到扩散采样复杂度量级

dlog3 ⁣(d+M22δ2)\mathsf{d}_\star \cdot \log^3\!\Bigl(\frac{d + \mathsf{M}_2^2}{\delta^2}\Bigr)dlog3(δ2d+M22)

当数据分布为对数光滑时,同一方案用 ddd 替代 d\mathsf{d}_\stard 可给出 KL 保证。

进一步的改进利用 非均匀 Lipschitz 条件。在分数雅可比矩阵 mτ(Yτ)\nabla m_\tau(Y_\tau)mτ(Yτ) 以高概率具有有界算子范数的假设下(该假设无条件成立,且 Lop,δ=O(d+log(1/δ))L_{\mathrm{op},\delta} = O(\mathsf{d}_\star + \log(1/\delta))Lop,δ=O(d+log(1/δ)),对于对数凹或混合模型变为常数),有效 Frobenius 范数界 LF,δL_{\mathrm{F},\delta}LF,δ 可能远小于 d\sqrt{d}d。这将所需步长条件缩减为 σk2/ηkLF,δlog(d/δ)+log2(1/δ)\sigma_k^2/\eta_k \gg L_{\mathrm{F},\delta} \log(\mathsf{d}_\star/\delta) + \log^2(1/\delta)σk2/ηkLF,δlog(d/δ)+log2(1/δ),得到低至 min{dLop,d2/3Lop1/3}\min\{\sqrt{d L_{\mathrm{op}}},\, \mathsf{d}_\star^{2/3} L_{\mathrm{op}}^{1/3}\}min{dLop,d2/3Lop1/3} 乘以多对数因子的复杂度。分析还验证了相同的 Lipschitz 参数控制着一步 DDPM 分布下分数的光滑性,从而闭合了归纳论证的循环。

最后,作者将同一高斯倾斜技术通过近端采样器应用于 对数凹采样。吉布斯步骤中出现的受限高斯谕示(RGO)恰好是一个倾斜分布,因此 FORS 用一阶查询替代了零阶查询。在对数凹性、强对数凹性或等周假设下,近端采样器外循环与 FORS 结合可提供高精度保证,其复杂度仅在平滑势函数时以 d\sqrt{d}d 随维度增长,或在 Lipschitz 情形下与维度无关。这在对数凹采样中统一了扩散模型与经典方法的处理,仅依靠一阶查询框架。

实验

实验部分将受限高斯谕示通过 FORS 实现的近端采样器应用于从对数凹和等周目标分布中仅使用一阶查询进行采样。对于光滑势函数(s=1),在对数 Sobolev、Poincaré 或对数凹条件下,在卡方散度或 KL 散度意义下获得高精度保证,查询复杂度随条件数、维度(通常为 d1/2d^{1/2}d1/2)以及初始分布距离多项式增长。对于 Lipschitz 势函数(s=0),在 Poincaré 和对数凹条件下导出类似保证,无需零阶查询且避免离散化误差劣化。总体而言,结果展示了一阶采样能够在不依赖扩散离散化或零阶谕示的限制下实现高精度。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供