1 Introduction↩︎

Equivariant machine learning interatomic potentials (MLIPs) have become indispensable tools for molecular dynamics simulations. However, the high computational cost of the Clebsch-Gordan tensor product (CGTP) has motivated the development of alternative tensor product formulations aimed at reducing computational complexity. Representative examples include the grid-based Gaunt Tensor Product (GTP) [1], which exhibits a path-merging property that avoids the explicit evaluation of multiple tensor product paths. Building upon this formulation, EquiformerV3 [2] incorporates nonlinearity together with many-body interactions [3], resulting in the SwiGLU-S\(_2\) architecture. Despite its computational efficiency, GTP inherently excludes antisymmetric interactions and introduces equivariance errors, although these limitations appear to have limited impact on practical fitting performance. The Matrix Tensor Product [4] similarly exploits path merging while accelerating computation through modified contraction schemes, providing an elegant and efficient simplification of conventional tensor products. By explicitly rotating features into an edge-aligned local frame, the SO(2) Linear framework not only provides a sparse implementation of CGTP operations, but also introduces a fundamentally new class of equivariant operators that go beyond conventional tensor products. This approach exhibits improved scaling behavior and enables additional operations on edges at relatively low computational cost. In contrast, irreducible Cartesian tensor products and contractions [5], [6] perform equivariant convolutions directly in Cartesian space and have demonstrated favorable efficiency-accuracy trade-offs at low angular degrees. Wigner-6j Convolutions [7] further maximize the fusion of edge operations into node-level computations, resulting in highly memory-efficient convolutions. Finally, Vector Spherical Tensor Product (VSTP) approaches [8], [9] combine vector spherical harmonics with tensor signals to address the missing antisymmetric interactions in grid-based methods.

Among these emerging tensor product formulations, the SO(2) theory has become the most widely adopted in practice. Nevertheless, the theoretical understanding of SO(2) methods remains incomplete. In particular, it is often observed that under the same angular resolution, SO(2)-based models exhibit lower extrapolation and accuracy than CGTP-based approaches. However, models built upon SO(2) operators demonstrate superior accuracy and faster convergence on the universal MLIP (uMLIP) benchmark. Although SO(2) Linear exhibits more favorable asymptotic scaling behavior in principle, this does not necessarily make it the preferred choice in practice. Beyond the aforementioned limitations above, practical implementations typically require two SO(2) Linear operations together with nonlinear transformations applied at the edge level. This reliance on complex edge-level operations prevents the development of highly optimized, general-purpose edge-fusion libraries analogous to OpenEquivariance [10] and CuEquivariance [11], which are available for CGTP-based models. In our view, the modest accuracy gains observed in uMLIP do not justify the additional implementation complexity and the loss of fusion opportunities introduced by these edge-level nonlinear operations. Conversely, if such nonlinear edge operations are omitted, the advantages of SO(2) Linear largely disappear, and one can simply employ CGTP directly without introducing the additional complexity of SO(2) operators. As we will demonstrate in subsequent experiments, the reduced extrapolation capability primarily arises from model-design choices rather than from the SO(2) operations themselves. In particular, an appropriate treatment of radial information enables a substantially broader model-design space without compromising extrapolation performance. By contrast, the accuracy improvements observed on the uMLIP benchmark can largely be attributed to the use of edge nonlinearities and higher edge body orders. In this work, we revisit the complex-valued nature of SO(2) equivariance. From this perspective, we analyze the design principles and constraints of learnable weights under complex-valued representations and characterize the class of operations that preserve SO(2) equivariance. In the experimental section, we first introduce a faster and numerically more stable method for constructing Wigner-D matrices by leveraging the relationship between spherical tensors and irreducible Cartesian tensors. Next, we compare the behavior of uuSO(2) Linear and uvSO(2) Linear, identifying the regimes in which each formulation exhibits its respective strengths and limitations. We then investigate the impact of complex-valued weight parameterizations in SO(2) Linear operators by examining three representative schemes: w1, w1_w1, and w1_w2. Building upon the SO(2) Tensor Product, i.e., complex-valued equivariant multiplication, we introduce an edge cluster expansion (ECE) block that performs many-body expansions directly on edges. In this part, we further demonstrate the advantages of ECE and SO(2) operations through experiments on M-fold rotationally symmetric structures [12] and k-body counterexamples [13]. In addition, we show how complex-conjugate multiplication can be used to construct an SO(2)-invariant (\(m=0\), complex-valued) from query and key representations, which naturally leads to the formulation of Radial Rotary Attention. Finally, we evaluate the proposed designs through experiments on the Matbench Discovery benchmark [14].

2 Preliminaries↩︎

2.1 SO(3) Irreps \(\leftrightarrow\) SO(2) Irreps↩︎

The theoretical framework related to SO(3)/O(3) operations in machine learning was first systematically implemented in e3nn [15]–[17]. Concretely, an irreducible representation (irrep) of SO(3) is indexed by the degree \(\ell\) and order \(m\), and learnable features can be represented as a tensor of dimension \(2\ell+1\). In machine learning applications, an additional channel dimension \(c\) is typically introduced. Therefore, a complete SO(3) irrep is indexed by \((c, \ell, m)\). In general, different values of \(\ell\) may correspond to different numbers of channels; however, in this work, we do not explicitly distinguish between them. The concept of SO(2)-related representations was first introduced in eSCN [18], where the authors defined the SO(2) Linear operation. Subsequently, QHNetV2 [19] further developed the theoretical framework of SO(2) for Hamiltonian prediction tasks, introducing operations such as SO(2) Tensor Product and SO(2) Gate. The SO(2) group, which describes rotations in the plane, is abelian. Therefore, over the complex numbers, all its irreducible representations are \(1\)-dimensional. For SO(2) features, the action of a rotation by \(\theta\) is given by multiplying features by the unit complex number \(e^{i m \theta}, m \in \mathbb{Z}\), \[\mathbf{x_m} \mapsto e^{i m \theta} \mathbf{x_m}.\] In geometric graph neural networks, for edge-level SO(3) features, if we rotate all edges to a common reference direction and apply the corresponding Wigner-D transformation to the features, the rotational degrees of freedom reduce to a single remaining angle. In this case, the features can be regarded as local SO(2) features, allowing us to apply SO(2)-related operations (from global SO(3) frame to local SO(2) frame). Similarly, by applying the transpose of the Wigner-D matrix, the features can be transformed back.

2.2 SO(2) Operations↩︎

In general, complex-valued SO(2) features are obtained by rotating real-valued SO(3) features. However, operations on SO(2) features are more naturally interpreted from the perspective of complex numbers. In this work, all SO(3) features are treated in the real-valued representation, while all SO(2) features are interpreted in the complex-valued representation. In practice, we can use torch.view_as_complex and torch.view_as_real [20] to convert between real-valued and complex-valued features.

Given complex-valued weights \(\mathbf{w} = w_1 + i w_2\) and complex-valued features \(\mathbf{x_m} = x_{+m} + i x_{-m}\), complex multiplication \(\mathbf{w} \cdot \mathbf{x_m}\) commutes with rotations and is therefore SO(2)-equivariant. Specifically, for a complex-linear map \(L_W(\mathbf{x_m})=W\mathbf{x_m}\), we have \[\label{eq:linear95map} L_W\!\left(e^{\mathrm{i}m\theta}\mathbf{x_m}\right) = We^{\mathrm{i}m\theta}\mathbf{x_m} = e^{\mathrm{i}m\theta}W\mathbf{x_m} = e^{\mathrm{i}m\theta}L_W(\mathbf{x_m}).\qquad{(1)}\] This implies that we simultaneously modulate both the magnitude and the phase of the complex-valued features. In particular, if we impose the constraint \(w_1 = w_2\), this implies that we only apply a fixed phase rotation together with an arbitrary magnitude transformation. If we instead constrain \(w_2 = 0\), then the operation only modulates the magnitude while preserving the phase unchanged. We refer to these three weighting schemes for complex-valued features as w1_w2, w1_w1, and w1, respectively.

