[论文解读] Annealed Flow Transport Monte Carlo

Michael Arbel, Alexander Matthews+1 more|arXiv (Cornell University)|Feb 15, 2021

Model Reduction and Neural Networks参考文献 73被引用 4

一句话总结

该论文提出了一种新型马尔可夫链蒙特卡洛方法——退火流传输（Annealed Flow Transport, AFT），结合序贯蒙特卡洛（SMC）与归一化流（NFs），以提升在复杂目标分布下对归一化常数和期望的估计性能。通过依次学习归一化流，将粒子逐步传输至后续的退火目标，AFT 有效降低了方差并缓解了模式崩溃问题，在更少温度步数下实现了优于标准 SMC 的性能，并在连续时间极限下展现出向受控扩散过程收敛的理论性质。

ABSTRACT

Annealed Importance Sampling (AIS) and its Sequential Monte Carlo (SMC) extensions are state-of-the-art methods for estimating normalizing constants of probability distributions. We propose here a novel Monte Carlo algorithm, Annealed Flow Transport (AFT), that builds upon AIS and SMC and combines them with normalizing flows (NFs) for improved performance. This method transports a set of particles using not only importance sampling (IS), Markov chain Monte Carlo (MCMC) and resampling steps - as in SMC, but also relies on NFs which are learned sequentially to push particles towards the successive annealed targets. We provide limit theorems for the resulting Monte Carlo estimates of the normalizing constant and expectations with respect to the target distribution. Additionally, we show that a continuous-time scaling limit of the population version of AFT is given by a Feynman--Kac measure which simplifies to the law of a controlled diffusion for expressive NFs. We demonstrate experimentally the benefits and limitations of our methodology on a variety of applications.

研究动机与目标

• 解决在目标分布间存在显著差异时，退火重要性采样（AIS）与 SMC 中重要性权重方差过高的问题。
• 通过在中间退火步骤而非直接从简单先验学习流映射，克服基于归一化流的传输映射中常见的模式崩溃与尾部低估问题。
• 结合 SMC（重采样、MCMC、重要性采样）与表达能力强的归一化流，实现更精确且稳定的推断。
• 为所提算法下的蒙特卡洛估计建立理论保证，包括弱大数定律与中心极限定理。
• 证明在连续时间极限下，AFT 收敛至一个费曼-卡茨测度，当流具有足够表达能力时，该测度简化为受控扩散过程。

提出的方法

• AFT 使用一系列退火目标 πₖ ∝ π₀^(1−βₖ)π_K^βₖ，其中 0 = β₀ < β₁ < ... < β_K = 1，从一个简单的初始分布 π₀ 逐步过渡至目标分布 π_K = π。
• 在每一步 k，训练一个归一化流 T_{k+1}，以最小化 (T_{k+1})#πₖ 与 π_{k+1} 之间的 KL 散度，从而将粒子从 πₖ 映射至更接近 π_{k+1} 的分布。
• 流映射完成后，利用 π_{k+1}/πₖ 的比值对粒子进行重新加权，形成对 π_{k+1} 的重要性采样近似。
• 随后对粒子进行重采样，并通过以 π_{k+1} 为不变分布的 MCMC 核心进行扰动，以维持多样性并校正流映射的近似误差。
• 通过使用未调整的朗之万算法（ULA）核，将算法推广至连续时间尺度极限，证明其收敛至一个费曼-卡茨测度。
• 当归一化流具备足够表达能力时，该极限测度对应于受控扩散过程的分布。

实验结果

研究问题

• RQ1 归一化流能否有效用于提升 SMC 采样器的效率与准确性，从而降低重要性权重的方差？
• RQ2 在中间退火步骤而非直接从先验学习传输映射，是否能缓解模式崩溃并提升采样质量？
• RQ3 在基于学习的传输方案下，SMC 中蒙特卡洛估计器可建立何种理论保证？
• RQ4 AFT 的连续时间极限如何与受控扩散等随机过程相关联？
• RQ5 AFT 是否能在更少温度步数下，实现比标准 SMC 更高的估计精度与更快的收敛速度？

主要发现

• 在基准问题上，AFT 以更少的温度步数实现了比标准 SMC 更低的估计误差，展现出更高的样本效率。
• 理论分析表明，尽管传输映射是学习得到的，AFT 估计器仍满足弱大数定律与中心极限定理。
• 在使用 ULA 核且不进行重采样的连续时间极限下，AFT 收敛至一个费曼-卡茨测度；当流具有足够表达能力时，该测度简化为受控扩散过程的分布。
• 在二值化 MNIST 与空间 Cox 过程上的实验表明，AFT 提供了比基线 SMC 和变分推断更精确的归一化常数估计与更优的后验近似。
• 通过利用中间退火目标，该方法有效缓解了直接训练归一化流时常见的模式崩溃与尾部低估问题。
• 消融实验表明，每步转换中足够的 HMC 更新次数对 AFT 性能至关重要，尤其在高维或复杂目标分布下。