Nat. Commun.: 新型欧几里得Transformer模型推动机器学习力场在分子动力学模拟中应用

摘要：2024年8月6日—— 近日，国际知名期刊《Nature Communications》发表了一项突破性研究，科学家们成功开发了一种名为SO3KRATES的欧几里得Transformer模型，显著提升了机器学习力场（MLFFs）在分子动力学（MD）模拟中的稳定

https://doi.org/10.1038/s41467-024-50620-6

2024年8月6日—— 近日，国际知名期刊《Nature Communications》发表了一项突破性研究，科学家们成功开发了一种名为SO3KRATES的欧几里得Transformer模型，显著提升了机器学习力场（MLFFs）在分子动力学（MD）模拟中的稳定性、准确性和计算效率。这一成果为未来在生物化学、材料科学等领域的广泛应用提供了新的可能性。

分子动力学模拟是研究物质微观相互作用及其宏观性质的重要工具。传统的分子动力学模拟依赖于力场（FF）或基于量子力学的从头计算（ab initio）方法。然而，传统力场虽然计算速度快，但精度有限；而从头计算方法虽然精度高，但计算成本极其昂贵，难以应用于大规模或长时间尺度的模拟。近年来，机器学习力场（MLFFs）作为一种折中方案，逐渐成为分子动力学模拟中的研究热点。MLFFs通过机器学习模型从量子力学计算数据中学习分子系统的势能面（PES），能够在保持较高精度的同时，显著降低计算成本。

尽管MLFFs在测试数据集上表现出较低的误差，但在长时间尺度的分子动力学模拟中，其稳定性问题逐渐受到关注。许多MLFFs在长时间模拟中会出现累积误差，导致模拟结果偏离物理真实情况。这一问题的根源在于MLFFs在训练数据分布之外的区域表现不佳，尤其是在高温或构象灵活的系统中的外推能力较弱。为了解决这一问题，研究人员提出了多种改进模型架构的方法，其中最具代表性的是基于等变表示（equivariant representations）的模型。等变表示能够捕捉分子系统中原子间的方向性信息，从而提高模型的稳定性和外推能力。然而，现有的等变模型通常计算复杂度较高，限制了其在实际应用中的广泛使用。

在这项研究中，来自德国柏林工业大学、柏林学习与数据基础研究所（BIFOLD）、Google DeepMind等机构的研究团队提出了一种新型的Transformer架构——SO3KRATES。该模型结合了稀疏的等变表示（欧几里得变量）和自注意力机制，能够在保持等变模型优势的同时，显著降低计算复杂度。

SO3KRATES模型的核心创新在于其欧几里得自注意力机制。传统的等变模型通常依赖于SO(3)卷积操作，该操作在处理高维等变表示时计算复杂度较高。SO3KRATES通过将分子几何结构分解为高维不变特征和等变欧几里得变量（EV），并利用自注意力机制将不变特征和等变信息分离，从而避免了昂贵的张量积操作。具体而言，SO3KRATES模型通过以下两个关键设计实现了高效的计算：

不变特征与等变信息的分离：SO3KRATES模型将不变特征和等变欧几里得变量分别处理，并通过自注意力机制将两者结合起来。这种设计使得模型能够在保持等变信息的同时，显著降低计算复杂度。

欧几里得自注意力机制：SO3KRATES模型利用自注意力机制来捕捉原子邻域的相对方向信息，从而避免了传统的SO(3)卷积操作。通过将SO(3)卷积的输出投影到不变分量，模型能够高效地处理高维等变表示，而无需进行复杂的张量积计算。

SO3KRATES模型的架构包括多个欧几里得Transformer块，每个块由自注意力块和交互块组成。自注意力块负责更新不变特征和等变欧几里得变量，而交互块则通过交换不变特征和等变信息来进一步增强模型的表达能力。最终，模型通过多层感知机（MLP）预测每个原子的能量，并通过自动微分计算原子力，从而确保能量守恒。

研究团队通过一系列数值实验验证了SO3KRATES模型在分子动力学模拟中的优越性能。实验结果表明，SO3KRATES模型在稳定性、准确性和计算效率方面均优于现有的等变模型和不变模型。

1. 稳定性与计算效率的平衡

SO3KRATES模型在保持高稳定性的同时，显著提升了计算效率。实验表明，SO3KRATES模型能够在亚毫秒级别的时间内完成每一步分子动力学模拟，比现有的等变模型快一个数量级。例如，在包含370个原子的双壁纳米管系统中，SO3KRATES模型的模拟速度比现有的等变模型快30倍。此外，SO3KRATES模型在长时间模拟中表现出极高的稳定性，能够生成纳秒级别的稳定分子动力学轨迹。

2. 数据效率与外推能力

SO3KRATES模型在数据效率和外推能力方面也表现出色。实验表明，等变模型在相同训练数据量下能够获得更低的测试误差，并且在高温模拟中表现出更好的外推能力。例如，在柔性分子DHA的模拟中，不变模型在300K时已经表现出不稳定性，而等变模型在500K时仍能保持稳定。这表明等变模型能够更好地捕捉分子系统中的方向性信息，从而提高模型的稳定性和外推能力。

3. 大分子系统的泛化能力

SO3KRATES模型在泛化到大分子系统方面表现出色。研究团队利用仅包含42个原子的小肽训练数据，成功将模型应用于包含151个原子的更大肽链系统。实验结果表明，SO3KRATES模型能够捕捉到肽链的折叠过程，并生成与实验数据一致的结构。这表明SO3KRATES模型能够有效地将局部相互作用信息泛化到更大的分子系统中，为未来在生物大分子模拟中的应用提供了新的可能性。

4. 势能面拓扑的探索

SO3KRATES模型在势能面（PES）拓扑探索方面也表现出色。研究团队利用SO3KRATES模型对DHA和Ac-Ala3-NHMe分子的势能面进行了大规模探索，发现了数千个局部极小值。实验结果表明，SO3KRATES模型能够有效地捕捉到分子系统中的折叠漏斗（folding funnels），并生成与DFT计算结果一致的势能面拓扑结构。这表明SO3KRATES模型在探索复杂分子系统的势能面方面具有巨大的潜力。

SO3KRATES模型的提出为机器学习力场在分子动力学模拟中的应用开辟了新的道路。通过结合欧几里得自注意力机制和稀疏等变表示，SO3KRATES模型在保持等变模型优势的同时，显著降低了计算复杂度，从而实现了稳定性、准确性和计算效率的平衡。这一成果为未来在生物化学、材料科学等领域的广泛应用提供了新的可能性，特别是在长时间尺度模拟和大分子系统模拟中的应用前景广阔。