2.2.0.1 SO(2) Linear.

Once the embedding mechanism of complex-valued weights is understood, it becomes immediately clear that the essence of the SO(2) Linear operation is to mix features with the same order \(m\) across different channels \(c\) and degrees \(\ell\) in the complex domain. This operation allows us to scalarize information with higher angular resolution, \((\ell > 0,\, m = 0) \rightarrow (\ell = 0,\, m = 0)\), thereby achieving an effect analogous to an SO(3) tensor product. See Section 3.7 for details.

2.2.0.2 SO(2) Tensor Product

Given two complex-valued SO(2) features \(\mathbf{x}_{m_1}\) and \(\mathbf{x}_{m_2}\), their complex multiplication remains equivariant under rotations when \(m_3 = m_1 + m_2\) or \(m_3 = m_1 - m_2\), because \[\mathbf{x}_{m_1} \mathbf{x}_{m_2} \mapsto e^{i(m_1+m_2)\theta} \mathbf{x}_{m_1} \mathbf{x}_{m_2},\] and \[\mathbf{x}_{m_1} \overline{\mathbf{x}}_{m_2} \mapsto e^{i(m_1-m_2)\theta} \mathbf{x}_{m_1} \overline{\mathbf{x}}_{m_2},\] where \(\overline{\mathbf{x}}_{m_2}\) denotes the complex conjugate of \(\mathbf{x}_{m_2}\). Although some previous work calls SO(2) Linear a “tensor product”, it is actually different from the SO(2) Tensor Product. They handle different parts of the model, and both are important.

2.2.0.3 SO(2) Gate.

Different from the SO(3) gate, where the \(\ell=0, m=0\) features are used as gates, the SO(2) Gate uses the \(m=0\) features from all degrees \(\ell\) as gates. This can be regarded as a special case of the SO(2) tensor product, so we do not discuss it in detail.

2.3 Special Cases of Node-wise SO(2) Operations↩︎

SO(2) operations are typically performed on edges, where the edge direction naturally provides a local frame. If equivalent operations could be carried out directly on node representations, additional computational savings might be achieved. One possible direction is represented by frame-averaging approaches [19], [21], which construct a local reference axis from the target atom and its nearby neighboring source atoms. While such formulations have proven effective for Hamiltonian prediction tasks, their applicability to interatomic potentials is less clear. The selection of a particular neighbor may introduce discontinuities when the local atomic environment changes. Averaging over all neighboring atoms can alleviate this issue; however, the resulting formulation becomes conceptually similar to aggregating edge-wise operations at the node level, thereby diminishing the computational benefits originally sought. A second direction originates from attempts to simplify Wigner-6j Convolutions [22]. Wigner-6j Convolutions exploit solid harmonics and binomial expansions to separate the angular dependence of source and target atoms. Since CGTP remains computationally expensive, the factorized representation enabled by solid harmonics naturally motivates the construction of local reference axes directly from atomic positions. Therefore, although node-wise SO(2) operations constitute a promising design space, a universally applicable design paradigm that generalizes across diverse architectures is still lacking.

3 Results↩︎

3.1 O(2) Irrep Notation↩︎

Although our model is only required to be local SO(2)-equivariant, it is useful to first introduce the notation for irreducible representations of O(2). This notation will be used in the subsequent discussion of Weight Types 3.5, Edge Cluster Expansion 3.6, and Radial Rotary Attention 3.8. Under a planar reflection, complex feature is mapped to its conjugate: \[\label{eq:reflection} \mathbf{x_m} \mapsto \overline{\mathbf{x_m}}.\tag{1}\] Therefore, for \(m>0\), the corresponding \(\mathrm{O}(2)\) irrep is still labeled by the order \(m\), while the special case is \(m=0\). Since the rotation action is trivial, reflection can act either as \(+1\) or \(-1\), giving two different one-dimensional irreps: \(0e\) and \(0o\). Thus, the irreps of \(\mathrm{O}(2)\) can be summarized as \[0e,\quad 0o,\quad 1m,\quad 2m,\quad 3m,\quad \ldots\] where \(m>0\) labels two-dimensional irreps, and only the \(m=0\) case splits into even and odd parity.

3.2 Cartesian Construction of Wigner-D Matrices↩︎

Spherical tensors and Cartesian tensors provide two equivalent descriptions of three-dimensional equivariant features. A spherical tensor of degree \(\ell\), denoted by \(\mathbf{x}^{(\ell)} \in \mathbb{R}^{2\ell+1}\), transforms under a rotation \(\mathbf{R}\in\mathrm{SO}(3)\) according to the Wigner-D matrix: \[\mathbf{x}^{(\ell)\prime} = \mathbf{D}^{(\ell)}(\mathbf{R})\mathbf{x}^{(\ell)}. \label{eq:spherical95rotation}\tag{2}\] In contrast, a rank-\(\ell\) Cartesian tensor \(\mathbf{T}^{(\ell)}\in\mathbb{R}^{3\times\cdots\times3}\) transforms by applying the rotation matrix independently to each Cartesian index: \[T^{(\ell)\prime}_{i_1\cdots i_\ell} = \sum_{j_1,\ldots,j_\ell} R_{i_1j_1}\cdots R_{i_\ell j_\ell} T^{(\ell)}_{j_1\cdots j_\ell}. \label{eq:cartesian95rotation}\tag{3}\] Here, we propose two methods for computing the Wigner-D matrix from rotation matrices. The first method is more intuitive and easier to understand, while the second provides a more efficient implementation.

3.2.0.1 Direct Cartesian Construction.

The irreducible Cartesian tensor decomposition (ICTD) [23] provides an orthonormal basis transformation between degree-\(\ell\) spherical tensors and rank-\(\ell\) symmetric traceless Cartesian tensors. Let \(\mathbf{C}_{\ell} \in \mathbb{R}^{3^\ell\times(2\ell+1)}\) denote the path matrix. The spherical-to-Cartesian and Cartesian-to-spherical transformations are respectively given by \[\operatorname{vec}\!\left(\mathbf{T}^{(\ell)}\right) = \mathbf{C}_{\ell}\mathbf{x}^{(\ell)}, \qquad \mathbf{x}^{(\ell)} = \mathbf{C}_{\ell}^{\mathsf T} \operatorname{vec}\!\left(\mathbf{T}^{(\ell)}\right). \label{eq:ictd95transform}\tag{4}\] Applying the Cartesian rotation in Eq. 3 between the two basis transformations gives a direct construction of the Wigner-D matrix: \[\mathbf{D}^{(\ell)}(\mathbf{R}) = \mathbf{C}_{\ell}^{\mathsf T} \mathbf{R}^{\otimes\ell} \mathbf{C}_{\ell}. \label{eq:cartesian95wigner}\tag{5}\]

3.2.0.2 Limitations of the Direct Cartesian Construction.

Although Eq. 5 is numerically stable, its direct evaluation introduces a Cartesian intermediate tensor containing \(3^\ell\) components. When evaluated for a batch of \(E\) edges, the intermediate storage scales as \(\mathcal{O}\left(E\,3^\ell(2\ell+1)\right)\). Consequently, the direct construction becomes increasingly expensive as \(\ell\) grows, despite the final Wigner-D block containing only \((2\ell+1)^2\) entries.

3.2.0.3 Recursive Clebsch-Gordan Construction.

