DeepMind 与柏林工业大学联合提出 SO3krates 欧几里得 Transformer 架构，用于提升机器学习力场（MLFF）在分子动力学（MD）模拟中的效率与稳定性。该方法利用欧几里得自注意力机制替代昂贵的 SO(3) 卷积，结合稀疏等变表示，无需张量积即可捕捉原子间方向信息。实验显示，相比现有等变模型速度提升约 30 倍，仅需 2.5 天即可完成原本需一年的 PES 拓扑结构分析。研究证实等变性有助于降低误差分布扩散，增强高温外推能力。未来工作将聚焦于多体展开、全局效应及长程相互作用的整合，以进一步扩大 MLFF 的应用范围。
近年来，基于从头算参考计算的机器学习力场（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) 旋转群进行的。通过固定架构中球谐函数的最大次数 $l_{max}$，可以避免相关函数空间的指数增长。
科学家已证明最大阶数与准确度、数据效率密切相关，并与 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 理论类似，其中映射到更高的维度可以产生更丰富的特征空间，这些特征空间更容易参数化。
传统的 SE(3) 等变网络通常依赖于球谐函数卷积，这在处理高阶信息时计算开销巨大。SO3krates 的核心创新在于引入了欧几里得自注意力机制。该机制允许模型直接关注原子间的相对位置和方向，而不必显式地进行高维张量运算。
具体而言，SO3krates 将原子邻域内的相对方向编码为过滤器。这种设计使得模型能够在保持旋转等变性的同时，显著减少参数量和计算复杂度。通过分离不变信息和等变信息，模型可以更有效地学习物理规律，特别是在处理长程相互作用时。
SO3krates 采用稀疏表示来描述分子几何形状。这意味着只有那些对当前计算至关重要的局部环境会被激活和处理。此外，通过将卷积响应投影到等变基函数中最相关的不变分量上，模型避免了全量张量积的计算。
这种投影利用了球谐函数的正交性质，将复杂的张量收缩转化为简单的内积操作。这不仅降低了内存占用，还加快了前向传播的速度，使得在大规模系统上进行高精度模拟成为可能。
在多个基准测试中，SO3krates 展现了卓越的性能。与现有的最先进等变 MLFF 相比，它在保持相似精度的前提下，速度提升了约 30 倍。这一提升主要得益于其优化的架构设计，减少了不必要的计算冗余。
更重要的是，SO3krates 在长时间模拟中表现出了极高的稳定性。传统的 MLFF 往往在模拟进行一段时间后出现能量漂移或结构崩溃，而 SO3krates 能够有效抑制这种误差累积。这对于药物发现和材料科学中的长期过程模拟至关重要。
通过对 DHA 和 Ac-Ala3-NHMe 的势能面（PES）拓扑结构进行研究，展示了 SO3krates 在处理复杂生物分子方面的潜力。这些系统涉及大量的原子和复杂的构象变化，传统的计算方法难以在合理时间内完成。
使用 SO3krates，研究人员成功地在 2.5 天内完成了原本需要一年时间的计算任务。这使得探索更大规模的化学空间和更复杂的反应路径成为现实，极大地加速了科研进程。
在最近的一系列研究中，提出了旨在降低 SO(3) 卷积计算复杂度的方法。它们可以作为完整 SO(3) 卷积的替代品，而本文介绍的方法可以在消息传递范式中完全避免使用昂贵的 SO(3) 卷积。
这些结果都表明，等变相互作用的优化是一个尚未完全成熟的活跃研究领域，可能提供进一步改进的途径。
虽然该团队的工作，使得使用现代 MLFF 建模范式可以实现稳定的扩展模拟时间尺度，但未来仍需进行优化，以使 MLFF 的适用性更接近传统的经典 FF。
目前，该方向出现了各种有潜力的途径：在当前设计中，EV 仅根据二体相互作用进行定义。通过将原子簇扩展纳入 MP 步骤，可以进一步提高准确性。同时，这可能有助于减少 MP 步骤的数量，从而降低模型的计算复杂度。
另一个尚未讨论的问题是全局效应的适当处理。通过使用低秩近似、可训练的 Ewald 求和或通过以物理启发的方式学习长程校正。当需要外推到更大的系统时，后一种类型的方法尤为重要。
虽然等变模型可以提高局部相互作用的外推能力，但对于超出训练数据中存在的长度尺度或超出模型有效截止值的相互作用，这并不适用。
由于上述方法依赖于局部属性，例如部分电荷、电负性或赫什菲尔德体积，因此可以通过学习 SO3krates 架构的不变特征分支中的相应局部描述符，将它们无缝集成到该团队的方法中。
因此，未来的工作将集中于将多体展开、全局效应和长程相互作用纳入 EV 形式，并旨在进一步提高计算效率，最终以高精度跨越 MD 时间尺度。
DeepMind 与柏林工业大学合作推出的 SO3krates 架构，代表了机器学习力场领域的一项重要突破。它通过创新的欧几里得自注意力机制，解决了传统等变网络计算成本高和模拟不稳定的问题。
这一成果不仅大幅缩短了分子动力学模拟的时间，还为探索复杂生物分子和材料的物理性质提供了强有力的工具。随着技术的进一步优化，SO3krates 有望在计算化学、药物研发和材料科学等领域发挥更大的作用，推动相关研究的快速发展。
相关内容：https://phys.org/news/2024-08-faster-coupling-ai-fundamental-physics.html
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online
