Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials

Abstract: In this paper, we provide a systematic investigation of SO(2) theory to machine learning interatomic potentials (MLIPs) and identify the limitations of conventional SO(2) Linear architectures relative to SO(3) Clebsch-Gordan Tensor Products (CGTP). Building on these insights, we propose direct Cartesian construction and recursive Clebsch-Gordan construction of Wigner D-matrices and introduce two novel interaction building blocks. First, we propose the Edge Complex Product Basis based on Generalized Asymmetric Contraction, a new formulation for many-body expansion that directly constructs higher-order interactions on edges through complex-valued equivariant multiplications. Second, we introduce Radial Rotary Complex Attention(RRA), which enhances extrapolation performance and surpasses existing attention vector formulations. We also introduce several improvements to the Atomic Cluster Expansion module. Building on these advances, we train our models on OMat24, sAlex, and MPTrj, and introduce TECE-OAM-RRA-1.0, which achieve state-of-the-art (SOTA) performance on the Matbench Discovery.

• The paper presents an Edge Cluster Expansion that enhances many-body interactions through asymmetric contractions, improving extrapolation in ML potentials.
• It introduces a recursive Clebsch-Gordan construction for efficient, low-error Wigner D-matrix generation, enabling scalable SO(2)-equivariant operations.
• Radial Rotary Complex Attention is developed to incorporate radial features into geometric graphs, achieving state-of-the-art RMSE benchmarks on molecular datasets.

Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials

Overview

The paper "Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials" (2607.10664) systematically advances equivariant architectures for machine learning interatomic potentials (MLIPs), focusing on efficient and expressive SO(2)-based tensor operations, novel many-body interaction schemes, and improved attention mechanisms. The work rigorously evaluates the theoretical, architectural, and practical limitations of conventional SO(2) Linear networks relative to SO(3) Clebsch-Gordan Tensor Products (CGTP), introduces direct Cartesian and recursive Clebsch-Gordan constructions of Wigner D-matrices, and develops new modules including Edge Cluster Expansion (ECE) and Radial Rotary Complex Attention (RRA). Empirical validation on major benchmarks demonstrates state-of-the-art accuracy and favorable computational scaling.

SO(2) Tensor Product Theory and its Limitations

SO(2)-equivariant methods, mapping global SO(3) features to local SO(2) frames, have become widely adopted due to their computational efficiency. However, under identical angular resolutions, SO(2) models generally exhibit inferior extrapolation and accuracy compared to CGTP-based approaches outside of universal MLIP benchmarks. The paper identifies that these discrepancies are primarily architectural rather than theoretical, with correct handling of radial information substantially broadening the feasible design space.

Key findings include:

• Complex-valued weights in SO(2) layers generally do not guarantee O(2)-equivariance, breaking reflection symmetry and harming extrapolation unless restricted to real-valued weights (w1 scheme).
• Channel mixing (uvSO2Linear) allows arbitrary transformations but weakens representational power without direct access to radial path information, whereas uuSO2Linear preserves strong inductive bias and extrapolation, achieving parity with uuu SO(3) Tensor Product given identical parameter counts.

Efficient Construction of Wigner D-Matrices

The evaluation and manipulation of Wigner D-matrices, necessary for equivariant operations, are typically computationally expensive due to the exponential scaling of Cartesian intermediates. This work proposes a recursive Clebsch-Gordan scheme that avoids $3^\ell$ complexity, yielding polynomial storage scaling and superior runtime as $\ell_{\max}$ increases.

Figure 1: Runtime comparison of Wigner-D matrix construction methods; recursive Clebsch-Gordan achieves best scaling with increasing $\ell_{\max}$.

The recursive approach also ensures lower numerical error and simplicity (<100 lines of code), with direct compatibility for specialized GPU kernels. This efficiency improvement is critical for scaling equivariant MLIPs to high angular degrees and large batch sizes.

Edge Cluster Expansion: Generalized Many-Body Interactions