We note that the path matrix itself is generated recursively by coupling a degree-one representation with a degree-\((\ell-1)\) representation \(1\otimes(\ell-1)\rightarrow\ell\). This structure allows us to avoid the explicit rank-\(\ell\) Cartesian tensor. Let \[\mathbf{K}^{(\ell)}_{abM} = \operatorname{Wigner3j}(1,\ell-1,\ell)_{abM} \label{eq:recursive95cg}\tag{6}\] denote the fixed Clebsch-Gordan coupling coefficients, where \(a\in\{1,2,3\}\), \(b\in\{1,\ldots,2\ell-1\}\), and \(M\in\{1,\ldots,2\ell+1\}\). Starting from \[\mathbf{D}^{(0)}(\mathbf{R}) = 1, \qquad \mathbf{D}^{(1)}(\mathbf{R}) = \mathbf{R}, \label{eq:recursive95base}\tag{7}\] higher-degree Wigner-D matrices can be constructed recursively as \[D^{(\ell)}_{MN}(\mathbf{R}) = (2\ell+1) \sum_{a,b,c,d} K^{(\ell)}_{abM} D^{(1)}_{ac}(\mathbf{R}) D^{(\ell-1)}_{bd}(\mathbf{R}) K^{(\ell)}_{cdN}. \label{eq:recursive95wigner}\tag{8}\] The factor \(2\ell+1\) accounts for the normalization convention of the Wigner-\(3j\) symbols. Equation 8 is algebraically equivalent to Eq. 5 , but it avoids the exponentially large Cartesian representation. The largest intermediate tensor in the recursive contraction scales as \(\mathcal{O} \left(E\,(2\ell-1)(2\ell+1)\right)\) which replaces the exponential scaling with polynomial scaling. Moreover, the recursive method has lower numerical error and avoids the inverse trigonometric functions and angular singularities present in the original eSCN formulation. It consists only of fixed-shape tensor contractions, making it suitable for specialized GPU kernels. For the asymptotic behavior of different methods, please refer to Fig 1. As shown in the figure, our method achieves the best asymptotic performance. Notably, UMA [24] and DPA4 [25] constructs the Wigner-D matrices using quaternions and custom operators implemented in more than 1000 lines of code. In contrast, our method requires fewer than 100 lines of code. Since the accuracy of the initial rotation matrix \(R\) is also crucial for the recursive construction, we likewise construct \(R\) from quaternions.

3.3 SO(3) Paths and SO(2) Paths↩︎

For SO(3) operations, we typically refer to a tuple \((\ell_{\text{in1}}, \ell_{\text{in2}}, \ell_{\text{out}})\) as a path. Each path is associated with a learnable weight parameter. The weights are defined with respect to the output degree \(\ell_{\text{out}}\), while the feature mixing over \(\ell_{\text{in1}}\) and \(\ell_{\text{in2}}\) is performed through a linear layer.

For SO(2) operations, in this work, we refer to a tuple \((\ell_{\mathrm{in}}, \ell_{\mathrm{out}}, m)\) as a path. Similarly, each path is associated with a learnable weight.

3.4 uvSO2Linear and uuSO2Linear↩︎

To the best of our knowledge, all existing SO(2) architectures currently adopt the SO(2) Linear introduced in eSCN [18]. However, considering that this operation is intended to replace tensor products with spherical harmonics, we believe that the correspondence between the two is not entirely straightforward. In conventional SO(3) frameworks, radial information is typically associated with the SO(3) paths arising from the uuu SO(3) tensor product, with each path receiving its own path-dependent weighting. In contrast, under the SO(2) framework, radial information is usually incorporated before the SO(2) Linear operation through SO(3) real multiplication or SO(2) complex multiplication where the radial weights are real-valued (i.e., with zero imaginary part). As a result, radial parameters are generally not directly constructed for each individual SO(2) path. We therefore argue that this mechanism more closely resembles uvw SO(3) tensor product, where radial dependencies cannot be incorporated directly for each path due to the large number of paths. Therefore, we observe that the SO(2) Linear layer introduced in eSCN is, in essence, a fully connected linear transformation that does not rely on path-wise radial weights. By analogy with the terminology used in SO(3) tensor product, we refer to this operation as uvSO2Linear. Therefore, we argue that the truly “equivalent” counterpart of the uuu SO(3) tensor product should instead be uuSO2Linear. For a fair comparison, we use a single SO(2) Linear operation throughout and do not introduce any edge-wise nonlinearities. Under this formulation, we compare the performance of the two approaches with the same number of parameters. As shown in Table ¿tbl:tab:uuSO2Linear?, uuSO2Linear achieves accuracy and extrapolation performance comparable to those of uuu SO(3) tensor product. In contrast, uvSO2Linear, which does not directly depend on radial information, exhibits a noticeable loss in expressive power.

3.4.0.1 Advantages and Limitations.

By more effectively leveraging radial information, uuSO2Linear offers stronger expressive power and extrapolation capabilities. However, this formulation does not permit channel mixing. In contrast, uvSO2Linear allows arbitrary channel transformations. When sufficient training data are available, the improvements in accuracy and extrapolation offered by uuSO2Linear may become marginal. In this regime, the arbitrary channel transformations enabled by uvSO2Linear can be combined with edge nonlinearities, making it more flexible from an architectural perspective.

We consider uuSO2Linear important because, when designing operators in an SO(2) framework, a single linear layer should ideally serve as a replacement for the uuuSO3TensorProduct. While both SO(2) Linear formulations are theoretically valid, using uvSO2Linear in this setting sacrifices compatibility with edge nonlinearities. On the other hand, introducing such nonlinearities precludes the memory-efficient optimization of fusing edge computations.

3.5 Complex-Valued Weights Break O(2)/O(3) Equivariance and Weaken Extrapolation↩︎

We seek to understand why complex-valued weights arise in force-field models whose inputs and outputs are both real-valued. This naturally raises the question of whether the complex phase provides meaningful modeling capacity or merely introduces redundant degrees of freedom. To investigate this issue, we compare the accuracy and extrapolation performance of three different SO(2) weight types. The comparison is conducted under three architectures: Radial Rotary Attention, the attention-vector formulation adopted by EquiformerV3 [2], and a non-attention architecture. For each architecture, the best result is highlighted in bold in Table ¿tbl:tab:RRCA?. From the results, we observe that introducing phase variations provides little benefit on this dataset. In several cases, it even leads to degradation in both accuracy and extrapolation performance.

In general, complex-valued weights do not guarantee \(\mathrm{O}(2)\)-equivariance. Taking the w1_w2 as an example, its complex weight is written as \(w=w_1+\mathrm{i}w_2\). Under the reflection \(-I\), the complex SO(2) feature \(\mathbf{x}_m\) is conjugated, namely \((-I)\mathbf{x}_m=\overline{\mathbf{x}_m}\). Therefore, \[\begin{align} L_w\!\left((-I)\mathbf{x}_m\right) &= (w_1+\mathrm{i}w_2)\overline{\mathbf{x}_m}, \\ (-I)L_w(\mathbf{x}_m) &= (w_1-\mathrm{i}w_2)\overline{\mathbf{x}_m}. \end{align}\] These two expressions are not equal unless \(w_2=0\). Thus, w1_w2 is generally only SO(2)-equivariant, while the real-weight w1 case is compatible with O(2) reflection symmetry. This implies that, restricting the model to use only real-valued weights allows us to enforce the correct behavior by construction. In contrast, an SO(3)-equivariant model does not impose this constraint and can only rely on the training data to learn the appropriate behavior under global inversion. With sufficiently large and diverse datasets, the model may approximately learn this property from data. However, in the low-data regime, the lack of an explicit constraint becomes much more pronounced, leading to a clear performance gap.

3.6 Generalized Symmetric Contraction and Uncoupled Edge Cluster Expansion↩︎

