E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products

Thank you for your detailed review and helpful feedback. We address your concerns below and describe the improvements we will make in our revision.

W1.1 Thanks for spotting those notation errors. You are right—we should use $k_m^{(l)}$ consistently throughout. We will fix these typos immediately. Also, following Reviewer ejRh’s suggestion, we will update our terminology to ‘solid spherical harmonics’ to reflect the unnormalized form we use.

W1.2 Our original writing attempts to convey that $m$ is tied to $l$. To avoid confusion, we will modify the summation to be $\sum_{m_1 = - l_1}^{l_1} \sum_{m_2 = - l_2}^{l_2}$.

W1.3 $Y^{(1)}(\vec{r}_{ij}) = Y^{(1)}(\vec{r}_i) - Y^{(1)}(\vec{r}_j)$ holds exactly for solid spherical harmonics, as clarified with Reviewer ejRh. Here, $Y^{(1)}(\vec{r}_i)$ follows the standard convention of denoting positions relative to a predefined origin.

W1.4 Figure 1: We appreciate your feedback on visual clarity. We have carefully redrawn the figure with distinct colors for different components, improved font consistency, and replaced the "xx dependent" labels with more descriptive terminology. We apologize that we cannot provide the revised figure here due to NeurIPS restrictions on sharing updated PDFs or external links, but these improvements will be incorporated in the final version of our paper.

W1.5 Complexity Clarification: Thank you for pointing out the difference between reducing tensor product calculations and overall complexity. We apologize for not highlighting this more clearly and will revise our wording to be more precise.

W1.6 Overall Presentation Logic: We sincerely thank the reviewer for the detailed and constructive feedback. We appreciate the reviewer's suggestions on the organization and agree that enhancing the clarity and coherence of our manuscript is important.

We initially chose to emphasize the mathematical derivations in the main text, given that the Wigner 6j convolution is a central contribution enabling the efficiency of our model, as demonstrated by the 7x-30x speedup in Section 4.1. However, we fully acknowledge the reviewer’s concern that a clearer connection between theory and practical architecture would greatly improve readability. Therefore, we will incorporate an overview and diagram of the E2Former architecture into the main body, highlighting key design choices like gated aggregation, and will clearly point readers to additional details provided in the appendix.

Regarding the proofs of Theorems 3.2 and 3.3, we respectfully believe they represent meaningful contributions. The introduction of the Wigner 6j convolution specifically addresses the reduction of tensor product calculations from $O(|\mathcal{E}|)$ to $O(|\mathcal{V}|)$. The theoretical justifications for these proofs are crucial: without the proof of Theorem 3.2, the factorization would only work for spherical harmonics of order 1, and it would not be clear why the binomial theorem applies in irreps space. Similarly, without introducing Wigner 6j symbols as done in the proof of Theorem 3.3, the decomposition from edge-level to node-level TPs would not be possible (for example $\sum_j h_j \otimes Y(r_i) \otimes Y(r_j)$ would still require edge-scaling TPs). The derivation of these results requires careful treatment of angular momentum coupling theory, which we believe will interest the broader ML community. However, we fully acknowledge the reviewer’s point that the proofs may appear abrupt in the current presentation. To address this, we will add an introductory explanation at the beginning of Section 3 to provide clear context and intuition. Additionally, we will greatly simplify the proof sketches to improve readability and ensure smoother integration within the main text.

We greatly appreciate the reviewer’s suggestions and believe these modifications will enhance the accessibility of our manuscript.

W2.1 Section placement: We agree with this suggestion and will move the “Related Work” section immediately after the introduction to better contextualize our contributions.

W2.2 Missing comparisons and discussion: We appreciate your insightful feedback, and to address your valuable suggestions, we will incorporate an expanded and detailed discussion of spherical-scalarization methods.

To frame our perspective within the broader context of EGNN, we note that scalarization is a special case of the tensor product: its characteristic steerable vector‑to‑scalar interaction is naturally expressed as a TP. For example, when an odd‑type vector interacts with a scalar, the operation can be written as $1\times L_o \otimes 1\times 0_e \rightarrow 1\times L_o$ (with similar derivations applicable to other operations). Studying the full tensor‑product setting is therefore still meaningful. Our aim is not to claim that tensor‑product models are superior, but to optimize a well‑established framework—one that underlies many state‑of‑the‑art molecular‑modeling systems (e.g., OMat24; ECD, ICLR ’25; QH9, NeurIPS ’23). Moreover, tensor‑product methods remain the preferred approach for systems with periodic boundary conditions (OMat24).

