arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2601.16622cs.LGcs.AI
收藏

E2Former-V2:具有线性激活内存的即时等变注意力

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

University of Science and Technology of China(中国科学技术大学)

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

AI总结:

提出E2Former-V2架构,通过等变轴对齐稀疏化(EAAS)和即时等变注意力机制,利用SO(3)到SO(2)基变换和自定义Triton内核,实现线性激活内存和20倍TFLOPS提升,在SPICE和OMol25数据集上加速推理并保持预测性能。

摘要:

等变图神经网络(EGNN)已成为建模3D原子系统的广泛使用的方法。然而,主流架构由于在每条边上显式构造几何特征或密集张量积而面临关键的可扩展性瓶颈。为克服这一问题,我们引入了**E2Former-V2**,一种将代数稀疏性与硬件感知执行相结合的可扩展架构。我们首先提出**等变轴对齐稀疏化(EAAS)**。EAAS基于Wigner-$6j$卷积,利用$\mathrm{SO}(3) \rightarrow \mathrm{SO}(2)$基变换,将计算昂贵的密集张量收缩转化为高效的稀疏奇偶重索引操作。基于这种表示,我们引入了**即时等变注意力**,一种通过自定义融合Triton内核实现的完全节点中心机制。通过消除物化的边张量并最大化SRAM利用率,我们的内核相比标准实现实现了**20倍的TFLOPS提升**。在SPICE和OMol25数据集上的大量实验表明,E2Former-V2在保持相当预测性能的同时显著加速推理。这项工作表明,大型等变Transformer可以使用广泛可用的GPU平台高效训练。代码可在https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2获取。

英文摘要:

Equivariant Graph Neural Networks (EGNNs) have become a widely used approach for modeling 3D atomistic systems. However, mainstream architectures face critical scalability bottlenecks due to the explicit construction of geometric features or dense tensor products on \textit{every} edge. To overcome this, we introduce \textbf{E2Former-V2}, a scalable architecture that integrates algebraic sparsity with hardware-aware execution. We first propose \textbf{E}quivariant \textbf{A}xis-\textbf{A}ligned \textbf{S}parsification (EAAS). EAAS builds on Wigner-$6j$ convolution by exploiting an $\mathrm{SO}(3) \rightarrow \mathrm{SO}(2)$ change of basis to transform computationally expensive dense tensor contractions into efficient, sparse parity re-indexing operations. Building on this representation, we introduce \textbf{On-the-Fly Equivariant Attention}, a fully node-centric mechanism implemented via a custom fused Triton kernel. By eliminating materialized edge tensors and maximizing SRAM utilization, our kernel achieves a \textbf{20$\times$ improvement in TFLOPS} compared to standard implementations. Extensive experiments on the SPICE and OMol25 datasets demonstrate that E2Former-V2 maintains comparable predictive performance while notably accelerating inference. This work demonstrates that large equivariant transformers can be trained efficiently using widely accessible GPU platforms. The code is avalible at https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