The Transformer architecture has achieved tremendous success in large language models (LLMs). However, we remain cautious about its direct application to machine learning interatomic potentials. There are two main reasons for this. First, adopting a Transformer architecture typically requires complicated operations explicitly at the edge level. For conventional CGTP frameworks, this can become inefficient and may prevent the direct use of mature operator-fusion libraries such as OpenEquivariance [10] and CuEquivariance [11]. Second, the success of models such as NequIP [26], MACE [3], EquFlash [27], and TACE [5] shows that strong performance can already be achieved without graph softmax mechanisms.

In light of expressive Transformer-based models such as EquiformerV3 [2] and PET [28], we believe that explicitly modeling source-target atom relationships during large-scale pretraining is at least unlikely to be harmful. Therefore, we introduce edge-level many-body interactions through an uncoupled edge cluster expansion (ECE). The key idea is to perform tensor contractions directly on edge features using generalized coupling coefficients. The symmetric contraction was first introduced in MACE [3] with O(3) equivariance. Here, we generalize this construction and extend it to SO(2). More broadly, the framework is not limited to SO(2) and O(3): any tensor representation equipped with Clebsch-Gordan-like coupling coefficients can be incorporated into the same formulation. We write the input features in layout as \(\mathbf{x}^{(r)} \in \mathbb{R}^{B \times F \times C}\), where \(B\) is the batch size, \(F\) is the feature dimension of the representation, and \(C\) is the channel dimension. For SO(2) features, \(l\) index are merged into channel dimension.

3.6.0.1 Generalized Coupling Coefficients.

Consider a representation decomposed into blocks indexed by \(\lambda\), where each block has dimension \(d_\lambda\). For O(3), \(\lambda\) may correspond to \((\ell,p)\); for SO(2), \(\lambda=m\). A pairwise coupling is described by coefficients \[C^{\lambda_{\mathrm{out}}}_{\lambda_1\lambda_2} \in \mathbb{R}^{d_{\lambda_1}\times d_{\lambda_2}\times d_{\lambda_{\mathrm{out}}}}.\] More generally, a correlation order \(\nu\) coupling path is denoted as \[\eta_\nu = (\lambda_1,\lambda_2,\ldots,\lambda_\nu; \lambda_{\mathrm{out}}),\] with generalized coefficients \[\mathcal{C}^{\lambda_{\mathrm{out}}}_{\eta_\nu} \in \mathbb{R}^{d_{\lambda_1}\times\cdots\times d_{\lambda_\nu} \times d_{\lambda_{\mathrm{out}}}}.\] These coefficients can be constructed recursively from pairwise couplings, as in the generalized Clebsch-Gordan construction used by MACE [3]. This makes symmetric contraction a general algebraic operation rather than a construction specific to O(3).

3.6.0.2 SO(2) as a Special Case.

For SO(2), the representation in real basis has \[d_m = \begin{cases} 1, & m=0,\\ 2, & m>0. \end{cases}\] The allowed pairwise output frequencies follow from complex multiplication: \[m_{\mathrm{out}} = m_1 + m_2, \qquad m_{\mathrm{out}} = |m_1-m_2|.\] For example, for the sum path with \(m_1,m_2>0\), \[\begin{align} (a+\mathrm{i}b)(c+\mathrm{i}d) = (ac-bd)+\mathrm{i}(ad+bc), \end{align}\] which gives the nonzero real basis coefficients \[C_{000}=1,\quad C_{110}=-1,\quad C_{011}=1,\quad C_{101}=1.\] To maintain stable output variance, we normalize the coefficients: \[\sum_{a,b} \left( C^{m_{\mathrm{out}}}_{m_1m_2,abo} \right)^2 = 1.\]

3.6.0.3 Limitations of Symmetric Contraction.

The standard symmetric contraction repeatedly contracts the same feature: \[\mathbf{x} \otimes \mathbf{x} \otimes \cdots \otimes \mathbf{x}.\] This has two limitations. First, symmetric contraction is not always variance stable. Even if the Clebsch-Gordan coefficients are normalized, repeatedly using the same feature tensor introduces correlations between different polynomial orders, such as \(\mathbf{x}\), \(\mathbf{x}^2\), and \(\mathbf{x}^3\). As a result, the variance of the summed output cannot be normalized in a straightforward manner. This effect becomes more severe for high-order expansions, large numbers of paths, or extreme local structures, and may lead to unstable training. Second, symmetric features inherently admit fewer computational paths than asymmetric implementations. In addition, the symmetry of the Clebsch-Gordan coefficients causes antisymmetric components to vanish automatically under O(3). Although the contributions from higher-order terms gradually decrease, and symmetric contraction may appear more intuitive, our objective is to perform a many-body expansion. Therefore, it is necessary to relax these constraints appropriately.

3.6.0.4 Generalized Asymmetric Contraction.

To relax these constraints, we introduce asymmetric contraction. The input can be a list of feature tensors \[\left[ \mathbf{x}^{(1)}, \mathbf{x}^{(2)}, \ldots, \mathbf{x}^{(\nu_{\max})} \right].\] If only one tensor is provided, it is reused for all orders. If fewer than \(\nu_{\max}\) tensors are provided, the last tensor is reused for the remaining orders. For a path \(\eta_\nu\), the product basis is \[\mathbf{B}^{(\nu)}_{b,\eta_\nu,c,o} = \sum_{a_1,\ldots,a_\nu} \mathcal{C}^{\lambda_{\mathrm{out}}}_{\eta_\nu,a_1\cdots a_\nu o} \prod_{r=1}^{\nu} \mathbf{x}^{(r)}_{b,\lambda_r,c,a_r}.\] The coefficients may be shared internal weights \(\mathbf{W}^{(\nu)} \in \mathbb{R}^{N_{\mathrm{path},\nu}\times C}\) or externally provided edge/node-dependent (non)linear weights \(\mathbf{W}^{(\nu)} \in \mathbb{R}^{B\times N_{\mathrm{path},\nu}\times C}\). The output is accumulated as \[\mathbf{y}_{b,\lambda,c,o} = \sum_{\nu=1}^{\nu_{\max}} \sum_{\eta_\nu\rightarrow \lambda} \frac{1}{ \sqrt{ N_{\mathrm{p}}(\lambda,\nu) N_{\mathrm{o}}(\lambda) }} \mathbf{W}^{(\nu)}_{b,\eta_\nu,c} \mathbf{B}^{(\nu)}_{b,\eta_\nu,c,o},\] where \(N_{\mathrm{p}}(\lambda,\nu)\) is the number of paths for order \(\nu\) contributing to output block \(\lambda\), and \(N_{\mathrm{o}}(\lambda)\) is the number of correlation orders contributing to \(\lambda\). This keeps the output variance close to one when different input branches are approximately independent. We refer to the recursive algorithm introduced in MACE [3] as the Horner-recursive implementation. It constructs features of successive correlation orders recursively by contracting the input features with a batched version of the generalized coupling coefficients associated with the highest correlation order. However, the size of the batched coefficient scales as \(F^{\nu}\). Consequently, the intermediate tensors also scale as \(F^{\nu}\), which becomes prohibitively expensive when \(F\) is large. For example, in the Cartesian representation, the feature dimension at angular degree \(\ell\) grows as \(3^\ell\), resulting in a total dimension of \(F_{\mathrm{cart}}=\sum_{\ell=0}^{\ell_{\max}}3^\ell\). This exponential growth quickly makes the computation impractical and largely defeats the purpose of the recursive formulation. In contrast, the spherical representation contains \(2\ell+1\) components at each angular degree \(\ell\), giving a total dimension of \(F_{\mathrm{sph}}=\sum_{\ell=0}^{\ell_{\max}}(2\ell+1)=(\ell_{\max}+1)^2\). Under the SO(2) basis, only order with \(m\leq m_{\max} \text{or } \ell_{\max}\) are retained, and the corresponding dimension is \(F_{\mathrm{SO(2)}}=1+2m_{\max}\). Therefore, this approach is most suitable for node-level operations, or for settings in which memory consumption is not a limiting factor, where it can provide a substantial computational speedup. This comparison also highlights the potential advantage of SO(2) ACE. However, as discussed below, there is currently no general node-level construction for SO(2) because of the discontinuity issue. Consequently, the Horner-recursive implementation, whose intermediate tensors scale exponentially with the correlation order, cannot be directly applied to edge-level operations. We therefore design a memory-efficient algorithm, presented in Algorithm 3, which avoids materializing the full high-order intermediate tensors. This design trades computational speed for substantially reduced memory consumption. It is particularly well suited to ECE, where the feature dimension remains small.

