Nat. Mach. Intell. | 用线性复杂度捕捉分子世界的 超距感应 ——Euclidean Fast Attention

论文信息 标题：Machine learning global atomic representations with Euclidean fast attention 作者：J. Thorben Frank, Stefan Chmiela, Klaus-Robert Müller, Oliver T. Unke 机构：Google DeepMind（柏林）/ 柏林工业大学 / BIFOLD / 马克斯·普朗克信息研究所 / 高丽大学 发表：Nature Machine Intelligence, 2026年3月25日 DOI：10.1038/s42256-026-01195-y 代码：github.com/thorben-frank/euclidean_fast_attention

一、研究背景与核心问题

1.1 机器学习力场的崛起与局限

分子动力学（Molecular Dynamics, MD）模拟是理解蛋白质折叠、材料相变、化学反应机理的核心工具。其准确性完全依赖于对原子间相互作用力的精确描述。传统第一性原理方法（如 DFT）精度高但计算代价极其昂贵，难以扩展到大体系或长时间尺度。

近年来，机器学习力场（Machine Learning Force Fields, MLFFs） 以逼近 DFT 精度、仅需其极小计算代价的特性迅速成为主流。MLFFs 已被成功应用于蛋白质动力学模拟、新材料发现等前沿任务。

然而，MLFFs 存在一个长期悬而未决的根本性缺陷：对长程相互作用的建模能力严重不足。

1.2 局部截断的两难困境

为了实现系统尺寸 N 下的线性计算复杂度 O(N)，几乎所有主流 MLFFs 都引入了局部截断半径：两个原子之间距离超过 时，彼此视为不相互作用。

这一设计在计算效率上高度合理，但在物理真实性上代价惨重。以下几类在实际化学与生物系统中普遍存在的相互作用，在局部模型中几乎无法正确描述：

这些作用力通常在量级上弱于局部共价相互作用，但对于体系的热力学稳定性、长时间动力学行为和响应性质却至关重要。

1.3 消息传递神经网络（MPNN）的有效感受野分析

目前最主流的 MLFF 骨干架构是消息传递神经网络（Message Passing Neural Networks, MPNNs），如 NequIP、SchNet、PaiNN、Equiformer 等。

MPNN 将化学结构表示为图，其中原子为节点，在 轮消息传递后，理论上的"有效感受野"上界为 。这比纯局部模型更大，但存在两个关键缺陷：

1. 1. 断联节点问题：若两原子之间不存在连续的"跳跃路径"（hop path），即便其距离在 以内，信息依然无法传递。
2. 2. 均场效应（Mean-field effect）：多跳间接信息传递本质上是对邻域信息的加权平均，对于需要精确捕捉特定原子对几何关系的长程效应，仅靠均场近似往往远远不够。

论文通过几何 Weisfeiler-Leman 测试等理论工具和受控实验，系统性地证明了上述局限。

二、方法：Euclidean Fast Attention（EFA）

EFA 的核心思想是：将 Transformer 中强大的自注意力机制引入 Euclidean 空间数据，同时克服其 复杂度问题，并保持原子相互作用所要求的物理对称性（平移不变、旋转不变/等变）。

2.1 Euclidean 旋转位置编码（ERoPE）

问题：如何在线性注意力框架中编码几何信息？

标准二次注意力可以通过显式计算原子对之间的距离、方向等几何特征来引入空间信息。但在线性注意力中，"键值积"需要先于查询独立计算，这使得引入依赖于原子对的几何信息变得非平凡。

解决方案：ERoPE

作者受自然语言处理中旋转位置编码（RoPE） 的启发，提出将三维坐标 编码进特征向量 的方法：

其中 为频率系数， 为单位方向向量。

在注意力框架中，将位置 、 分别编码进查询 与键 后，其内积自然包含相对位移信息：

这一构造优雅地将相对位移 嵌入了注意力权重，且对平移天然不变。

旋转不变性的获取：对 积分

上述内积中的指数项 仍然依赖于方向 的选取，因而对旋转不具不变性。解决方案是对所有可能的方向进行积分：

积分结果是仅依赖于距离 的 sinc 函数，从而实现 SO(3) 旋转不变性。通过在不同特征维度上使用不同频率系数 ，可以将对距离的依赖表达为多个 sinc 函数的混合，大幅增强表达能力。

数值上，该积分通过 Lebedev 求积法（球面高斯求积）实现：以 G 个离散网格点 近似球面积分，时间复杂度为 ，当 时相比 具有实质性优势。

2.2 EFA 的完整形式

将 ERoPE、球面积分与线性注意力三者结合，得到等变 EFA 的核心公式：

