今天为大家介绍的是来自Stefan Chmiela团队的一篇论文。近年来，基于从头参考计算（ab-initio reference calculations）的机器学习力场（MLFFs）取得了巨大的进展。尽管这些力场在测试中表现出较低的误差，但由于在分子动力学（MD）模拟中长期不稳定性的问题，其可靠性正受到越来越多的关注。作者的研究发现，累积误差的稳健性与MLFF中等变表示的使用之间存在潜在联系，但这种表示的计算成本较高。为了解决这一问题，作者提出了一种名为的transformer架构，它结合了稀疏等变表示（欧几里得变量）和一种将不变信息与等变信息分离的自注意机制，避免了复杂的张量积运算。实现了准确性、稳定性和速度的权衡，使得对物质量子特性的深入分析成为可能，适用于更长的时间尺度和更大的系统规模。为展示其能力，作者生成了包含数百个原子的柔性肽和超分子结构的稳定MD轨迹。此外，作者通过探索数千个能量最低点来研究中等大小链状分子（如小肽）的势能面（PES）拓扑结构。值得注意的是，能够在稳定性和与最低能量构象这两个冲突的需求之间取得平衡，这对于生物化学领域中真实的探索任务至关重要。
原子级建模依赖于长时间尺度的分子动力学（MD）模拟，以揭示实验观测到的系统宏观特性如何从微观尺度上的相互作用中产生。此类模拟的预测精度取决于驱动它们的原子间力的准确性。传统上，这些力要么来自近似的机械力场（FF），要么来自准确但计算量极大的从头电子结构计算（ab-initio electronic structure calculations）。近年来，机器学习（ML）势函数开始解决这个问题，利用分子系统的统计依赖性，展现出前所未有的灵活性。
在MLFF发展的过程中，不断取得的进展带来了许多越来越复杂的模型架构。其中，消息传递神经网络（MPNNs）已成为一种特别有效的架构类型。MPNNs可以看作是卷积操作的推广，能够处理非结构化数据域，如分子图。这种操作为从输入数据中提取特征提供了一种有效的方法，广泛应用于许多现代机器学习架构中。
最近在这一领域的进展集中在引入具有物理意义的几何先验，这导致了所谓的“等变MPNNs”。研究人员发现这类模型可以减少近似误差，并且在数据效率方面优于不变（invariant）模型。不变模型依赖于成对的距离来描述原子间的相互作用，因为这些距离在旋转时不会发生变化。然而，随着系统规模的增大、灵活性或化学异质性的增加，在这种有限的表示中推导出正确的相互作用模式变得越来越困难。因此，等变模型能够整合额外的方向信息，捕捉取决于相邻原子相对方向的相互作用。“等变”性质使得这些模型能够区分那些在简单模型中无法分离的相互作用，并能从相同的训练数据中学习到更具可迁移性的相互作用模式。大多数等变架构的一个基本构建模块是张量积（tensor product）。张量积在卷积操作中被计算为两个函数 f(x)和g(x)在线性基中展开后进行的乘积(f∗g)(x)。结果定义在原始基函数集的乘积空间中。因此，随着卷积操作次数的增加，相关的乘积空间会迅速呈指数级增长，计算花销非常大。
在SO(3)等变架构中，卷积操作在球谐函数的基下进行SO(3)旋转群的运算。通过固定球谐函数的最大阶数，避免了相关函数空间的指数级增长。然而，SO(3)卷积的计算复杂度随变化，与不变模型相比，可能使每个构象的预测时间增加到原来的100倍。这导致了在准确性、稳定性和速度之间不得不进行权衡的情况，这在高通量或大规模探索任务中可能带来显著的实际应用问题。
作者以此为动机，提出了一种欧几里得自注意力机制（图1c），用以替代SO(3)卷积，通过过滤原子邻域的相对方向来表示原子相互作用，无需进行高成本的张量积运算。
作者的解决方案基于神经网络架构设计的最新进展以及几何深度学习领域的发展。模型（图2）采用稀疏表示来描述分子几何，并将所有卷积响应的投影限制在等变基函数中最相关的不变分量上。由于球谐函数的正交归一性，这种投影相当于计算乘积张量的迹（trace），可以通过线性尺度的内积表示，从而实现对高阶等变表示的高效扩展，同时不牺牲计算速度和内存成本。力的预测通过结果的不变能量模型的梯度获得，该模型代表了一种自然等变的分段线性化过程。在整个过程中，自注意力机制用于在模型中解耦“不变”和“等变”基元素。
SO3KRATES在准确率、稳定性与速度之间做到了权衡
在接下来的数值实验中，作者展示了SO3KRATES如何克服MD稳定性、准确性和计算效率之间的权衡（图3）。最近的一项研究比较了不同先进MLFFs在短MD模拟中的稳定性，发现只有基于SO(3)卷积的NEQUIP架构能够提供可靠的结果。然而，这类模型的出色稳定性需要付出大量的计算成本（图3a）。这就需要在MLFF的稳定性和计算效率之间进行权衡，而SO3KRATES现在能够克服这一问题（图3b）。作者的模型允许每秒预测多达一个数量级以上的帧数（FPS）（图3c），实现亚毫秒速度的步进时间，同时在MD模拟中不牺牲可靠性或准确性（图3b）。然而，作者指出，测试准确性和稳定性并不一定相关（例如，比较图3a中的GEMNET和SPHERENET），但在实际应用中，只有同时准确且稳定的模型才具有意义。作者发现，SO3KRATES能够提供准确的力预测，有效地克服了这一权衡（图3b）。
任何经验运行时间的测量都依赖于特定的硬件和软件条件。参考文献32中报告的运行时间是为在PYTORCH+ASE中实现的MLFF模型测量的。为了确保与（在JAX中实现）的可比性，作者在与相同的设置下重新实现了两个具有代表性的模型（图3a）。由于上述研究报告中显示，最快的模型（FORCENET）在其基准测试中给出了错误的可观测量，作者选择了第二快的模型（SCHNET）进行重新实现。作为最稳定和最准确的模型，作者选择了NEQUIP进行重新实现。这个架构的选择也代表了不变模型（SCHNET）和基于等变SO(3)卷积的模型（NEQUIP），分别构成了计算复杂性的上限和下限（表1）。
所有模型都使用E3X库在JAX中重新实现。MD步进时间通过相同的用JAX-MD编写的MD代码在相同的物理设备（Nvidia V100 GPU）上测量。有趣的是，从PYTORCH+ASE（图3b中的紫色虚线）到JAX+JAX-MD（图3b中的紫色实线）的转换，NEQUIP实现了28倍的加速，SCHNET则实现了15倍的加速。这说明了相同设置的重要性以及JAX生态系统的潜力。需要注意的是，步进时间的测量不包括输入输出（IO）操作所需的时间，因为这些时间高度依赖于本地高性能计算（HPC）基础设施。因此，作者报告的完整模拟的总时间相对于报告的步进时间有一个固定的偏移量。
对于最多包含21个原子的有机小分子，相比于NEQUIP架构实现了平均5倍的加速，而步进时间比不变模型SCHNET稍长（大约1.4倍）。随着原子总数的增加，相对于SO(3)卷积的加速效果也在增加（图3c），这与理论尺度分析中较小的前因子一致（表1）。因此，对于双壁纳米管（370个原子），相比于NEQUIP，加速效果已经增加到30倍。与不变卷积相比，作者的方法预测速度略慢，这也符合理论预期。
Frank, J. T., Unke, O. T., Müller, K. R., & Chmiela, S. (2024). A Euclidean transformer for fast and stable machine learned force fields. Nature Communications, 15(1), 6539.
