E2Former-V2让3D分子模型更快更省资源,推理速度提升20倍。
E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory
- 用几何稀疏化技术把复杂计算转为高效索引操作
- 自研GPU核实现每秒20倍的浮点运算效率提升
- 适合在普通显卡上训练大规模分子预测模型
等变图神经网络(EGNN)广泛用于建模三维原子系统,但主流架构因在每条边上显式构建几何特征或稠密张量积而面临严重可扩展性瓶颈。为此,我们提出可扩展的E2Former-V2架构,融合代数稀疏性与硬件感知执行。首先提出等变轴对齐稀疏化(EAAS),基于威格纳-6j卷积,利用SO(3)→SO(2)基变换将计算昂贵的稠密张量收缩转化为高效的稀疏奇偶重索引操作。在此表示基础上,引入原位等变注意力机制,通过定制融合的Triton内核实现完全节点中心化计算。该机制避免生成边缘张量,最大化SRAM利用率,相比标准实现达到20倍的TFLOPS提升。在SPICE和OMol25数据集上的大量实验表明,E2Former-V2在保持相当预测性能的同时显著加速推理。本工作证明,大规模等变变换器可在普遍可用的GPU平台上高效训练。代码已开源:https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。