Scalable Equilibrium Sampling with Sequential Boltzmann Generators
作者: Charlie B. Tan, Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, Alexander Tong
发布时间: 2025-02-27
来源: arxiv
研究方向: 统计物理与机器学习
主要内容
本文研究了一种名为Sequential Boltzmann Generators (SBG)的采样方法，旨在解决统计物理中分子状态在热力学平衡下的可扩展采样问题。SBG通过结合强大的正态化流和重要性采样，在目标分布下获得统计独立的样本。
主要贡献
1. 提出了一种名为Sequential Boltzmann Generators (SBG)的采样方法，该方法通过引入两个关键改进来扩展Boltzmann生成器框架：一个高效的非等变Transformer正态化流，直接在所有原子笛卡尔坐标上操作；以及在推理时对流样本进行缩放，使用退火Langevin动力学将样本转移到目标分布，从而降低方差（退火）的重要性权重，使得能够通过顺序蒙特卡罗进行更高保真度的重采样。
2. 通过使用现代化的非等变Transformer架构作为主干，SBG实现了对提议流的可扩展预训练，并展示了在所有数据集上比连续Boltzmann生成器具有更高的计算效率。
3. SBG是第一个能够解决三肽、四肽、六肽在笛卡尔坐标下的平衡采样问题的方法，而过去的方法在这些分子系统上是无能为力的。
研究方法
1. 非等变Transformer正态化流
2. 退火Langevin动力学
3. 顺序蒙特卡罗
4. 重要性采样
5. 数据增强
实验结果
SBG在所有评估指标上均实现了最先进的性能，特别是在小肽分子系统上，SBG是第一个能够解决三肽、四肽、六肽在笛卡尔坐标下的平衡采样问题的方法。实验结果表明，SBG在计算效率上远超连续Boltzmann生成器，并且在生成样本的质量上也有显著提升。
未来工作
未来工作可以探索将SBG应用于更大规模的分子系统，并研究混合方法，例如将CNFs通过蒸馏到一个可逆架构或基于一致性的目标。此外，可以考虑其他可扩展生成模型，例如自回归模型，这些模型也允许精确的似然估计。