[论文解读] Transferable Boltzmann Generators

Leon Klein, Frank Noé|arXiv (Cornell University)|Jun 20, 2024

Lattice Boltzmann Simulation Studies被引用 7

一句话总结

本文提出基于连续正则化流（CNF）和等变图网络的可迁移 Boltzmann 生成器（TBGs），实现对未见双肽的零样本 Boltzmann 采样，并能够高效地重加权以得到目标分布。

ABSTRACT

The generation of equilibrium samples of molecular systems has been a long-standing problem in statistical physics. Boltzmann Generators are a generative machine learning method that addresses this issue by learning a transformation via a normalizing flow from a simple prior distribution to the target Boltzmann distribution of interest. Recently, flow matching has been employed to train Boltzmann Generators for small molecular systems in Cartesian coordinates. We extend this work and propose a first framework for Boltzmann Generators that are transferable across chemical space, such that they predict zero-shot Boltzmann distributions for test molecules without being retrained for these systems. These transferable Boltzmann Generators allow approximate sampling from the target distribution of unseen systems, as well as efficient reweighting to the target Boltzmann distribution. The transferability of the proposed framework is evaluated on dipeptides, where we show that it generalizes efficiently to unseen systems. Furthermore, we demonstrate that our proposed architecture enhances the efficiency of Boltzmann Generators trained on single molecular systems.

研究动机与目标

• 动机：需要高效的生成模型来对 Boltzmann 分布的分子构型进行采样。
• 提出一个可迁移的 Boltzmann 生成器框架，使其能够在不重新训练的情况下推广到未见分子。
• 利用带等变性的图神经网络的连续正则化流来保持物理对称性。
• 在二肽上展示可迁移性，并通过消融研究评估数据效率。
• 证明该框架支持对目标 Boltzmann 分布的准确重加权。

提出的方法

• 使用连续正则化流通过一个时变向量场 v_theta(t,x) 将高斯先验映射到目标 Boltzmann 分布。
• 用流匹配训练 CNF，以学习向量场，而无需特定系统的分子动力学数据。
• 采用一个 O(D)-和 S(N)-等变的图神经网络实现该向量场，确保置换/旋转不变性。
• 在推断阶段，重构拓扑和手性以过滤有效样本并实现对真实 Boltzmann 分布的重加权。
• 考虑多种可迁移架构（TBG、TBG + backbone、TBG + full），并使用拓扑感知编码。
• 通过重要性采样讨论重加权观测量，并评估有效样本量（ESS）。

Figure 1: Results for the alanine dipeptide system simulated with a classical force field (a) Ramachandran plots for the biased MD distribution (left) and for samples generate with the TBG + full model (right). (b) Energies of samples generated with different methods. (c) Free energy projection alon

实验结果

研究问题

• RQ1 Boltzmann 生成器是否可以在一组分子上训练后推广到未见分子而无需重新训练？
• RQ2 可迁移的 Boltzmann 生成器在未见双肽上从目标 Boltzmann 分布抽样的效果如何，是否能有效地重加权？
• RQ3 架构选择和训练数据对可迁移性和采样质量有何影响？
• RQ4 在多大程度上，较小的训练数据集或偏向采样能恢复正确的自由能景观？

主要发现

模型	NLL (↓)	ESS (↑)
BG + backbone (alanine dipeptide, semi-empirical)	-107.56±0.09	0.50±0.13%
TBG + full (alanine dipeptide, semi-empirical)	-124.71±0.08	1.03±0.17%
BG + backbone (alanine dipeptide, classical)	-109.02±0.01	1.56±0.30%
TBG + full (alanine dipeptide, classical)	-127.06±0.12	6.03±1.34%

• 本文提出基于 CNFs 与等变图网络的首个可迁移 Boltzmann 生成器（TBG）框架。
• 在 alanine dipeptide 上，TBG + full 的负对数似然更低、ESS 更高，优于基线（NLL: -124.71±0.08 vs -107.56±0.09；ESS: 1.03±0.17% vs 0.50±0.13%）。
• 对于经典力场，TBG + full 达到 NLL -127.06±0.12 和 ESS 6.03±1.34%，优于其他方案。
• 在 2AA dipeptides 中，TBG + full 一直产生更高的 ESS（最高 15.29±9.27%），以及更高的正确构型比率（最高 98±2%）。
• 有偏的训练倾向于稀有亚稳态可以改善自由能投影，尽管 ESS 可能下降。
• 即使训练数据较少，TBG + full 仍保持相当的可迁移性，尽管可能遗失一些亚稳态。

Figure 2: Results for the KS dipeptide (a) Sample generated with the TBG + full model (b) Ramachandran plot for the weighted MD distribution (left) and for samples generate with the TBG + full model (right). (c) TICA plot for the weighted MD distribution (left) and for samples generate with the TBG