[论文解读] Sampling, Diffusions, and Stochastic Localization

Andrea Montanari|arXiv (Cornell University)|May 18, 2023

Markov Chains and Monte Carlo Methods被引用 4

一句话总结

本文建立了去噪扩散模型与随机定位之间的理论联系，表明标准去噪扩散模型是随机定位过程的一个特例。它将随机定位推广至更广泛的设置，展示了该视角如何简化模型学习，并在高斯模型中通过线性估计近似证明了对目标分布的收敛性。

ABSTRACT

Diffusions are a successful technique to sample from high-dimensional distributions. The target distribution can be either explicitly given or learnt from a collection of samples. They implement a diffusion process whose endpoint is a sample from the target distribution. The drift of the diffusion process is typically represented as a neural network. Stochastic localization is a successful technique to prove mixing of Markov Chains and other functional inequalities in high dimension. An algorithmic version of stochastic localization was recently proposed in order to sample from certain statistical mechanics models. This expository article has three objectives: $(i)$~Generalize the algorithmic construction to other stochastic localization processes. This construction is both simple and broadly applicable; $(ii)$~Clarify the connection between diffusions and stochastic localization. This allows to derive several known sampling schemes in a unified fashion; $(iii)$~Describe the insights that follow from this unified viewpoint.

研究动机与目标

• 将[EAMS22]中的算法化随机定位框架推广至更广泛的随机过程类别。
• 阐明基于扩散的采样与随机定位之间的理论关系，表明标准去噪扩散模型是后者的一个特例。
• 展示随机定位视角如何简化去噪扩散模型中漂移函数的学习。
• 通过随机定位的视角，为采样算法的设计与分析提供新见解。
• 在高维高斯模型中，为基于线性估计器的该方法实例建立收敛性保证。

提出的方法

• 将随机定位过程的框架推广至原始框架之外，允许使用时变且不可逆的动力学。
• 推导扩散过程的倒向时间SDE，以从已知初始分布构建通往目标测度的采样路径。
• 引入一个连续时间随机过程，其漂移为给定当前状态时目标变量线性变换的条件期望。
• 采用线性、平移不变的估计器结构来近似最优漂移，利用循环矩阵结构实现计算效率。
• 对最优估计器应用有限窗口卷积近似，降低复杂度，同时保持渐近一致性。
• 采用Kullback-Leibler散度作为损失函数，指导随机定位框架中漂移分量的学习。

实验结果

研究问题

• RQ1 随机定位框架能否超越原始构造，扩展至包含更广泛类别的扩散过程？
• RQ2 标准去噪扩散模型与随机定位过程之间有何关系？
• RQ3 随机定位视角能否通过降低所学习漂移函数的复杂度，简化基于得分的生成模型的训练？
• RQ4 在高维高斯模型中，基于线性估计器的随机定位过程的收敛行为如何？
• RQ5 使用有限窗口卷积近似最优条件期望是否会保持渐近收敛至目标分布？

主要发现

• 本文证明了标准去噪扩散模型是随机定位过程的一个特例，建立了两种主要采样范式之间的正式理论联系。
• 在高斯模型中，采用有限窗口卷积近似的线性、平移不变估计器可实现对目标分布的渐近收敛。
• 当时间 t 趋近于无穷时，目标分布与生成分布之间的极限Wasserstein距离和总变差距离均收敛至零。
• 在高斯设定下，该漂移分量的估计器在 O(1/n) 项内为贝叶斯最优，从而为其作为合理近似提供了理论依据。
• 该方法通过用具有循环结构的简单线性函数近似最优漂移，显著简化了学习任务，降低了对复杂神经网络架构的依赖。
• 该框架通过统一扩散与定位视角，使得能够设计出更具可解释性与可分析性的采样算法。