🤖 AI

E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

E2Former-V2 是一种可扩展的等变图神经网络架构，它通过引入等变轴对齐稀疏化（Equivariant Axis-Aligned Sparsification）和即时等变注意力机制（On-the-Fly Equivariant Attention mechanism），克服了主流模型的计算瓶颈，在保持对 3D 原子系统预测性能的同时，实现了 20 倍的 TFLOPS 提升。

原作者： Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

发布于 2026-06-08

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图预测一大群人（原子）在房间里会如何移动和互动。在化学和物理学领域，这被称为“原子级建模”（atomistic modeling）。长期以来，科学家们一直使用一种特殊的 AI，称为等变图神经网络（Equivariant Graph Neural Network, EGNN）。这些网络非常出色，因为它们理解如果旋转整个房间，其中的物理规律是不应发生改变的。

然而，目前的这些网络版本有一个主要问题：当人群变得庞大时，它们会变得笨拙且缓慢。

以下是该论文提出的解决方案——E2Former-V2 的简单拆解，使用了日常生活的类比。

问题：“以边为中心”的交通拥堵

想象一场派对，每个人都想和他们的邻居聊天。

旧方法（传统的 EGNN）： 系统会为每一场正在发生的对话创建一个巨大的、实体的清单。如果有 1,000 个人，且每个人都与 50 个邻居交谈，系统必须在巨大的白板（内存）上写下 50,000 条单独的对话笔记，然后才能进行数学运算。
瓶颈： 随着派对规模的扩大（更多原子），这个白板会变得巨大，导致计算机耗尽空间。它会把所有时间都花在搬运这些笔记上，而不是真正进行思考。这就是为什么旧模型在模拟蛋白质等大型分子时会崩溃或变得极其缓慢。

解决方案：E2Former-V2

作者构建了一个新系统，不再记录每一条对话笔记。相反，它使用了两个聪明的技巧来使过程变得极速。

技巧 1：“神奇指南针”（等变轴对齐稀疏化）

在旧系统中，计算原子如何相互作用就像是在解一个复杂的 3D 拼图，而每个碎片都可以向任何方向旋转。这很混乱，且需要大量的数学计算。

类比： 想象你有一个复杂的 3D 拼图。与其在拼图旋转时尝试解决它，不如使用一个“神奇指南针”来旋转整个拼图，使其中一个特定的碎片始终指向正上方。
结果： 一旦拼图对齐，大部分碎片就会变成空白或零。你不再需要对这些空白部分进行数学运算了。论文称之为 EAAS。它将混乱的、稠密的计算转变为一个简单的“重新索引”游戏（仅仅是移动碎片），这要快得多。

技巧 2：“流式厨师”（即时等变注意力机制）

这是最大的突破。

旧方式（自助餐）： 旧系统会准备好整个自助餐（所有的 50,000 条对话笔记），并在任何人开吃之前将其摆在桌子上。如果桌子太小，食物就会溢出，派对也会停止。
新方式（流式厨师）： 新系统 E2Former-V2 就像一位厨师，在顾客面前一次只做一份菜，做完后立即端上桌。
- 它不写下笔记。
- 它不存储整个列表。
- 它计算两个人之间的相互作用，将结果加入最终得分，然后立即忘记细节，以便为下一对组合腾出空间。
结果： 这被称为“即时”（On-the-Fly）处理。因为它从不建立那个巨大的白板，所以它几乎不占用内存。它能将所有内容放入计算机超快的内部缓存（SRAM）中，使其比旧方法快 20 倍。

他们究竟取得了什么成就

论文测试了这个新系统在两个主要方面的表现：

小分子（SPICE）： 他们在类药物分子上进行了测试。新系统与现有的最佳模型一样准确，但速度更快。
巨型分子（OMol25）： 他们在巨大的化学结构上进行了测试。
- 极限： 旧模型在尝试模拟超过 10,000 个原子的系统时会崩溃（内存溢出）。
- 胜利： E2Former-V2 成功模拟了拥有 100,000 个原子 的系统而没有崩溃。