其中：

• • ， 为任意保等变性的特征映射
• • 为包含所有 阶球谐函数的向量
• • 为张量积，使输出特征兼容等变表示

关键计算技巧：先计算"键值外积" （对所有原子求和，），再与查询 作用，全程 ，无需存储或计算 的注意力矩阵。

离散化后的实现公式为：

2.3 与 SO(3) 卷积的关联

对于等变输入特征（），EFA 更新后的原子特征在解析形式上近似为：

其中 为向量值径向函数， 为单位方向向量。这一形式与 NequIP、TFN 等等变 MPNN 中 SO(3) 卷积的基本构件完全类似，但 EFA 实现全局范围（无截断半径）的 SO(3) 卷积，而传统 MPNN 只能在局部邻域内进行。

2.4 周期边界条件（PBC）的处理

对于晶体材料模拟，单元胞内的 SO(3) 对称性被破坏。此时论文提供了一种替代方案：选取晶格向量而非随机单位向量作为 ，使 ERoPE 退化为一种平面波展开，与 Ewald 求和方法建立了深刻的联系，可无缝处理周期系统的长程静电相互作用。

2.5 数值精度与实用参数指导

EFA 的旋转不变/等变精度受两个参数控制：Lebedev 网格点数 和最大频率 。两者通过约束 相互关联， 定义为数值积分误差低于 （接近单精度浮点极限）时的最大参数值。

实践经验表明：长程作用通常对应低频成分，因此往往可以使用很小的 （例如 或 ），保持极高效率。

三、理想化系统上的分析实验

在应用于真实分子之前，论文通过一系列受控实验深入分析了 EFA 的理论性质。

3.1 几何表达力分析

实验一：区分局部邻域 以不同几何构型的局部邻域对（Pozdnyakov 等人提出的测试集）为对象，比较 EFA 和 SO(3) 卷积所需的最小球谐函数阶数 。结果表明两者具有相同的阶数需求规律，验证了 EFA 在几何表达能力上与 SO(3) 卷积等价。

实验二：k-链图同构性测试 构造两个依据几何 Weisfeiler-Leman（GWL）测试不同构的链状图 、（差异仅体现在链端节点的相对朝向上）。带局部截断的 MPNN 至少需要 层才能区分两者，与 GWL 理论预测完全吻合。而 EFA 仅需单次更新即可区分，表现与全连接图上的 MPNN 等价。

3.2 两体势与多体 NaCl 系统

对简单的各向同性衰减势（类静电）和各向异性的电荷-偶极势进行测试：

• • 标准 MPNN 在粒子间距超出截断后给出定性错误的常数预测
• • 的 EFA 可准确描述仅依赖距离的势函数
• • 只有 的 EFA 才能正确描述依赖距离与方向的电荷-偶极势

对大小从 10 Å 到 30 Å 直径的 NaCl 类团簇系统，使用单层 MP + EFA 模型：精度随系统尺寸的增长远优于任意层数的纯 MPNN；推理时间随原子数 N 线性增长，线性拟合 ，实证验证了理论分析。

四、真实化学体系上的全面验证

4.1 非局域电荷转移基准（Ko et al.）

该基准测试集由带非局域效应的三类系统构成，涵盖有/无周期边界的情况，是目前公认的最具挑战性的 MLFF 基准之一。

RMSE 结果对比（meV/atom 或 meV/Å）：

EFA 在 8 个指标中的 7 个上超越包括 SpookyNet（专门设计处理非局域效应的模型）在内的所有对比方法。

4.2 SN2 反应模拟

选取 （X, Y ∈ {F, Cl, Br, I}）这一经典亲核取代反应。

定量结果：

• • 加入 EFA 后，能量 MAE 降低 34倍，力 MAE 降低 8倍
• • 将局部截断从 5 Å 扩大至 10 Å，或加入色散修正，均无法解决问题

定性结果（最为关键）： 在以相同初始条件运行 MD 轨迹时：

• • 标准 MP 模型：离子与分子之间超出截断距离后没有任何相互作用力，反应物直接"飞过"彼此，反应根本不发生
• • MP+EFA 模型：正确感知到长程吸引力，反应物相互靠近，甲基卤化物分子经历正确的朝向调整，发生完整的背面攻击，生成目标产物

这一对比生动说明：在涉及长程导向效应的化学反应场景中，局部模型的错误是定性的，不可通过参数调整弥补。

4.3 分子二聚体非共价相互作用

以 DES370K 基准（源自 SPICE 数据集，含 76 种二聚体系统）为测试对象，评估 EFA 对非共价长程相互作用（静电、感应、色散及其混合）的泛化能力。

