E2Former-V2：基于线性激活内存的即时等变注意力
机器学习
2026-01-26 v1 人工智能
摘要
等变图神经网络 (EGNN) 已成为建模 3D 原子系统的广泛使用方法。然而，主流架构因在每条边上显式构建几何特征或稠密张量乘积而面临严重的可扩展性瓶颈。为克服这一问题，我们引入了 \textbf{E2Former-V2}，一种集成代数稀疏性与硬件感知执行的可扩展架构。我们首先提出了 \textbf{E}quivariant \textbf{A}xis-\textbf{A}ligned \textbf{S}parsification (EAAS)。EAAS 基于 Wigner- 卷积，通过利用 基变换，将计算昂贵的稠密张量收缩转化为高效的稀疏奇偶性重新索引操作。在此表示之上，我们引入了 \textbf{On-the-Fly Equivariant Attention}，一种完全基于节点的机制，通过自定义融合 Triton 内核实现。通过消除物化边缘张量并最大化 SRAM 利用率，本内核相较于标准实现实现了 \textbf{20 倍的 TFLOPS 提升}。在 SPICE 和 OMol25 数据集上的广泛实验表明，E2Former-V2 在保持可比预测性能的同时显著加速了推理。本工作表明，使用广可访问的 GPU 平台可以高效训练大型等变变换器。代码已可用于 https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2。
引用
@article{arxiv.2601.16622,
title = {E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory},
author = {Lin Huang and Chengxiang Huang and Ziang Wang and Yiyue Du and Chu Wang and Haocheng Lu and Yunyang Li and Xiaoli Liu and Arthur Jiang and Jia Zhang},
journal= {arXiv preprint arXiv:2601.16622},
year = {2026}
}