文中提到的现实世界测试

作者不仅运行了数字，还模拟了真实的物理过程：

液态水： 他们模拟了一个水分子盒。新模型几乎完美地预测了水分子的排列方式（径向分布函数），比以往的模型更贴近真实世界的实验数据。
蛋白质： 他们模拟了一个在水中漂浮的大型蛋白质（约 30,000 个原子）。该模型保持了蛋白质的稳定性，并帮助其折叠成符合真实实验数据的形状，证明了它能够处理复杂的生物结构。

总结

E2Former-V2 就像是从一个需要将每场对话都写在巨大的、缓慢的白板上的系统，升级到了一个智能厨师能即时烹饪并提供每一场对话的系统。这使得科学家能够在标准计算机上模拟大规模分子系统（如巨大的蛋白质），而不会耗尽内存，同时保持物理学的准确性。

技术摘要：E2Former-V2：具有线性激活内存的即时（On-the-Fly）等变注意力机制

1. 问题陈述

等变图神经网络（EGNNs）因其能够保持旋转和平移对称性，已成为建模 3D 原子系统的标准。然而，主流架构（如 MACE、EquiformerV2、eSCN）受困于其**以边为中心（edge-centric）**设计的关键效率瓶颈。这些模型显式地构建几何特征或在每条边上执行密集的张量积，导致：

内存复杂度： $O(|E|)$ 的激活内存，其中 $|E| \approx kN$ （ $k$ 为平均邻居数）。对于大型系统，这会导致巨大的内存占用，从而超出 GPU 容量。
硬件效率低下： 显式物化（materialization）边张量的过程迫使数据在高速带宽内存（HBM）与片上 SRAM 之间频繁移动，产生了一个限制吞吐量的“内存墙”。
可扩展性限制： 随着系统规模（ $N$ ）的增加，延迟呈非线性增长，使得在大规模分子动力学（MD）模拟中进行计算在普通硬件上变得极其昂贵。

尽管最近在 Transformer 领域（如 FlashAttention）通过流式、分块执行解决了序列数据的类似瓶颈，但这些原则在 SO(3) 等变架构中仍未得到充分探索。

2. 方法论

E2Former-V2 通过将代数稀疏化（algebraic sparsification）与硬件感知的以节点为中心（node-centric）的执行模型相结合来解决这些问题。该架构包含两个核心创新：

2.1. 等变轴对齐稀疏化 (EAAS)

标准的 SO(3) 张量收缩涉及对 Clebsch–Gordan 系数的密集求和，导致 $O(L^6)$ 的复杂度以及不规则的内存访问模式。E2Former-V2 引入了 EAAS，将这些密集操作转化为基于置换的稀疏操作：

SO(3) $\to$ SO(2) 基底变换： 该方法将全局坐标系与局部向量 $\vec{r}$ （连接两个原子的轴）对齐。在此对齐框架下，实球谐函数变得极大稀疏（引理 4.1），其中只有 $m=0$ 分量非零。
重索引算子（Re-indexing Operator）： 通过将旋转通过张量积进行交换，密集的收缩坍缩为一个确定性的重索引规则（公式 13）。该算子通过稀疏的分块线性映射将输入特征映射到输出特征，消除了对磁性指数（magnetic indices）进行密集求和的需求。
结果： 这降低了卷积阶段的复杂度，并实现了比标准实现约 $1.5\times$ 的加速。

2.2. 即时（On-the-Fly）等变注意力机制

为了消除 $O(|E|)$ 的内存瓶颈，作者提出了一个通过自定义 Triton 内核 实现的全融合、以节点为中心的注意力机制：