W2.3 Equivariant transformer discussion: We appreciate this suggestion and will create a dedicated subsection discussing equivariant transformer architectures to provide a more comprehensive overview of the field.

W3 We appreciate the reviewer’s feedback on our experimental design. Our primary goal is to demonstrate that E2Former achieves performance comparable to state-of-the-art spherical equivariant methods (such as Equiformer V2), but with significantly enhanced efficiency. The key contribution of our work is this scalability, enabling large-scale simulations previously impractical for most spherical equivariant models.

W3.1: OC-22 Performance and Ablation Study While E2Former does not set a new state-of-the-art on the OC-22 benchmark, it achieves competitive results with substantially greater efficiency. E2Former’s training requires only 1,500 GPU hours, a threefold reduction compared to the 4,500 hours needed for the leading Equiformer V2 model. This highlights our method’s core design goal: enabling scalability and efficiency rather than extracting marginal performance gains on existing benchmarks.

As requested, we conducted an ablation study on the OC20 2M dataset. The table below compares our complete model to several variants, demonstrating that each architectural component contributes to either improved accuracy or enhanced efficiency.

W3.2: Small-Scale System Performance We appreciate the reviewer's observation on small-scale benchmarks. Our model is specifically optimized for efficiency on large-scale systems, where its primary advantages are most evident. The QM9 benchmark was included in the appendix as a quality check rather than a primary evaluation metric. As expected, E2Former's performance on this smaller benchmark is comparable to similar architectures like Equiformer and EquiformerV2, which aligns with our design objectives. We thank the reviewer for raising this point and will work to further enhance performance on small-scale benchmarks in the future.

W3.3: Real-World Large-Scale Experiments Following the reviewer's suggestion, we added new experiments on large-scale systems. We trained the model on a large dataset consist of diverse molecules with DFT force labels and evaluated it using MD simulations of systems containing up to 6,000 atoms. Our results closely match DFT calculations and outperform baseline methods.

We first tested the efficiency and scalability on those systems. E2Former consistently achieved the highest throughput for systems up to 6,400 atoms, with its performance advantage becoming clearer as system size increased. At 3,200 atoms, E2Former processes data nearly three times faster than MACE. Importantly, E2Former is the only model capable of handling simulations at 6,400 atoms, a scale where other models fail.

To check both accuracy and speed, we tested our model on a 6,000-atom water cluster and looked at the vibration patterns of the atoms. E2Former matched the reference DFT results very closely, capturing all the important features: low-frequency movements (below 1000 cm⁻¹), the H–O–H bending motion (1650 cm⁻¹), and most importantly, the O–H stretching pattern. MACE-Large, on the other hand, had significant errors, especially for the high-frequency stretching modes. We also tested on a more complex system - a Chignolin peptide in water (~2,000 atoms) - and successfully optimized its structure while keeping force predictions highly accurate (0.484 kcal/mol/A error).

Note: We apologize that we cannot share the figure due to NeurIPS restrictions. The table below describes the key peaks from our analysis.

I've raised my rating to Strong Accept and increased my confidence level to 5. The authors deserve it, enjoy it.

3. Enhanced Experimental Validation on Challenging Systems. As requested, we provide new evaluations on the MD22 benchmark to demonstrate E2Former's strong performance on complex small-scale systems, validating our design choices. E2Former is highly competitive with specialized SOTA models.

Note: Selected from MD22 as representative systems due to limited rebuttal time; full table will be in the revised manuscripts.

We also wish to clarify a minor point: we apologize for previously misunderstanding the reviewer's intended meaning regarding decentralization. Indeed, our implementation preprocesses coordinates by subtracting the center of mass, ensuring the system is centralized before training. We will clearly highlight this detail in the revised manuscript.