The use of ECE allows the edge basis to depend either on an individual source or target atom, as discussed in Section 3.7, or jointly on both the source and target atoms. Here, we focus on the latter formulation. Since many-body expansions are typically implemented using channel-wise operations, all angular orders \(m\) are required to share the same number of channels. This requirement can be satisfied either by applying uvSO2Linear or padding. More specifically, given the edge-aligned source and target features, \(\mathbf{x}_{j}^{t-1}\) and \(\mathbf{x}_{i}^{t-1}\), from the \((t-1)\)-th layer, we construct the first-order edge basis, \({}^{\nu=1}\mathbf{E}_{ij}^{t}\), at the \(t\)-th layer as \[{}^{\nu=1}\mathbf{E}_{ij}^{t} = \text{uvSO2Linear} \left( \mathbf{x}_{i}^{t-1} \;\Vert\; \mathbf{x}_{j}^{t-1} \right),\] where \(\Vert\) denotes channel-wise concatenation. Besides introducing complex-valued weights, we also have what appears to be an additional option: allowing the real and imaginary parts of the weight to depend separately on the source and target atoms. However, as discussed above, we believe that real-valued weights are preferable. Therefore, the greatest flexibility of ECE lies in the choice of edge-dependent real weights. \[\tikz[baseline=(eq.base)]{ \node[inner sep=0pt] (eq) { \underbrace{w_1}_{\text{source atoms}} + \quad i\underbrace{w_2}_{\text{target atoms}} }; \draw[red, line width=0.8pt] ([yshift=5pt]eq.west) -- ([yshift=5pt]eq.east); }\]

3.6.0.5 Nonlinearity and Extrapolation.

As one of the early influential equivariant models, NequIP [26] applies a gating operation after message aggregation. Subsequently, BoTNet [29] conducted an ablation study by removing the gating mechanism and other nonlinear components. The results showed that comparable accuracy, and in some cases slightly improved accuracy, could be achieved without gating. Motivated by this observation, subsequent models, such as MACE [3], also avoid introducing additional nonlinear transformations after aggregation. In fact, as the scale of the dataset increases, the fitting capacity of purely linear models becomes limited. Therefore, nonlinear modules are necessary. However, improved accuracy should not be obtained at the cost of a substantial loss in extrapolation capability or by introducing unconstrained modules. Instead, appropriate nonlinearities should be incorporated while preserving the model’s extrapolation performance. Accordingly, the design principle of TACE is not to perform module ablations solely based on in-distribution training accuracy. Instead, we evaluate each candidate module on the 3BPA benchmark at 1200K. Any module that increases the force RMSE beyond 65 meV/Åis excluded from the final architecture.

3.6.0.6 Differences Between Edge/Atomic Cluster Expansion.

Here, we discuss the differences between ECE and ACE. In a local frame, edge features (\(m > 0\)) cannot be aggregated arbitrarily. Consequently, we cannot employ a bilinear operation analogous to that used in Allegro [30] (global SO(3) frame) to introduce an explicit aggregation step before ECE (local SO(2) frame). Compared with ACE, the many-body information in ECE is represented more implicitly. Moreover, ECE is generally not applied in the first layer, since the corresponding edge features have not yet aggregated information from all neighboring atoms. The role of ECE is to increase the effective body order of the model while explicitly modeling the relationship between the source and target atoms. When combined with appropriate nonlinearities, ECE can build nonlinear messages. The incorporation of nonlinearities in ECE follows the same procedure as in ACE and details can be found in Section 3.9.

3.7 Angular Resolution and Body Order↩︎

In this subsection, we analyze why Edge Cluster Expansion can increase angular resolution and body order, namely, whether it is complete in this sense. Specifically, we study the ability of these modules to resolve \(M\)-fold rotationally symmetric structures and whether they can increase the effective body order of the representation, and, if so, through what mechanism this increase is achieved. Joshi et al. identify two key factors that are critical to the a [12] identifies two key factors that are critical to the accuracy of MLIPs: the angular resolution \(L_{max}\) (or, more precisely, \(m_{max}\) in this work) and the body order (\(\nu\)). In particular, they validate this conclusion through two experiments. We redesign these experiments and use them to illustrate the importance of ECE from the perspectives of global SO(2) frame and global SO(3) frame.

3.7.0.1 Global SO(2) Frame.

The first experiment compares models with different values of \(m_{\max}\) in distinguishing 2-dimension \(M\)-fold rotationally symmetric structures with different global orientations [12]. Since the structures considered here are planar, we describe them using the SO(2) group rather than the SO(3) group adopted. Moreover, we provide an explicit mathematical derivation that offers a more rigorous treatment, replacing the less formal discussion in the original work. Since the two structures differ only in their global orientations and are otherwise symmetric, we denote the orientation by \(\phi\) and ignore the radial component. We directly use the edge vector \(\hat{\mathbf{r}}_{ij}\) to encode angular information: specifically, we compute its azimuthal angle \(\theta_{k}\) (Due to the rotational symmetry of the structure, the edges that were originally distinguished by the pair index \(ij\) can now be indexed by \(k\)), which represents the angular positions of the \(M\) neighbors surrounding the central atom and construct the corresponding complex SO(2) angular features: \[\mathbf{x}_{m}^{k}(\phi) = e^{\mathrm{i}m\theta_k(\phi)} = e^{\mathrm{i}m\left(\phi+\frac{2\pi k}{L}\right)}, \qquad k=0,\ldots,M-1.\] Assume that the model constructs the final representation of the central atom by directly aggregating information from its neighboring atoms: \[A_m(\phi) = \sum_{k=0}^{M-1} e^{\mathrm{i}m\theta_k(\phi)} = \begin{cases} M e^{\mathrm{i}m\phi}, & m\equiv 0\pmod M,\\ 0, & \text{otherwise}. \end{cases}\] It follows immediately that if the input angular resolution satisfies \(m_{\max}<M\), then the aggregated information at center atom is independent of the global orientation \(\phi\). Moreover, applying an SO(2) Atomic Cluster Expansion after the aggregation cannot resolve this issue, because the relevant angular information has already vanished during the aggregation: \[A_{m_1}(\phi)A_{m_2}(\phi)=0, \qquad m_1,m_2 < M .\] In contrast, when the SO(2) Edge Cluster Expansion product is applied before aggregation, complex multiplication combines angular frequencies and thereby generates higher angular orders. As a result, the model can obtain representations that depend on the global orientation: \[E_{m_1}(\phi)E_{m_2}(\phi)\neq0, \qquad m_1+m_2 = M .\] It is worth noting that although we do not explicitly include learnable weights in the notation here, the conclusion remains unchanged when such weights are used. Since all features are represented in a global SO(2) frame, applying SO(2) Linear before or after aggregation is equivalent. Accordingly, this experiment exposes a potential limitation of SO(2) ACE in this setting, while also highlighting the importance of SO(2) ECE, even though body order is not involved in this case.

3.7.0.2 Global SO(3) Frame.

