2.5天完成1年的MD计算？DeepMind团队基于欧几里得Transformer的新计算方法

近年来，基于从头算参考计算的机器学习力场 (MLFF) 的开发取得了巨大进展。虽然实现了较低的测试误差，但由于担心在较长的模拟时间范围内会出现不稳定性，MLFF 在分子动力学 (MD) 模拟中的可靠性正面临越来越多的审查。

研究表明，对累积不准确性的稳健性与 MLFF 中使用等变表示之间存在潜在联系，但与这些表示相关的计算成本可能会在实践中限制这种优势。

为了解决这个问题，Google DeepMind、柏林工业大学（TU Berlin）的研究人员提出了一种名为 SO3krates 的 transformer 架构，它将稀疏等变表示（欧几里得变量）与分离不变和等变信息的自注意力机制相结合，从而无需昂贵的张量积。

SO3krates 实现了精确度、稳定性和速度的独特组合，能够对长时间和系统尺度上的物质量子特性进行深入分析。

该研究以「A Euclidean transformer for fast and stable machine learned force fields」为题，于 2024 年 8 月 6 日发布在《Nature Communications》。

背景与挑战性

分子动力学（MD）模拟通过长时间尺度的模拟，可以揭示系统从微观相互作用到宏观性质的演变，其预测精度取决于驱动模拟的原子间力的精确度。传统上，这些力来源于近似的力场（FF）或计算复杂的从头计算电子结构方法。

近年来，机器学习（ML）势能模型通过利用分子系统的统计依赖性，提供了更灵活的预测手段。

然而，研究表明，ML模型在基准数据集上的测试误差与长时间尺度MD模拟中的表现相关性较弱。

为改善外推性能，消息传递神经网络（MPNNs）等复杂架构被开发，特别是等变MPNNs，通过引入张量积捕捉原子间的方向信息，提高了数据的可转移性。

在 SO(3) 等变架构中，卷积是在球谐函数基础上对 SO(3) 旋转群进行的。通过固定架构中球谐函数的最大次数，可以避免相关函数空间的指数增长。

科学家已证明最大阶数与准确度、数据效率密切相关，并与 MD 模拟中模型的可靠性相关。然而，SO(3) 卷积的规模为，与不变模型相比，这可以将每个构象的预测时间增加多达两个数量级。

这导致了一种必须在准确性、稳定性和速度之间做出妥协的情况，还可能会带来重大的实际问题。必须先解决这些问题，这些模型才能在高通量或广泛的探索任务中发挥作用。

性能强劲的新方法

Google DeepMind、柏林工业大学的研究团队以此为动机，提出了一种欧几里得自注意力（Euclidean self-attention）机制，用原子邻域相对方向的过滤器代替 SO(3) 卷积，从而无需昂贵的张量积即可表示原子相互作用；该方法称为称为 SO3krates。

图示：SO3krates架构和构建块。（来源：论文）

该解决方案建立在神经网络架构设计和几何深度学习领域的最新进展之上。SO3krates 使用稀疏表示来表示分子几何形状，并将所有卷积响应的投影限制在等变基函数中最相关的不变分量上。

图示：学习不变量。（来源：论文）

由于球谐函数的正交性，这种投影对应于乘积张量的迹，可以用线性缩放内积来表示。这可以高效地扩展到高阶等变表示，而不会牺牲计算速度和内存成本。

力预测是根据所得的不变能量模型的梯度得出的，该模型代表自然等变的分段线性化。整个过程中，自注意力机制用于分离模型中的不变和等变基础元素。

该团队比较了 SO3krates 模型的与当前最先进的 ML 模型的稳定性和速度，发现该解决方案克服了当前等变 MLFF 的局限性，同时又不损害它们的优势。

研究人员提出的数学公式可以实现高效的等变架构，从而实现可靠稳定的 MD 模拟；与具有相当稳定性和准确性的等变 MPNN 相比，其速度可提高约 30 倍。

为了证明这一点，研究人员仅在几个小时内就对超分子结构运行了精确的纳秒级 MD 模拟，这使他们能够计算出从具有 42 个原子的小肽到具有 370 个原子的纳米结构的结构的收敛速度自相关函数的傅里叶变换。

图示：结果概述。（来源：论文）

研究人员进一步应用该模型，通过使用最小值跳跃算法研究 10k 最小值，探索二十二碳六烯酸 (DHA) 和 Ac-Ala3-NHMe 的 PES 拓扑结构。

这样的研究需要大约 30M FF 评估，这些评估在几百 K 到 1200 K 之间的温度下进行。使用 DFT 方法，这种分析需要一年以上的计算时间。现有的具有类似预测精度的等变 MLFF 需要运行一个多月才能完成此类分析。

相比之下，该团队仅用 2.5 天就能完成模拟，从而有可能在实际时间尺度上探索数十万个 PES 最小值。

此外，SO3krates 能够检测训练数据中未包含的物理有效最小构象。外推到 PES 未知部分的能力对于将 MLFF 扩展到大型结构至关重要，因为可用的从头算参考数据只能覆盖构象丰富的结构的子区域。

该团队还研究了在网络架构中禁用等方差属性的影响，以更深入地了解其对模型特性及其在 MD 模拟中的可靠性的影响。

研究人员发现，等变性与所得 MD 模拟的稳定性以及更高温度的外推行为有关。能够证明，即使测试误差估计平均相同，等变性也会降低误差分布的扩散。

图示：计算效率和 MD 稳定性。（来源：论文）

因此，通过等变表示使用方向信息在精神上与经典 ML 理论类似，其中映射到更高的维度可以产生更丰富的特征空间，这些特征空间更容易参数化。

未来的研究

在最近的一系列研究中，提出了旨在降低 SO(3) 卷积计算复杂度的方法。它们可以作为完整 SO(3) 卷积的替代品，而本文介绍的方法可以在消息传递范式中完全避免使用昂贵的 SO(3) 卷积。

这些结果都表明，等变相互作用的优化是一个尚未完全成熟的活跃研究领域，可能提供进一步改进的途径。

图示：稳定、长时间尺度的 MD 模拟和对更大生物分子的外推。（来源：论文）

虽然该团队的工作，使得使用现代 MLFF 建模范式可以实现稳定的扩展模拟时间尺度，但未来仍需进行优化，以使 MLFF 的适用性更接近传统的经典 FF。

目前，该方向出现了各种有潜力的途径：在当前设计中，EV 仅根据二体相互作用进行定义。通过将原子簇扩展纳入 MP 步骤，可以进一步提高准确性。同时，这可能有助于减少 MP 步骤的数量，从而降低模型的计算复杂度。

另一个尚未讨论的问题是全局效应的适当处理。通过使用低秩近似、可训练的 Ewald 求和或通过以物理启发的方式学习长程校正。当需要外推到更大的系统时，后一种类型的方法尤为重要。

虽然等变模型可以提高局部相互作用的外推能力，但对于超出训练数据中存在的长度尺度或超出模型有效截止值的相互作用，这并不适用。

由于上述方法依赖于局部属性，例如部分电荷、电负性或赫什菲尔德体积，因此可以通过学习 SO3krates 架构的不变特征分支中的相应局部描述符，将它们无缝集成到该团队的方法中。

因此，未来的工作将集中于将多体展开、全局效应和长程相互作用纳入 EV 形式，并旨在进一步提高计算效率，最终以高精度跨越 MD 时间尺度。

论文链接：https://www.nature.com/articles/s41467-024-50620-6

相关内容：https://phys.org/news/2024-08-faster-coupling-AI-fundamental-physics.html

文章来自于微信公众号机器之心 作者ScienceAI