DeepMind 提出欧几里得 Transformer：2.5 天完成 1 年分子动力学模拟
近年来，基于从头算参考计算的机器学习力场（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 的局限性，同时又不损害它们的优势。
