Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics
Timewarp：通过学习时间粗化动力学实现分子动力学的可转移加速
https://arxiv.org/abs/2302.01170
摘要
分子动力学（MD）模拟是一种广泛使用的模拟分子系统的技术，最常见的是在全原子分辨率下进行，其中运动方程以飞秒量级的时间步长进行积分。MD 通常用于计算平衡性质，这需要从平衡分布（如玻尔兹曼分布）中进行采样。然而，许多重要过程，如结合和折叠，发生在毫秒或更长时间尺度上，无法通过常规 MD 有效采样。此外，每次研究新的分子系统时，都需要进行新的 MD 模拟。我们提出了 Timewarp，一种增强采样方法，它使用归一化流作为马尔可夫链蒙特卡罗方法中的提议分布，目标是对玻尔兹曼分布进行采样。该流在 MD 轨迹上离线训练，并学会在时间上进行大步长，模拟
的分子动力学。至关重要的是，Timewarp 在分子系统之间是可转移的：一旦训练完成，我们展示了它对未见过的短肽（2-4 个氨基酸）在全原子分辨率下进行泛化，探索其亚稳态并提供与标准 MD 相比的采样加速。我们的方法构成了向通用、可转移的加速 MD 算法迈出的重要一步。
1 引言
分子动力学（MD）是一种在原子水平上模拟物理系统的成熟技术。当准确执行时，它提供了对分子运动的详细机理的无与伦比的洞察，无需湿实验室实验。MD模拟已被用于理解分子生物物理学中的核心兴趣过程，如蛋白质折叠[1, 2]、蛋白质-配体结合[3]和蛋白质-蛋白质结合[4]。MD的许多关键应用归结为有效地从玻尔兹曼分布中采样，即分子系统在温度T下的平衡分布。
设为时间t时分子的状态，包括N个原子在笛卡尔坐标下的位置
和速度
。玻尔兹曼分布由以下公式给出：
标准分子动力学模拟不会在分子系统之间传递信息：对于每个研究的系统，都必须进行新的模拟。这是一个浪费机会：许多分子系统表现出密切相关的动态，模拟一个系统应该能够产生与类似系统相关的信息。特别是，蛋白质，由20种氨基酸序列组成，是研究这种可转移性的首要候选者。我们提出了一种名为 Timewarp 的通用、可转移的增强采样方法，它使用一个归一化流 [9] 作为针对玻尔兹曼分布的马尔可夫链蒙特卡洛（MCMC）方法的提议分布。我们的主要贡献是：
我们提出了第一个在一般笛卡尔坐标中工作的机器学习算法，它展示了对训练期间未见过的小型肽系统的可转移性。
我们首次展示了对未见过的肽系统的玻尔兹曼分布的无偏 MCMC 采样的实时加速。
我们定义了一个使用条件归一化流作为提议分布，通过 Metropolis-Hastings（MH）校正步骤确保详细平衡（第3.4节）的 MCMC 算法，目标是玻尔兹曼分布。
我们设计了一个基于变换器的、排列等变的归一化流。
我们制作了一个包含数千种小型肽的分子动力学轨迹的新训练数据集。
我们展示了即使在没有 MH 校正（第3.5节）的情况下部署，Timewarp 也可以用来比 MD 更快地探索新肽的亚稳态。
2 相关工作
最近，分子上的深度学习引起了广泛关注。玻尔兹曼生成器 [10, 11, 12] 使用流从玻尔兹曼分布中进行渐近无偏采样。有两种生成样本的方法：(i) 从流中生成独立同分布样本，并使用统计重加权来去偏期望值。(ii) 在 MCMC 框架中使用玻尔兹曼生成器 [13]，如 Timewarp。目前，玻尔兹曼生成器缺乏在多个分子之间泛化的能力，这与 Timewarp 不同。唯一的例外是 [14]，他们提出了一个在扭转空间中的扩散模型，并使用底层 ODE 作为可转移的玻尔兹曼生成器。然而，与 Timewarp 不同，他们使用内部坐标，不在全原子系统中操作。此外，[15, 16] 最近在笛卡尔坐标中为分子引入了玻尔兹曼生成器，可能实现可转移训练。最近，[17] 提出了 GeoDiff，一个从分子图预测分子构象的扩散模型。像 Timewarp 一样，GeoDiff 在笛卡尔坐标中工作，并泛化到未见过的分子。然而，GeoDiff 未应用于蛋白质，而是小分子，并且不以玻尔兹曼分布为目标。与 Timewarp 不同，[18] 学习了多个时间分辨率的转移概率，准确捕捉了动力学。然而，他们没有展示系统间的可转移性。与 Timewarp 最相似的是最近的工作 [19]，他们训练了一个可转移的 ML 模型来模拟聚合物的时间粗粒度动力学。然而，与 Timewarp 不同，他们的模型作用于粗粒度表示。此外，它未应用于蛋白质，并且没有 MH 步骤，这意味着错误可以在模拟中累积而未被纠正。
马尔可夫状态模型（MSMs）[20, 21, 22] 通过运行许多短 MD 模拟来工作，这些模拟用于定义离散状态空间，以及估计的转移概率矩阵。与 Timewarp 类似，MSMs 估计时间 t 和时间 t + τ 之间的状态转移概率，其中 τ ≫ ∆t。最近的工作将深度学习应用于 MSMs，导致了 VAMPnets [23] 和深度生成 MSMs [24]，它们用深度网络取代了 MSM 数据处理管道。与 Timewarp 不同，这些模型不可转移，并且在粗粒度、离散状态空间中建模动力学，而不是在全原子坐标表示中。
之前有许多关于神经自适应采样器的工作 [25, 26, 27]，它们使用深度生成模型作为提议分布。A-NICE-MC [25] 使用体积保持流，通过似然自由对抗方法进行训练。其他方法使用旨在鼓励探索的目标函数。我们目标函数中的熵项受到 [28] 的启发。与这些方法不同，Timewarp 专注于泛化到新的分子系统而不重新训练。
存在许多增强采样方法用于 MD，如并行温度提升 [29, 30] 或沿过渡路径提议集体变量的更新 [8, 31]。鉴于 Timewarp 加速 MD 的能力，它通常提供了与这些技术集成的机会。
3方法
3.1 条件归一化流
3.2 数据集生成
3.3 增强归一化流
3.4使用渐近无偏MCMC瞄准玻尔兹曼分布
迭代这些更新会产生一个样本。要获得位置的玻尔兹曼分布样本
，我们只需丢弃辅助变量
。由于将m 送入无穷是不可行的，我们模拟链直到达到固定的预算。在实践中，我们发现我们模型的接受率可能很低，大约为1%。然而，我们强调即使接受率低，如果提出的变更足够大，我们的模型也可以导致更快的探索，正如我们在第6节中展示的。此外，我们引入了一种批量采样程序，该程序在保持详细平衡的同时显著加快了采样速度。这种程序用单个前向传递采样一批提议，并接受第一个满足MH接受标准的提议。MCMC算法的伪代码在附录C中的算法1。
3.5 快速但有偏的状态空间探索，不使用MH校正
与其使用MH校正来保证渐近无偏的样本，我们可以选择在简单的探索算法中使用Timewarp。在这种情况下，我们忽略了MH校正，并接受所有能量变化低于某个阈值的提议。这使得状态空间的探索速度大大加快，并且在第6节中我们展示了，尽管在技术上是有偏的，但这通常会导致定性准确的自由能估计。它还成功地比算法1和标准MD快几个数量级地发现了所有的亚稳态，这可以用于为后续的MSM方法提供初始化状态。伪代码在附录D的算法2中给出。
4 模型架构
4.1 对称性
MD动力学尊重某些物理对称性，这些对称性在算法中加以考虑会带来优势。现在我们描述如何在Timewarp中整合这些对称性。
置换等变性 设σ是N个原子的一个置换。由于原子没有内在的顺序，对x(t)进行置换对未来状态x(t + τ)的唯一影响是类似地置换原子，即
我们的条件流满足置换等变性。为了证明这一点，我们使用附录A.1中证明的以下命题，这是对[40, 41]中条件流的扩展：
5培训目标
6 实验
我们在小肽系统上评估Timewarp。为了与MD进行比较，我们关注亚稳态之间的最慢过渡，因为这些是最难穿越的。为了找到这些过渡，我们使用时间滞后独立成分分析（TICA）[42]，这是一种线性降维技术，最大化变换坐标的自相关性。最慢的成分，TIC 0和TIC 1，特别值得关注。为了衡量Timewarp实现的速度提升，我们计算TICA成分在每秒实际时间（ESS/s）内的有效样本大小。ESS/s由以下公式给出：
接下来，我们在2AA数据集上展示Timewarp的可转移性。在训练了2AA中的二肽之后，我们在测试二肽上部署Timewarp，并使用算法1进行采样，链长为200万。Timewarp的接受概率在0.03%到2%之间，并探索了所有亚稳态（附录B.1）。图3ef显示了二肽QW和HT的结果，显示了Timewarp与长时间MD链（1微秒=2×步）之间的密切一致性。对于这些二肽，Timewarp在ESS/s方面比MD分别实现了5倍和33倍的加速（附录B.4）。图4左图显示了Timewarp与MD在每个100个测试二肽上的加速因子。Timewarp在这些肽上提供了大约五倍的中位加速因子。此外，我们还在没有MH校正的情况下生成了Timewarp模型的样本，如第3.5节所述。我们从每个测试肽的相同初始状态开始，只对100条并行链进行10000步采样。对于每个肽，我们只选择找到所有亚稳态状态的链进行评估。与之前一样，我们计算ESS/s以与MD进行比较，显示出大约600倍的中位加速因子（图4c）。请注意，当使用所有并行采样的链时，实际加速将大得多。Timewarp探索得到的自由能估计与MD定性相似，但不如Timewarp MCMC准确（图3f）。
最后，我们研究了更具挑战性的4AA数据集。在训练集上训练后，我们使用算法1对每个测试四肽采样2000万个马尔可夫链状态，并与长时间的MD轨迹（1微秒）进行比较。与更简单的二肽相比，Timewarp MCMC和长时间MD轨迹都错过了一些亚稳态。然而，探索模式下的Timewarp（算法2）可以用作验证工具，快速验证整个状态空间的探索。图5a显示，MD和Timewarp MCMC未探索的亚稳态可以被Timewarp探索算法发现。我们通过在它们的附近运行更短的MD轨迹来仔细确认这些发现状态的物理有效性（见附录B.5），以确保它们不仅仅是模型发明的伪影。与2AA一样，我们再次报告Timewarp相对于MD的加速因子，如图4b，d所示。尽管Timewarp MCMC未能为大多数四肽加速采样，但Timewarp探索显示出中位数加速因子约为50。对于8个测试四肽，MD未能探索所有亚稳态，而Timewarp成功——这些用绿色标记。对于10个四肽，Timewarp MCMC未能找到MD发现的所有亚稳态——这些用灰色标记。图5b显示，当Timewarp MCMC发现所有亚稳态时，其自由能估计与MD非常吻合。然而，它有时会错过亚稳态，导致在这些区域的自由能估计较差。图5c显示，Timewarp MCMC还导致势能分布与MD非常接近。相比之下，Timewarp探索发现了所有亚稳态（甚至是MD错过的），但自由能图的准确性较低。它还有一个相对于MD和Timewarp MCMC略高的势能分布。
7 局限性 / 未来工作
Timewarp MCMC算法为大多数肽生成了较低的接受概率（< 1%）（见附录B.1）。然而，这本身并不是一个限制。通常，更大的提议时间步长τ会导致更小的接受率，因为预测问题变得更加困难。然而，由于算法1，我们可以在几乎没有额外成本的情况下并行评估多个样本。因此，较低的接受率，当与更大的时间步长τ结合时，通常是一个有利的权衡。尽管在使用MH校正时，我们仅在约三分之一的4AA肽上加速，但在全原子表示中在未见过的肽上以实际时间击败MD是一项具有挑战性的任务，以前没有被ML方法证明过。此外，可以考虑使用半经验力场而不是经典力场来针对系统。鉴于Timewarp比MD模拟需要少得多的能量评估，可以预期在此背景下会有更显著的加速。
尽管MD和Timewarp MCMC未能找到一些被Timewarp探索发现的亚稳态，但由于高计算成本（附录F），我们没有运行更长时间的MD和Timewarp MCMC。与传统的MD模拟相比，Timewarp在相同时间内生成的样本更少。因此，这种样本稀缺性在过渡态中变得更加明显，这使得Timewarp难以应用于化学反应。
Timewarp可以与其他增强采样方法集成，如并行温度或过渡路径采样。在并行温度的情况下，有效的集成需要在多个温度下训练Timewarp模型，然后允许使用Timewarp而不是MD对所有副本进行采样。我们还可以交替使用Timewarp提议与学习到的集体变量更新，如二面角角度。这些组合步骤仍然允许从目标分布中进行无偏采样[31]。
此外，我们仅在本工作中研究了小肽系统。将Timewarp扩展到更大系统仍然是未来研究的一个主题，有几个有前途的方向可以考虑。一种方法是探索不同的网络架构，可能捕捉系统中固有的所有对称性。另一种选择是研究粗粒化结构而不是全原子表示，以降低问题的维度。
8 结论
我们提出了Timewarp，这是一种可转移的增强采样方法，使用深度网络在模拟分子系统时提出大的构象变化。我们展示了使用MH校正的Timewarp可以在许多未见过的二肽上加速渐近无偏采样，从而更快地计算平衡期望值。尽管这种加速仅在我们考虑的少数四肽中是可能的，但我们展示了不使用MH校正的Timewarp比标准MD更快地探索二肽和四肽的亚稳态，并且我们验证了发现的亚稳态在物理上是合理的。这提供了一种有前途的方法，可以快速验证MD模拟是否访问了所有亚稳态。尽管需要进一步的工作将Timewarp扩展到更大、更有趣的生物分子，但这项工作清楚地展示了深度学习算法利用可转移性来加速MD采样问题的能力。
代码可在以下链接获取：[https://github.com/microsoft/timewarp](https://github.com/microsoft/timewarp)。数据集可按需提供。