将二聚体间长程作用拟合为物理动机的展开形式：

其中 涵盖了从电荷-电荷（）到色散（）的全部相互作用类型。

MP+EFA 不仅数值精度更高，其系数分布与真实值分布几乎完全重叠，而标准 MP 模型的系数分布严重失真——这说明 EFA 正确学习到了长程相互作用的物理函数形式，而非仅仅拟合了某些局部特征。模型还展现出对训练集外 4 个全新二聚体的强鲁棒泛化能力。

4.4 累积烯（Cumulene）电子离域效应

累积烯 （k=7）是测试长程非局域效应最具代表性的有机分子之一。其能量对链两端 CH₂ 基团之间二面角 Θ 的依赖极为敏感，但这一依赖关系在空间上跨越了整个碳链骨架。

关键结论：

• • 标准 MP（T=3 层）：感受野不足，能量预测几乎为平坦曲线，能垒完全消失
• • 标准 MP（T=5 层）：感受野勉强覆盖，但因"过度压缩"（over-squashing）现象，能垒高度仍被严重低估
• • MP+EFA（T=3，）：仍无法描述，因为纯距离信息不足以分辨二面角变化
• • MP+EFA（T=3，）：准确描述完整能量曲线；更高 进一步改善能垒预测精度

MD 模拟的后续分析揭示了错误能量曲面带来的灾难性后果：

• • 基于标准 MP 的轨迹：所有二面角均等概率被采样（因为预测为平坦势面），构象分布完全错误
• • 基于 MP+EFA 的轨迹：二面角稳定维持在 90° 附近的势能极小值处，符合物理预期
• • 功率谱分析：基于 MP 的轨迹在 ~2300 cm⁻¹ 处出现虚假峰，这一伪影在 MP+EFA 轨迹中完全消失

4.5 局域主导体系：无副作用

在 BIGDML 数据集（以局域作用为主的材料体系）上，MP 与 MP+EFA 性能相当，证明在不需要长程效应的场景下，EFA 模块不引入额外误差，是真正意义上的即插即用增强。

五、方法实现细节

5.1 ERoPE 的高效实现

ERoPE 以实数运算等价实现，避免复数计算：对于 维的特征向量，采用 个 旋转矩阵的块对角矩阵，并利用稀疏结构以逐元素乘法高效计算：

每个旋转矩阵使用独立频率系数 ，使得不同特征维度可以以不同频率感知距离，极大丰富了模型的径向表达能力。

5.2 等变特征的张量积框架

论文采用 E3x 库实现 O(3) 不可约表示（irreps）框架。特征 按宇称 和阶数 组织，不同阶数特征之间通过 Clebsch-Gordan 系数进行张量积耦合，严格保持 O(3) 等变性。ERoPE 沿"特征轴"独立作用，宇称轴和阶数轴作为批处理维度处理。

5.3 模型整体架构

初始化原子嵌入（按原子序数）
    ↓
循环 T 次 MessagePassing 层（局部 MPNN，截断半径 r_cut）
    + EFA 块（全局注意力，无截断）
    → 两者输出相加
    ↓
最终不变特征 → MLP → 原子能量贡献
    ↓
对所有原子求和 → 总能量
    ↓
对坐标自动微分 → 负梯度 = 原子受力

MP 块采用等变连续卷积，径向滤波函数在 处平滑衰减至零，保证表示的连续性。EFA 块默认参数为：，，，（用于不变特征系统）。

六、计算复杂度分析

当 时（实践中通常满足），EFA 优于标准自注意力两个数量级，且是目前唯一在保持 SO(3) 等变性的同时实现全局感知与线性复杂度三者兼得的机制。

七、局限性与未来方向

论文对 EFA 的局限性有充分的自我认知：

1. 1. 高频短程相互作用：球面积分的数值精度在高频（ 较大）时下降，导致 EFA 对急剧变化的短程排斥势描述不佳。这正是论文坚持将 EFA 作为局部 MPNN 的增强模块而非完全替代品的原因。
2. 2. 纯 EFA 架构的可行性：目前尚不清楚能否构建完全基于 EFA 块的端到端模型，仍需进一步算法改进。
3. 3. 积分格点的选择：当系统尺寸增大导致原子间最大距离增加时，若要维持相同的数值精度，可能需要相应增大 ，在极端情况下将破坏严格的线性扩展性。
4. 4. 等变特征下的计算开销：高阶等变特征（大 ）的张量积运算本身计算量不小，其性能与成本的权衡仍需针对具体任务优化。