4. Positioning within Related Literature. We sincerely appreciate the reviewer’s insightful suggestions and constructive feedback on framing our discussion of related literature. We fully agree that explicitly highlighting the strengths of spherical-scalarization methods and clearly positioning our work within the broader landscape of equivariant model design would significantly enhance our manuscript’s narrative clarity. Indeed, as perceptively noted by the reviewer, specialized scalarization approaches such as ViSNet and GotenNet have demonstrated impressive empirical advantages across diverse tasks compared to their more general tensor-product-based counterparts (e.g., TFN, Equiformer).

Following the reviewer’s valuable guidance, we will revise our manuscript to briefly acknowledge the proven speed and empirical effectiveness of spherical-scalarization methods. Additionally, inspired by the insightful suggestions drawn from D-spanning [A], we will clarify that while scalarization methods offer substantial efficiency in specific applications, the tensor product framework maintains broader theoretical completeness. Thus, accelerating full tensor product computations—as achieved through our proposed Wigner 6j convolution—remains an important and relevant research direction.

We greatly appreciate this thoughtful recommendation, as it provides us with a clearer path to accurately contextualize and strengthen the presentation of our contributions.

We sincerely thank the reviewer for their insightful and constructive feedback, which significantly enhances the clarity and depth of our manuscript.

1. The Wigner 6j Convolution as a General, Numerically-Faithful Speedup. We agree with the reviewer regarding the potential general applicability of the Wigner 6j convolution as a numerically faithful solution to accelerate tensor products in various equivariant architectures.

To rigorously substantiate this claim, we performed an additional targeted ablation study on numeric equivalence. Specifically, we compared outputs from a 12-layer, 384-hidden-channel E2Former initialized identically with and without the Wigner 6j convolution (untrained, single random seed). Results demonstrate numerical equivalence within single-precision limits:

• Energy MAE: 2.13e-5 ± 4.35e-5 (relative error: 3.84e-4 ± 8.67e-4)
• Force MAE: 5.90e-8 ± 22.1e-8 (relative error: 5.43e-5 ± 2.05e-4)

Note: This experiment was performed with random system of 50 atoms using a 12-layer model with a 384-channel hidden dimension, initialized with the same random seed. The model was intentionally left untrained to isolate the numerical equivalence of the operation itself.

2. Architectural Rationale, and Plug and Play Implementation. We agree with the reviewer’s perspective on the broader potential of our Wigner 6j convolution as a 'plug-and-play' solution. We are especially thankful for your perceptive question about our architectural choices, as it allows us to articulate a key insight that was not sufficiently highlighted before: although the mathematical operation is equivalent, the complex, edge-centric designs of existing models make a simple drop-in replacement difficult, which is what motivated the development of E2Former.

In Equiformer, the architecture is inherently edge-centric and hierarchical. To compute an attention weight ($\alpha_{ij}$), it first combines node features ($h_i, h_j$) into an edge feature $m_{ij}$, followed by computing an expensive tensor product between $m_{ij}$ and geometric information $Y(r_{ij})$, followed by a set of linear and non-linear layers. Then Equiformer calculates a second stage tp $\alpha_{ij} m_{ij} \otimes Y(r_{ij})$, further compounding model’s complexity.

E2Former, by contrast, is inherently node-centric. We use a standard, efficient inner-product attention where the weights ($\alpha_{ij}$) are calculated without any tensor products. The Wigner 6j tensor product is decoupled from the attention and is performed efficiently at the node level. This architectural choice is crucial; it preserves the node-scaling tensor-products of our method.

We appreciate the reviewer’s suggestion regarding integration into existing architectures. While we acknowledge, as noted by Reviewer ejRh, that implementing in modern architectures like Equiformer or MACE is non-trivial, we commit to integrating Wigner 6j into the simpler architecture (SE(3)-Transformer variant/TFN) and will report results in the final revised manuscripts.

Thank you for your detailed review and helpful feedback. We address your concerns below and describe the improvements we will make in our revision.

W1 - 2 We sincerely thank the reviewer for their exceptionally thoughtful and precise observation. The reviewer is correct to highlight the distinction between spherical and solid harmonics, and we confirm that our model does, in fact, use solid harmonics.

Our implementation realizes this by computing the normalization factor $\|r\|^l$ externally and applying it to the spherical harmonic basis functions $Y^l(\hat{r})$. The empirical correctness of our model, which the reviewer correctly inferred, stems from this choice.