以节点为中心的分解： 注意力机制不再物化边特征，而是解耦计算。Query 和 Key 向量使用 SO(3) 不变的标量投影（公式 6）构建，而等变几何信息仅通过 Value 路径进行传播。
流式归约（Streaming Reduction）： 内核以流式方式处理邻居。它迭代遍历邻居 $j \in \mathcal{N}(i)$ ，在 SRAM 内实时计算注意力分数并累加 Value 更新。
消除边物化： 形状为 $N \times K \times H \times D$ 的中间张量永远不会被写入 HBM。内核在片上维护运行统计数据（最大值、归一化和 Value 累加器），将激活内存复杂度从 $O(|E|)$ 降低到 $O(|V|)$ 。
在线 Softmax： 为了确保流式处理过程中的数值稳定性，使用了在线 Softmax 公式来增量更新累加器。

3. 核心贡献

等变轴对齐稀疏化 (EAAS)： 一种新颖的代数技术，统一了 Wigner-6j 重结合（recoupling）与 SO(2) 轴对齐，将密集的 SO(3) 张量收缩转换为稀疏的重索引操作。
即时等变注意力机制： 一个全融合的 Triton 内核，避免了边物化，最大限度地提高了 SRAM 复用率，并实现了比标准实现高出达 20 $\times$ 的 TFLOPS。
可扩展性能： 证明了能够处理单 GPU 上高达 100,000 个原子 的系统，在这一规模下，之前的等变 Transformer 会因显存溢出（OOM）而失败。
物理保真度： 在分子动力学中的验证表明，该模型能准确复现液态水的结构，并能稳定地精炼 AlphaFold 预测的蛋白质结构（约 30k 原子）。

4. 实验结果

作者在 SPICE（药物化学）和 OMol25（基础模型规模）数据集以及 MD 模拟中对 E2Former-V2 进行了评估。

预测精度：
- 在 SPICE 上，E2Former-V2 在大多数子集上实现了最低的能量和力平均绝对误差（MAE），在二聚体（Dimers）的能量 MAE 上优于 MACE-Large 达 48%。
- 在 OMol25 上，保守版本实现了 1.27 meV/atom 的总能量 MAE，与 eSEN-small 持平，并显著优于 UMA-S (3.62 meV/atom)。
- 在 OC20 S2EF-2M 上，其精度与 EquiformerV2 和 eSCN 等强基准模型相当。
效率与可扩展性：
- 吞吐量： 融合内核实现的平均 TFLOPS 比朴素 PyTorch 实现快 20 $\times$ 。
- 内存： 与基准模型相比，峰值内存使用量降低了 16.9 $\times$ 至 19.5 $\times$ 。
- 可扩展性： 先前的模型（如 EquiformerV2、MACE-Large）在超过 10k–50k 原子时会遇到 OOM 错误，而 E2Former-V2 能可靠地扩展至 100k 原子。
- 推理速度： 在 100k 原子规模下，E2Former-V2 (Conservative) 达到 0.12 steps/s，几乎比 UMA-S 快 3 倍。在 Direct 设置下，它在较小规模时比 Allegro 和 EquiformerV2 快一个数量级。
分子动力学：
- 对周期性体相水（216 个分子）的模拟显示，E2Former-V2 与实验径向分布函数（RDF）高度吻合，表现优于 MACE-OFF。
- 一个溶剂化蛋白质（约 30k 原子）的 2 ns 模拟保持了相对于原生结构的高 TM-score，展示了稳定的动力学特性。

5. 意义与主张

论文声称，E2Former-V2 展示了在大规模通用 GPU 上运行大型等变 Transformer 的实际可行性。通过将执行范式从以边为中心转向以节点为中心，并利用硬件感知的流式处理，该方法解决了限制等变模型可扩展性的根本内存瓶颈。

作者强调，这不仅仅是算法上的改进，更是一种系统级的实现，它弥合了理论上的分解（将复杂度从 $O(|E|)$ 降至 $O(|V|)$ ）与硬件效率之间的鸿沟。这项工作表明，大规模分子基础模型和长时标 MD 模拟（此前在计算上是难以实现的）现在可以使用标准的 GPU 基础设施进行研究。代码已发布，以促进社区的进一步采用。