E2Former用新方法让分子建模快7到30倍,还保持精度。
E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products
- 把计算重心从边移到节点,用Wigner 6j卷积降低复杂度
- 相比传统SO(3)卷积提速7至30倍,仍能捕捉精细几何信息
- 适合需要高效分子模拟的研究者,尤其关注大规模系统
等变图神经网络(EGNN)在化学、生物和材料科学的微观系统建模中表现优异,但其构建边特征所需的球面张量积带来高昂计算成本,难以应用于大规模系统。为此,我们提出E2Former,一种具备等变性与高效性的Transformer架构,引入Wigner $6j$卷积(Wigner $6j$ Conv)。该方法将计算负担从边转移至节点,使复杂度从 $O(|/mathcal{E}|)$ 降至 $ O(| /mathcal{V}|)$,同时保持模型表达能力和旋转等变性。实验表明,该方法相较传统 $ ext{SO}(3)$ 卷积实现7倍至30倍的速度提升。结果证明,E2Former有效缓解了现有方法的计算瓶颈,且不牺牲对几何细节的捕捉能力,为可扩展、高效的分子建模提供了新方向。
原文摘要 · Abstract (English)
Equivariant Graph Neural Networks (EGNNs) have demonstrated significant success in modeling microscale systems, including those in chemistry, biology and materials science. However, EGNNs face substantial computational challenges due to the high cost of constructing edge features via spherical tensor products, making them impractical for large-scale systems. To address this limitation, we introduce E2Former, an equivariant and efficient transformer architecture that incorporates the Wigner $6j$ convolution (Wigner $6j$ Conv). By shifting the computational burden from edges to nodes, the Wigner $6j$ Conv reduces the complexity from $O(|\mathcal{E}|)$ to $ O(| \mathcal{V}|)$ while preserving both the model's expressive power and rotational equivariance. We show that this approach achieves a 7x-30x speedup compared to conventional $\mathrm{SO}(3)$ convolutions. Furthermore, our empirical results demonstrate that the derived E2Former mitigates the computational challenges of existing approaches without compromising the ability to capture detailed geometric information. This development could suggest a promising direction for scalable and efficient molecular modeling.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。