We agree this is an important detail that was insufficiently clarified. We will revise the manuscript to state our use of solid harmonics explicitly. Furthermore, we will add a clarification on the design of the filter functions, noting the need to adjust the radial component via scaling ($r^{-l}g_l(r)$) to counteract the implicit weighting from the solid harmonic basis.

We are grateful for this feedback, which significantly sharpens both the theoretical exposition and practical guidance of our paper.

W3 We sincerely thank the reviewer for their insightful and constructive comments. The reviewer has astutely highlighted the accessibility challenges associated with the advanced mathematical formalism in our paper. While these concepts are integral to the rapidly evolving field of equivariant deep learning at NeurIPS, we fully agree that enhancing accessibility is essential to maximize the impact of our work.

In response to this valuable feedback, we propose the following steps to improve clarity and accessibility:

First, as suggested, we will develop a standalone tutorial script. This script will demonstrate a practical “toy task,” clearly illustrating the mechanics of our method and numerically validating our theoretical claims.

Second, following the reviewer’s advice, we will introduce a new appendix section explicitly connecting abstract concepts—such as solid harmonics and Clebsch-Gordan coefficients—to straightforward polynomial functions and tangible numerical tensors. For example, we will explicitly present the solid harmonics for $l=1$ and $l=2$ in their Cartesian polynomial forms to effectively bridge abstract mathematical concepts to familiar, concrete functions. Additionally, we will enhance the clarity of these concepts in the main text itself.

We believe this approach will improve the transition from theory to practical implementation. We thank the reviewer once again for this excellent guidance, which will substantially enhance both the impact and reach of our work.

Q2 Thank you for your insightful observation regarding the relationship between our Wigner 6j convolution approach and the Euclidean Fast Attention method by Frank et al. (2024). Your understanding is exactly correct!

The fundamental distinction lies in the scope and target of optimization. Our method provides a general factorization framework for the tensor product operation itself, which constitutes the computational bottleneck in equivariant neural networks. This factorization is agnostic to the specific architecture in which it is deployed—it can accelerate tensor products whether they appear in attention mechanisms, message-passing layers, or other equivariant operations such as those in MACE-like models. In contrast, the Euclidean Fast Attention method specifically targets the linearization of the dot-product attention mechanism through kernel approximations.

These two acceleration strategies operate at different levels of the computational hierarchy and are indeed orthogonal to each other. The Euclidean Fast Attention reduces the quadratic complexity of attention mechanisms to linear complexity in sequence length, while our Wigner 6j convolution reduces the complexity of the underlying tensor product operations that appear within each attention computation (or any other equivariant operation). Therefore, the two methods can be naturally combined: one could apply Euclidean Fast Attention to achieve linear scaling in sequence length while simultaneously using our Wigner 6j factorization to accelerate the tensor products within the attention computation.

We recognize the potential for such combinations and have explicitly identified this as a promising direction in our future work section. The integration of both techniques could potentially yield multiplicative performance improvements, particularly for large-scale equivariant models that employ attention mechanisms. We will clarify this distinction and the complementary nature of these approaches in our revised related work section to better position our contribution within the broader landscape of acceleration methods for equivariant neural networks.

Thank you for your detailed review and helpful feedback. We address your concerns below and describe the improvements we will make in our revision.

W1-2 Q2-3 We agree with the reviewer that a more thorough comparison with relevant baselines strengthens the paper. We apologize for the omission of the SCN architecture [1] in our discussion and bibliography. In our revised manuscript, we have:

1. Added a proper citation and discussion of the SCN model in the related work and experiments sections.
2. Expanded our results on the Open Catalyst 2022 (OC22) dataset to include the S2EF-Total test split. This allows for a direct comparison with more models, including eSCN.

The new results are presented in the table below and have been integrated into the main paper. E2Former demonstrates highly competitive performance, outperforming eSCN and GemNet-OC on energy predictions and achieving strong results on forces, all while maintaining its significant computational efficiency advantages.

Q1 You are correct. Thank you for this insightful comment and for spotting this inconsistency.