Using the 2-body, 3-body, and 4-body counterexamples from Ref. [13], we illustrate how ECE progressively constructs higher-body-order features. Each \(k\)-body counterexample consists of two geometric graphs that are indistinguishable under \(k\)-body scalarization. In general, without an explicit many-body expansion, a single message-passing layer is equivalent to a \(2\)-body scalarization and therefore cannot distinguish the \(2\)-body counterexample. Increasing the correlation order in the many-body expansion progressively improves the completeness of the model. For an ACE-based architecture, the computation follows the pipeline: \[\begin{align} \mathrm{Convolution} &\rightarrow \mathrm{Aggregation} \rightarrow \mathrm{ACE} \rightarrow \mathrm{Node\;Readout}. \end{align}\] For the edge-based setting, however, additional design is required to properly test the effective body order of the model. Specifically, we use the following pipeline: \[\begin{align} &\mathrm{Convolution} \rightarrow \mathrm{Aggregation} \rightarrow \mathrm{Gather} \rightarrow \mathrm{Rotate} \rightarrow &&\\ &\mathrm{ECE} \rightarrow \mathrm{SO(2) Linear} \rightarrow \mathrm{Rotate\;Back} \rightarrow \mathrm{Edge\;Readout}. && \end{align}\] Using ECE also requires first aggregating the two-body scalarized information. Therefore, the pipeline up to the aggregation step is the same as in the ACE setting. However, since ECE is performed on edges and under an \(\mathrm{SO}(2)\) representation, additional gather, rotate, and rotate-back operations are required. It is important to note that once an SO(2) Linear operation is applied in this rotated frame, higher-order components with \(l>0,m=0\) can be mixed into the scalar component with \(l=0,m=0\). This is equivalent to introducing an interaction with the edge \(ij\), thereby increasing the effective body order by one. As a result, the model becomes capable of distinguishing the \(2\)-body counterexample. However, it should be emphasized that stacking multiple SO(2) Linear layers still does not further increase the effective body order, and therefore the model remains limited to distinguishing the \(2\)-body counterexample. The reason why SO(2) Linear is necessary in this pipeline is that channel-wise ECE alone, although capable of increasing the body order, cannot transfer the resulting higher-order information into the scalar (l=0, m=0). Due to the block-diagonal structure of the Wigner-D matrices, information in the \(l>0,m=0\) components cannot flow into the \(l=0,m=0\) component without an SO(2) Linear mixing step. In this case, the operation only improves the angular resolution, as in the previous experiment, while the higher-order scalarized information is not incorporated into the final readout function. This motivates the pipeline design described above. Therefore, since the SO(2) Linear operation introduces exactly one additional effective body order, the results in Table 1 should be interpreted carefully. Finally, in our experiments, the SO(2) Linear operation is implemented as a special case that permutes the components \(l>0,m=0\) with scalar component \(l=0,m=0\). As shown in Table 1, ECE indeed increases the effective body order, thereby demonstrating its effectiveness and completeness in distinguishing higher-body-order geometric structures.

3.8 Radial Rotary Complex Attention↩︎

We argue that geometric graphs are inherently weighted by radial information. Discarding this information and relying solely on simple QK inner products or contractions with attention vectors may impair the model’s extrapolation capability. Ideally, attention should explicitly depend on both the similarity between nodes and their relative distances. Drawing inspiration from Rotary Position Embeddings (RoPE) [31], we propose Radial Rotary Complex Attention (RRA). Given the complex-valued nature of SO(2) representations, we define the QK similarity as the real part (0e) of the complex inner product. We discard the imaginary component because it transforms as \(0o\). We further introduce a phase term and a bias term, which directly modulate the QK inner product as a scale and a shift, respectively. Notably, although the scale term is, in principle, allowed to be complex-valued, using a complex-valued phase term restricts the model to SO(2) equivariance in the local frame and breaks O(2) equivariance. In contrast to the weights in SO(2) Linear layers, we observe that the use of a complex-valued phase term does not affect model extrapolation and accuracy. We denote SO(2) features in local frame as \(\mathbf{x}^{ij}_{j,m,h,c}\) and \(\mathbf{x}^{ij}_{i,m,h,c}\), respectively, where \(m\) is the SO(2) order, \(h\) denotes the attention head and \(c\) denotes the channel within each head. The complex query, key, and value features are obtained using independent uvSO2Linear with weight type w1, respectively. Notably, the phase term and bias shift are deliberately designed as zero-initialized linear mappings that depend only on radial basis. We intentionally exclude element information from these two terms, so that they encode purely distance-dependent effects. Element-dependent feature similarity is instead captured through the complex inner product. \[\begin{align} \mathbf{Q}_{m,h,c}^{ij} &= \text{uvSO2Linear} \left(\mathbf{x}^{ij}_{i,m,h,c}\right), \\ \mathbf{K}_{m,h,c}^{ij} &= \text{uvSO2Linear} \left(\mathbf{x}^{ij}_{j,m,h,c}\right), \\ \mathbf{V}_{m,h,c}^{ij} &= \text{uvSO2Linear} \left(\mathbf{x}^{ij}_{i,m,h,c}, \mathbf{x}^{ij}_{j,m,h,c}\right). \end{align}\] The attention logit is defined as: \[s^{ij}_{h} = b^{ij}_{h} + \frac{\tau_h}{\sqrt{D}} \sum_{m=0}^{m_{\max}} \sum_{c} \operatorname{Re} \left[ e^{\mathrm{i}\theta_{h}} \, \overline{\mathbf{Q}}^{ij}_{m,h,c} \, {\mathbf{K}}^{ij}_{m,h,c} \right].\] In particular, we define radial phase \(\theta_{h}\), radial bias \(b_{h}\) and and temperature \(\tau_h\) as: \[\theta_{m,h}(r) = 0 \, \text{or} \, m \pi \cdot \tanh \left(b_{h} + \sum_{n}W_{h}^{n}j_0^n(r)\right),\] \[b_{h}(r) = b_{h} + \sum_{n}W_{h}^{n}j_0^n(r),\] \[\tau_h = \tau_{\min} + \frac{\tau_{\max}-\tau_{\min}}{1+e^{-w_h}}.\] where W and b represent the learnable weight and bias, respectively, n represent the number of radial basis. Here, \(\operatorname{Re}(\cdot)\) denotes the real-part operator. The optional temperature term \(\tau_{h}\) controls the sensitivity of the softmax function across different attention heads. In practice, we typically set \(\tau_{\min}=0.25\) and \(\tau_{\max}=4.0\). The attention weight for each head is then computed over all incoming edges of node \(i\). Meanwhile, an envelope function is incorporated to imporve the discontinuity issue reported in EquiformerV3 [2], \[\alpha^{ij}_{h} = \frac{ \exp \left( s^{ij}_{h} \right)\cdot f_{\text{cut}}(r_{ij}) }{ \sum_{k \in \mathcal{N}(i)} \exp \left( s^{ik}_{h} \right)\cdot f_{\text{cut}}(r_{ik}) }.\] To examine the effect of RRA, we compare different attention mechanisms, with the results summarized in Table ¿tbl:tab:RRCA?. The best results are highlighted in green. Notably, even when using uvSO2Linear, which may compromise extrapolation performance, our model still achieves strong results. This suggests that the strong inductive bias introduced by RRA can effectively compensate for this limitation, thereby supporting our claim.

3.9 Coupled Atomic Cluster Expansion.↩︎