The Edge Cluster Expansion module introduces generalized asymmetric contractions of edge features using Clebsch-Gordan coefficients, enabling flexible many-body expansions directly on edges. Unlike conventional symmetric contractions, ECE supports variance stabilization and antisymmetric interactions, overcoming limitations in completeness and expressivity found in grid-based Gaunt Tensor Product and standard atomic cluster expansions.

Strong numerical results demonstrate:

• ECE with SO(2) Linear can distinguish higher-body geometric counterexamples otherwise indistinguishable for lower body-order expansions, validating the completeness of the architecture (see Table~\ref{tab:body_order_experiments} in the paper).
• Application of ECE before aggregation ensures retention and generation of high angular frequencies, critical for resolving structures with $M$-fold rotational symmetry; ACE post-aggregation cannot recover lost angular information.

Radial Rotary Complex Attention (RRA)

The Radial Rotary Complex Attention module, inspired by RoPE, explicitly incorporates radial information into geometric graph attention, modulating QK similarities by radial phase and bias terms. Empirical ablation shows that:

• Radial modulation is essential for extrapolation, outperforming standard attention vector and non-attention schemes.
• Complex-valued phase terms in RRA do not harm extrapolation, as long as weights in SO(2) Linear layers remain real-valued.

Quantitative benchmarks (Table~\ref{tab:RRCA} in the paper) indicate minimum RMSEs for energies and forces in molecular datasets using RRA, with further improvements via edge nonlinearities.

Architectural and Implementation Refinements

Additional model enhancements include:

• Optimized coupled Atomic Cluster Expansion (ACE) via group linear experts, reducing memory and computational overhead for large channel counts.
• Equivariant residual network and layer norm strategies to improve training stability, especially in deep ACE-based models.
• Careful choice of optimizer: Muon and SOAP accelerate convergence and in-domain accuracy but compromise extrapolation, with Adam providing the best extrapolation.

The final model, TECE-OAM-RRA-1.0, achieves state-of-the-art results on Matbench Discovery, ranking first or second across all major metrics (F1 score, thermal conductivity, and structure optimization) with the highest composite performance score among leading architectures.

Computational Scaling and Practical Runtime

Benchmarking shows that EquflashV2, TECE-OAM-RRA-1.0, and EquiformerV3+DeNS-OAM represent the top three models in terms of computational efficiency on large-scale molecular dynamics, with SO(2)-based operator fusion remaining promising for future developments in speed and memory scaling.

Practical and Theoretical Implications

The findings affirm the critical importance of tensor product completeness, explicit many-body modeling, and physical inductive biases (e.g., radial attention) for achieving high extrapolation and accuracy in MLIPs. The recursive Clebsch-Gordan construction and ECE methodology provide effective scalability, facilitating the adoption of SO(2) architectures in large-scale materials foundation models. Architectural choices and weight parameterization have nontrivial effects on extrapolation, with real-valued SO(2) weights strongly preferred for robustness. The demonstrated SOTA performance on universal benchmarks and efficient runtime renders SO(2)-ECE models highly competitive for practical deployment.

Future Directions

The main open challenges include operator fusion for large-scale SO(2) Linear operators, comprehensive node-level SO(2) constructions (to address discontinuity), and robust handling of dataset inconsistencies (e.g., Hubbard U corrections, DFT parameter mismatches). Given the demonstrated architectural flexibility and computational scalability, future work should focus on deep integration with efficient edge-fusion libraries, comprehensive benchmarking in low-data regimes, and further theoretical exploration of equivariant contraction schemes.

Conclusion

This paper consolidates and extends SO(2) theory in MLIPs, introducing efficient Wigner-D constructions, Edge Cluster Expansion, and Radial Rotary Complex Attention. These innovations resolve practical limitations in completeness, extrapolation, and computational scaling, validated by strong numerical benchmarks and SOTA accuracy on Matbench Discovery (2607.10664). The work highlights the importance of theoretical rigor and architectural choices for robust interatomic potential modeling, setting the stage for future developments in operator fusion, equivariant graph networks, and materials foundation models.