The relationship should indeed be a strict equality, not an isomorphism. The main point of our derivation, particularly for Theorem 3.3, is that the expressions are strictly equal, and using the isomorphism symbol $\cong$ is misleading and detracts from this key fact.

We will revise the manuscript to change the expression on page 5 from $A \otimes (B \otimes C) \cong (A \otimes B) \otimes^{\mathrm{6j}} C$ to $A \otimes (B \otimes C) = (A \otimes B) \otimes^{\mathrm{6j}} C$. As you rightly pointed out, this equality follows directly, and the subsequent mention of "effective commutativity" is unnecessary in this context. We will remove this justification to avoid confusion and improve the clarity of the argument.

Minor comment Thank you for identifying these notational inconsistencies. You are correct—the notation should be $l$. These typographical errors will be corrected promptly. We appreciate your thorough review and attention to detail.

Thank you for your detailed review and helpful feedback. We address your concerns below and describe the improvements we will make in our revision.

W.1 The reviewer raised an important and insightful question regarding the clarity of expressiveness of the Wigner 6j convolution. We deeply appreciate this thoughtful observation and acknowledge that this fundamental aspect warrants thorough clarification. To address this valuable concern, we emphasize that our Wigner 6j convolution maintains the exact expressive power equivalent to conventional SO(3) convolution methods.

Theoretical Equivalence: As demonstrated in Theorems 3.2 and 3.3, the Wigner 6j convolution provides exact representation of the complete angular basis used by traditional CG-based convolutions. The methods are therefore mathematically equivalent by construction. The Wigner 6j recoupling represents a unitary transformation of angular momentum channels, ensuring no loss in representational capability while delivering significant computational benefits.

Empirical Validation: The results presented in Figure 2(b–d) provide evidence of this theoretical equivalence, demonstrating that our convolution produces numerically identical outputs compared to conventional SO(3) convolution. This empirical validation confirms the expressiveness equivalence while achieving significant efficiency improvements that make the approach practically advantageous.

Following the reviewer's feedback, we will add a dedicated subsection to the main text that directly addresses expressiveness equivalence. This addition will clarify the theoretical foundation and demonstrate that representational power is fully preserved.

W.2 We greatly appreciate the reviewer's insightful perspective on existing CG acceleration methods. Their observation highlights an important distinction that underscores the unique value of our contribution.

Prior methods such as eSCN (ICML ’23), Gaunt Tensor Product (ICLR ’24), and SPHNet (ICML ’25) optimize individual tensor product computations, significantly accelerating performance yet remaining constrained by $O(|\mathcal{E}|)$ scaling, as each edge still necessitates a separate tensor product calculation. E2Former introduces a paradigm shift by drastically reducing the total number of tensor products. Leveraging Wigner 6j recoupling theory, we transform computations from an edge-based approach to a node-based formulation, reducing complexity from $O(|\mathcal{E}|)$ to $O(|\mathcal{V}|)$. Importantly, our method is orthogonal to and complements existing Clebsch–Gordan (CG) optimizations; the smaller set of node-based tensor products can further exploit these optimized implementations, offering compounded efficiency gains. We will explicitly highlight this synergistic potential as part of our future work.

W.3 We appreciate the reviewer's valuable feedback regarding the need for clearer organization and more prominent definition of "linear scaling" in the introduction. To address this concern, we will add an explicit definition of "linear scaling" at the beginning of the introduction, clarifying that this refers to computational complexity that scales linearly with respect to the number of tensor products, which is $O(\|\mathcal{V}\|)$ where $\|\mathcal{V}\|$ represents the number of nodes. We apologize that this fundamental distinction was not made sufficiently clear in the original manuscript.

W.4 We appreciate the reviewer's suggestion to include recent works on invariant polynomial networks, specifically "Weisfeiler Leman for Euclidean Equivariant Machine Learning" (ICML 2024) and "A New Perspective on Building Efficient and Expressive 3D Equivariant Graph Neural Networks" (NeurIPS 2024). These works offer complementary perspectives through different invariance structures and polynomial expansions. We will incorporate discussion of these works in Section 6 (Related Work) and identify their integration with our node-centric convolution as a promising future research direction.

Dear reviewer,

It seems you decided to increase your score to 4 but forgot to do so.

AC