We still adopt the original coupled ACE implementation used in the Cartesian TACE. Specifically, “coupled” means that, after the tensor product, element-dependent weights are applied to couple information across all channels. In practice, using only \(48\)-\(64\) channels is sufficient to achieve an accuracy comparable to that obtained by the uncoupled implementation with \(128\)-\(256\) channels. However, its limitations are also evident. First, the element-dependent parameterization introduces a large number of weights. Second, when the number of channels exceeds \(64\), the coupled formulation leads to unnecessary computational and memory overhead. Therefore, we improve the ACE implementation mainly in two aspects. First, we make coupled ACE feasible when increasing the channel dimension. Second, we introduce additional nonlinearities without destroying the resulting many-body structure.

3.9.0.1 From Dense Linear to Group Linear.

Unlike standard sparse Mixture-of-Experts models [32], where each selected expert is still a linear transformation over all input channels and using top-\(k\) selection, our goal is to keep computation feasible as the number of channels increases. Therefore, such a MoE formulation does not directly address our problem. Instead, we uniformly partition the input and output channels into \(E\) element-aware expert groups, so that each expert only operates on a subset of channels. To compensate for the reduced cross-channel communication, we additionally introduce an element-agnostic shared expert that couples all channels, \[\underbrace{ \begin{bmatrix} \mathbf{W}^{(z)}_1 & \mathbf{0} & \cdots & \mathbf{0} \\ \mathbf{0} & \mathbf{W}^{(z)}_2 & \cdots & \mathbf{0} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{0} & \mathbf{0} & \cdots & \mathbf{W}^{(z)}_E \end{bmatrix} }_{\text{aware group experts}} + \underbrace{ \mathbf{W}_{\mathrm{shared}} }_{\text{agnostic shared expert}}.\] Here, \(\mathbf{W}\) denotes learnable weights, and \(z\) denotes the element type on which the group-expert weights depend.

3.9.0.2 GLU-style implementation and Nonlinear Coefficients.

To introduce additional nonlinearity into the product basis block without disrupting its many-body structure, we adopt an EquflashV2-inspired strategy, in which scalar 0e features are used to parameterize the tensor product weights [33], [34]. An linear projection first generates two asymmetric atomic bases together with scalar, feature-dependent ACE coefficients: \[\operatorname{Linear}_{\mathrm{up}}(\mathbf{A}_{i}) = \mathrm{A}_{i}^{0e} \oplus \mathbf{A}_{i}^{\mathrm{node}} \oplus \mathbf{A}_{i}^{\mathrm{base}},\] where both \(\mathbf{A}_{i}^{\mathrm{node}}\) and \(\mathbf{A}_{i}^{\mathrm{base}}\) contain all retained irreps, including their \(0e\) components. The first-order correlation is defined directly by the node features, while the second-order correlation is obtained by coupling them with the base features augmented by a constant scalar: \[\mathbf{B}_{i}^{(1)} = \mathbf{A}_{i}^{\mathrm{node}}\] \[\mathbf{B}_{i}^{(2)} = \underbrace{ \mathbf{A}_{i}^{\mathrm{node}} \otimes_{\mathrm{SO(3)}} \left( \mathbf{1}_{0e} \oplus \mathbf{A}_{i}^{\mathrm{base}} \right) }_{ \text{modulated by } \operatorname{SiLU}\!\left(\mathrm{A}_{i}^{0e}\right) }.\] The constant scalar \(\mathbf{1}_{0e}\) and the base features are used to approximately produce second- and third-order correlations, respectively, despite involving only a single tensor-product operation. The resulting correlation features are then concatenated by \(\mathbf{B}_{i} = \mathbf{B}_{i}^{(1)} \oplus \mathbf{B}_{i}^{(2)}\). Finally, the grouped and shared experts are combined using \[\frac{ \operatorname{Linear}_{\mathrm{group}}(\mathbf{B}_{i}) + \operatorname{Linear}_{\mathrm{shared}}(\mathbf{B}_{i}) }{\sqrt{2}}.\]

3.10 Residual Network and Layer Norm↩︎

In equivariant networks, an important detail is that the first layer in the message-passing process differs from subsequent layers. During the first edge-level tensor-product operation, the tensor product is typically formed between 0e features (such as element embedding) and spherical harmonics. Consequently, prior to the first message-passing step, only scalar information is available. Under such circumstances, directly treating the 0e featrues as residual connections [35] is generally unfavorable. We further observed that, in the absence of the LayerNorm layer, employing two residual connections independently in both the interaction module (\(\mathbf{A}\) basis) and the product module (\(\mathbf{B}\) basis) does not yield optimal performance in our framework. Moreover, this design becomes more sensitive to the number of layers and is more prone to training instability. Instead, superior accuracy is consistently achieved when only a single residual pathway is used. In particular, both the \(\mathbf{A} \rightarrow \mathbf{B}\) and \(\mathbf{B} \rightarrow \mathbf{B}\) connection schemes outperform the \(\mathbf{B} \rightarrow \mathbf{A} \rightarrow \mathbf{B}\) design, regardless of whether the identity branch is implemented through direct addition or (element-dependent) linear transformations. We also note that recently proposed residual-network designs, such as Attention Residuals (AttnRes) [36], have been shown to further improve model accuracy. However, since the depth of current uMLIPs remain substantially smaller than that of modern LLMs, the benefits brought by such residual formulations are comparatively limited. Therefore, although these approaches may provide additional accuracy gains, we do not adopt them in the present work in order to reduce computational overhead. We also observed that, compared with non-ACE architectures, ACE-based models are generally more sensitive to outliers and tend to exhibit greater training instability when multiple layers are stacked. To address this issue, we adopt the EquivariantMergeLayerNorm implementation introduced in EquiformerV3 [2]. In our view, this normalization strategy currently provides one of the most effective solutions for equivariant architectures, as it preserves the scale differences across different irreps while improving training stability. Following the same considerations applied to the residual design, we omit LayerNorm in the first layer and do not apply it after the final layer. This is because LayerNorm is primarily used to stabilize training, whereas an additional normalization at the output stage of MLIPs may limit the model’s expressive accuracy. Hence, unless the training process is highly unstable, introducing LayerNorm after the final layer is unnecessary.

3.11 Effects of Optimizers on Model Extrapolation↩︎

Recently, emerging optimizers such as Muon [37] and SOAP [38] have attracted increasing attention as potential alternatives to conventional optimizers such as Adam [39] and AdamW [40]. For example, Harari et al. [41] systematically investigated SOAP and Muon and showed that both optimizers, whether used individually or in combination, consistently achieved higher accuracy. Similarly, DPA-4 [25], [42] employed the Muon optimizer and demonstrated strong performance on the MPtrj dataset. These emerging optimizers often provide improved in-domain accuracy and faster convergence. Nevertheless, we emphasize that these optimizers should be used with caution, as both Muon and SOAP may compromise extrapolation performance to some extent. In terms of extrapolation capability, Adam performs best, followed by SOAP and then Muon.

3.12 Matbench Discovery↩︎

Matbench Discovery [14] has become a widely used benchmark for evaluating foundation models in materials science. In particular, OMat24, sAlex, and MPtrj [43]–[45] are among the most advanced and widely used materials datasets currently available. We therefore train our model on these datasets to demonstrate its representational capacity. As shown in Table ¿tbl:tab:matbench?, our model achieves state-of-the-art (SOTA) performance among compared architectures. Notably, among the three most important metrics, F1 score for crystal stability prediction, thermal conductivity prediction, and structure optimization, our model achieved the best performance on two tasks and ranked second on the remaining one, resulting in the highest overall score.

\begin{table}[!h]

