Iterated Energy-based Flow Matching for Sampling from Boltzmann Densities

📄 arXiv: 2408.16249v1 📥 PDF

作者: Dongyeop Woo, Sungsoo Ahn

分类: cs.LG, stat.ML

发布日期: 2024-08-29

💡 一句话要点

提出迭代能量基流匹配(iEFM)，用于从玻尔兹曼密度中采样。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 能量基模型 连续归一化流 蒙特卡罗估计 非归一化密度 概率建模

📋 核心要点

  1. 现有方法难以直接从能量函数或非归一化密度中训练生成模型，尤其是在高维复杂系统中。
  2. iEFM通过训练模型预测基于能量函数的边际向量场的蒙特卡罗估计，实现了从非归一化密度中进行离策略训练。
  3. 实验表明，iEFM在GMM和DW-4能量函数上优于现有方法，证明了其在高维系统中的有效性。

📝 摘要（中文）

本文研究了从能量函数或非归一化密度的评估中训练生成器的问题。这是概率推理中的一个基本问题，对于科学应用至关重要，例如学习分子的3D坐标分布。为了解决这个问题，我们提出了迭代能量基流匹配（iEFM），这是第一个用于从非归一化密度训练连续归一化流（CNF）模型的离策略方法。我们引入了无模拟的能量基流匹配目标，该目标训练模型来预测从已知能量函数构建的边际向量场的蒙特卡罗估计。我们的框架是通用的，可以扩展到方差爆炸（VE）和最优传输（OT）条件概率路径。我们在二维高斯混合模型（GMM）和八维四粒子双阱势（DW-4）能量函数上评估了iEFM。我们的结果表明，iEFM优于现有方法，展示了其在复杂高维系统中进行高效且可扩展的概率建模的潜力。

🔬 方法详解

问题定义：论文旨在解决从能量函数或非归一化密度中采样的问题，这在分子动力学等科学领域至关重要。现有方法，特别是基于生成对抗网络（GANs）的方法，在高维空间中训练困难，且对能量函数的梯度信息敏感。此外，直接训练连续归一化流（CNF）模型也面临挑战，因为需要精确的密度估计或复杂的对抗训练过程。

核心思路：iEFM的核心思路是利用能量函数构建一个边际向量场，并训练CNF模型来预测这个向量场的蒙特卡罗估计。通过这种方式，模型可以直接从能量函数学习，而无需显式地计算密度或进行对抗训练。这种方法是一种离策略学习方法，允许从任意分布中采样，并使用蒙特卡罗方法估计目标向量场。

技术框架：iEFM的整体框架包括以下几个主要步骤：1) 从任意分布中采样数据点；2) 使用能量函数计算每个数据点对应的边际向量场的蒙特卡罗估计；3) 使用能量基流匹配目标训练CNF模型，使其能够预测这些向量场；4) 迭代地更新模型，直到收敛。该框架可以扩展到不同的条件概率路径，如方差爆炸（VE）和最优传输（OT）路径。

关键创新：iEFM的关键创新在于提出了无模拟的能量基流匹配目标，这是一种离策略训练CNF模型的方法，可以直接从非归一化密度中学习。与现有方法相比，iEFM不需要显式地计算密度或进行对抗训练，从而避免了高维空间中的训练困难。此外，iEFM还引入了蒙特卡罗估计来近似边际向量场，从而降低了计算复杂度。

关键设计：iEFM的关键设计包括：1) 使用蒙特卡罗方法估计边际向量场，这需要选择合适的采样策略和样本数量；2) 设计能量基流匹配损失函数，该函数衡量CNF模型预测的向量场与蒙特卡罗估计之间的差异；3) 选择合适的CNF模型结构，例如使用神经网络来参数化向量场。此外，还需要选择合适的优化器和学习率，以确保模型的收敛。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，iEFM在二维高斯混合模型（GMM）和八维四粒子双阱势（DW-4)能量函数上均优于现有方法。在DW-4任务上，iEFM能够更准确地学习到能量函数的分布，并生成更真实的样本。这些结果证明了iEFM在复杂高维系统中进行高效且可扩展的概率建模的潜力。

🎯 应用场景

iEFM在多个领域具有广泛的应用前景，包括分子动力学模拟、材料科学、药物发现等。它可以用于学习复杂分子的3D结构分布，预测蛋白质的折叠构象，以及生成具有特定性质的新材料。此外，iEFM还可以应用于概率推理和贝叶斯推断等领域，解决高维概率分布的建模和采样问题。

📄 摘要（原文）

In this work, we consider the problem of training a generator from evaluations of energy functions or unnormalized densities. This is a fundamental problem in probabilistic inference, which is crucial for scientific applications such as learning the 3D coordinate distribution of a molecule. To solve this problem, we propose iterated energy-based flow matching (iEFM), the first off-policy approach to train continuous normalizing flow (CNF) models from unnormalized densities. We introduce the simulation-free energy-based flow matching objective, which trains the model to predict the Monte Carlo estimation of the marginal vector field constructed from known energy functions. Our framework is general and can be extended to variance-exploding (VE) and optimal transport (OT) conditional probability paths. We evaluate iEFM on a two-dimensional Gaussian mixture model (GMM) and an eight-dimensional four-particle double-well potential (DW-4) energy function. Our results demonstrate that iEFM outperforms existing methods, showcasing its potential for efficient and scalable probabilistic modeling in complex high-dimensional systems.