[论文解读] Flow Annealed Importance Sampling Bootstrap

Laurence Illing Midgley, Vincent Stimper+3 more|arXiv (Cornell University)|Aug 3, 2022

Machine Learning in Materials Science被引用 14

一句话总结

该论文提出Flow AIS Bootstrap（FAB），一种新颖方法，通过将归一化流与冷凝重要性采样（AIS）结合，仅使用未归一化的目标密度训练流模型，避免了昂贵的分子动力学（MD）模拟。通过最小化α=2时的质量覆盖α-散度，FAB降低了重要性权重的方差，并实现了对多模态分布（包括丙氨酸二肽的玻尔兹曼分布）的精确近似——在仅使用MD样本进行最大似然训练的100分之一目标评估次数下，取得了更优结果。

ABSTRACT

Normalizing flows are tractable density models that can approximate complicated target distributions, e.g. Boltzmann distributions of physical systems. However, current methods for training flows either suffer from mode-seeking behavior, use samples from the target generated beforehand by expensive MCMC methods, or use stochastic losses that have high variance. To avoid these problems, we augment flows with annealed importance sampling (AIS) and minimize the mass-covering $α$-divergence with $α=2$, which minimizes importance weight variance. Our method, Flow AIS Bootstrap (FAB), uses AIS to generate samples in regions where the flow is a poor approximation of the target, facilitating the discovery of new modes. We apply FAB to multimodal targets and show that we can approximate them very accurately where previous methods fail. To the best of our knowledge, we are the first to learn the Boltzmann distribution of the alanine dipeptide molecule using only the unnormalized target density, without access to samples generated via Molecular Dynamics (MD) simulations: FAB produces better results than training via maximum likelihood on MD samples while using 100 times fewer target evaluations. After reweighting the samples, we obtain unbiased histograms of dihedral angles that are almost identical to the ground truth.

研究动机与目标

• 解决现有流模型训练方法依赖昂贵MD模拟或在多模态分布中出现模式搜寻行为的局限性。
• 开发一种训练目标，最小化重要性权重方差，同时实现对复杂、多模态目标分布的精确近似。
• 仅使用未归一化的目标密度，无需预先进行MD采样，即可训练玻尔兹曼生成器。
• 通过优先级重放缓冲区重用AIS样本，降低计算成本。
• 在具有挑战性的基准测试中（包括丙氨酸二肽分子）展示最先进性能。

提出的方法

• FAB使用α=2的α-散度作为训练目标，该目标具有质量覆盖性，且在推理过程中最小化重要性权重的方差。
• 采用冷凝重要性采样（AIS），将流模型作为初始分布，目标分布作为方差最小的分布，以估计α-散度损失。
• AIS能够从流模型近似效果较差的区域生成高质量样本，从而发现新模式并提升训练信号质量。
• 引入优先级重放缓冲区以重用AIS样本，降低计算成本并提高样本效率。
• 通过在AIS估计的α=2散度目标上使用随机梯度下降，端到端训练流模型。
• 推理时，使用训练好的流模型进行重要性采样，可实现对目标分布下期望值的低方差估计。

实验结果

研究问题

• RQ1 能否仅使用未归一化的目标密度，无需依赖昂贵的MD模拟，有效训练归一化流？
• RQ2 与标准KL散度相比，最小化α=2的α-散度是否能带来更优的样本质量与更低的重要权重方差？
• RQ3 FAB能否在高度多模态分布（如丙氨酸二肽的拉氏图）中发现并准确建模所有模式？
• RQ4 与基于MD生成样本的最大似然训练相比，FAB在性能和样本效率方面表现如何？
• RQ5 使用重放缓冲区是否能显著减少训练所需的目标密度评估次数？

主要发现

• FAB仅使用未归一化的能量函数，无需任何MD样本，成功学习了丙氨酸二肽的玻尔兹曼分布，达到最先进水平。
• 该方法在仅使用最大似然训练MD样本的1/100目标评估次数下，对真实分布的近似效果更优。
• 重加权后，FAB生成的二面角直方图与真实值几乎完全一致，表明其在模式覆盖方面具有高度保真度。
• 在32维的“Many Well”问题中，FAB在有效样本量（ESS）、对数似然和KL散度方面均优于竞争方法。
• 在α=2时，FAB在拉氏图上实现了最低的KL散度（2.51×10⁻³）和最高的ESS（52.2%），显著优于其他方法。
• 使用重放缓冲区显著降低了计算成本并提升了训练稳定性，实现了对高质量AIS样本的高效重用。