\begin{center}
  
  \begin{tabular}{lcccccr}
    \toprule
    Model & CPS & F1($\uparrow$) & $\kappa$SRME ($\downarrow$)& RMSD& Params \\
    \midrule
    \rowcolor{green!10}
    TECE-OAM-RRA-1.0 & \textbf{0.908} & 0.929 & \textbf{0.093} & \textbf{0.058} & 222M \\
    EquFlashV2 & 0.907 & 0.929 & 0.094 & \textbf{0.058} & 44.9M \\
    EquiformerV3+DeNS-OAM & 0.902 & \textbf{0.931} & 0.118 & 0.060 & 30.3M \\
    GRACE-3L-OAM-L & 0.900 & 0.925 & 0.121 & \textbf{0.058} & 42.1M \\
    PET-OAM-XL & 0.898 & 0.924 & 0.119 & 0.060 & 730M \\
    TACE-OAM-L & 0.889 & 0.910 & 0.126 & 0.061 & 82.9M \\
    eSEN-30M-OAM & 0.888 & 0.925 & 0.170 & 0.061 & 30.2M \\
    EquFlash & 0.888 & 0.919 & 0.158 & 0.060 & 28.7M \\
    Nequip-OAM-XL & 0.886 & 0.906 & 0.125 & 0.063 & 32.1M \\
    MatRIS-10M-OAM & 0.877 & 0.921 & 0.218 & 0.060 & 10.4M \\
    SevenNet-Omni-i12 & 0.873 & 0.906 & 0.192 & 0.062 & 54.9M \\
            GRACE-2L-OAM-L & 0.865 & 0.883 & 0.169 & 0.064 & 26.4M \\
    ORB v3 & 0.860 & 0.905 & 0.210 & 0.075 & 25.5M \\
    DPA-4.0.1-Pro-MPtrj & 0.840 & 0.857 & 0.211 & 0.069 & 22.8M \\
    Allegro-OAM-L & 0.840 & 0.895 & 0.319 & 0.065 & 9.7M \\
                                                                                                                                                                                                            \bottomrule
  \end{tabular}
  
\end{center}
\vskip\caption{Top models on MatBench Discovery~\cite{Matbench}. For models with the same architecture, we retain only the variant with the highest accuracy. CPS is a weighted combination of the F1 score, thermal-conductivity metric ($\kappa$SRME), and RMSD, and thus represents the model with the strongest overall performance. Full results are provided in Appendix Table~\ref{tbl:tab:full95matbench}.}\label{tab:matbench} -0.1in

\end{table}

3.13 Model Speed↩︎

We are pleased to see the emergence of systematic evaluations of both the accuracy and computational efficiency of uMLIPs, as these benchmarks provide valuable guidance for the community when selecting models across a wide range of model sizes [46]. Nevertheless, we would like to emphasize an important point. Improving both model extrapolation and fitting accuracy remains essential (see Section 3.6.0.5). A more expressive architecture can achieve accuracy with a smaller and faster model, while its highest-accuracy variant can also serve as an effective teacher for knowledge distillation. In practice, smaller variants of SOTA architectures can therefore offer a better combination of accuracy and speed, since increasing model size often yields only marginal improvements on accuracy-oriented benchmarks. Existing speed benchmarks do not always evaluate models under their optimal configurations. In some cases, appropriate acceleration techniques may not be enabled because of limited familiarity with the model implementation. To provide a fairer and more representative comparison, we benchmarked the top three models [2], [34] on Matbench Discovery using the best available configuration for each model and compared their practical runtime performance. As shown in Table ¿tbl:tab:speed?, the top three models are ranked in terms of computational speed and memory efficiency as follows: EquflashV2 \(>\) TECE-OAM-RRA-1.0 \(>\) EquiformerV3+DeNS-OAM. This result is consistent with our expectations. As discussed previously, architectures based on SO(2) should employ uuSO2Linear to enable operator fusion. In comparison, EquflashV2 adopts more efficient SO(3) edge-fused operators, resulting in improved computational performance. Nevertheless, SO(2) architectures may still offer a higher performance ceiling and more favorable scaling behavior. We therefore look forward to further developments in efficient SO(2) kernel.

4 Discussion, Limitations, and Outlook↩︎

4.1 Diatomic Benchmark↩︎

We also observed several interesting phenomena in the diatomic benchmark. The widely used OAM-family datasets exhibit discontinuities in their potential energy surfaces when O or F occurs with certain transition-metal elements, due to the use of different Hubbard U treatments [48]. In addition, OrbV3 [49] also reported that the rattle subset of OMat24 can adversely affect performance on diatomic benchmarks. Indeed, the top 3 models on Matbench Discovery, including TECE-OAM-RRA-1.0, EquFlashV2, and EquiformerV3+DeNS-OAM consistently exhibit relatively poor performance on the diatomic benchmark. By contrast, models with lower fitting capacity, such as TACE-OAM-L, NequIP-OAM-XL, and MACE-MPA-0, tend to perform better on this benchmark. We attribute this behavior primarily to limitations in the training data: as model capacity and fitting accuracy increase, the model becomes more likely to fit dataset-specific artifacts and unphysical features, thereby reducing its robustness in regimes such as diatomic configurations.

In summary, we systematically review, consolidate, and extend SO(2) theory for machine learning interatomic potentials, while validating several implementation details that are important in practice. We further introduce Cartesian direct and recursive Wigner-D constructions, Edge Cluster Expansion (ECE), and Radial Rotary Complex Attention (RRA), and release a SOTA materials foundation model built upon these components. Our results suggest that the proposed design choices provide complementary benefits rather than a universally optimal configuration. When pursuing the highest accuracy in large-scale pretraining, ECE with edge nonlinearities and RRA provides the most effective strategy. When memory efficiency is the primary concern, uuSO2Linear is generally preferable because of its stronger inductive bias, expressive capability, and compatibility with operator fusion. From an architectural perspective, an important direction for future work is operator fusion for single-layer SO(2) Linear operators. During the course of this work, we also confirmed several limitations in the existing datasets, including data redundancy, inconsistent DFT parameter settings, and inappropriate treatments of Hubbard U corrections. For example, TACE-OAM-L, an intermediate model developed during the construction of our SO(2) model family, was trained on only approximately random 20% of the training data due to a code bug. Despite this substantial reduction in training data, it still exhibited remarkably strong performance.

Author contribution↩︎

W.X. and P.H. conceived the project and guided the research. Z.X. prepared code, equations, figures and tables. All authors edited and revised the manuscript.

Acknowledgment↩︎

This work was supported by the National Natural Science Foundation of China NSFC (22433004 and 22403064), the open research fund of Key Laboratory of Precision and Intelligent Chemistry, and ShanghaiTech University. We are also grateful for the computing time provided by the the HPC Platform of ShanghaiTech University.

5 Training Details↩︎

All TECE models were trained in torch.float32 precision using 32 NVIDIA H20 GPUs. During pretraining, we employed DeNS [50]and stochastic depth [51] as additional regularization techniques. DeNS was mainly used with the direct model to accelerate convergence, rather than to improve the final accuracy. With appropriately chosen hyperparameters, the direct and conservative models achieved comparable performance upon convergence. The DeNS hyperparameters were kept consistent with those used in EquiformerV3 [2]. For optimization, we used the Muon optimizer implemented in DeepMD-kit [25], [37], [42] to accelerate convergence. Muon is not essential. In fact, AdamW may be a better choice because it generally provides better extrapolation performance. However, considering the computational cost, we used Muon to accelerate convergence. Training the TECE-OAM-RRA-1.0 model consumed approximately 7680 NVIDIA H20 GPU-hours in total. This figure should not be interpreted as representative of the intrinsic computational efficiency of the TACE/TECE architecture, which achieves SOTA accuracy and computational performance. The computing platform used in our experiments exhibited suboptimal multi-GPU scaling and relied on outdated CUDA and PyTorch versions, preventing us from using compilation-based acceleration techniques. Consequently, the observed training cost was approximately three times higher than necessary. To provide a more representative evaluation, we benchmarked the computational performance of various TACE architectures on a newer platform. Further details on computational efficiency and inference speed are provided in the main text.

6 Matbench Discovery↩︎

References↩︎