Equivariant Axis-Aligned Sparsification (EAAS)

• Equivariant Axis-Aligned Sparsification (EAAS) is an algebraic technique that converts dense Clebsch–Gordan tensor products into ultra-sparse, axis-aligned block operations for efficient equivariant neural architectures.
• It significantly reduces computational complexity and memory usage, achieving empirical 5–6× speedups and enabling scalable performance in models like E2Former-V2.
• EAAS maintains exact SO(3)-equivariance through strategic axis alignment, sparse coupling, and Wigner-D rotations, proving critical for large-scale biomolecular and materials modeling.

Equivariant Axis-Aligned Sparsification (EAAS) is an algebraic approach for accelerating equivariant neural architectures, notably those involving $\mathrm{SO}(3)$ message passing, by restructuring the heavy Clebsch–Gordan (CG) tensor products into ultra-sparse, axis-aligned blockwise operations plus Wigner-D rotations. EAAS achieves cubic asymptotic, and empirically, order-of-magnitude reduction in compute and memory per edge while maintaining exact $\mathrm{SO}(3)$-equivariance and high modeling expressivity. This innovation is central to the scalability and performance of modern linear-scaling equivariant transformers such as E2Former-V2 and is a key component in universal biomolecular foundation models like UBio-MolFM (Huang et al., 13 Feb 2026, Huang et al., 23 Jan 2026).

1. Algebraic Formulation and Theoretical Underpinnings

Clebsch–Gordan tensor products appear ubiquitously in $\mathrm{SO}(3)$-equivariant message-passing, coupling node features $h^{(\ell_i)} \in \mathbb{R}^{2\ell_i+1}$ and solid harmonics $R^{(\ell_f)}(r)$ into outputs of degree $\ell_o$. Conventionally, this is performed by the full contraction:

$$(h^{(\ell_i)} \otimes R^{(\ell_f)}(r))^{(\ell_o)}_{m_o} = \sum_{m_i, m_f} C^{\ell_o, m_o}_{\ell_i, m_i; \ell_f, m_f} h^{(\ell_i)}_{m_i} R^{(\ell_f)}_{m_f}(r),$$

resulting in an inherent computational and memory complexity of $\mathcal{O}((2\ell_i+1)(2\ell_f+1)(2\ell_o+1))$ per edge.

EAAS exploits the observation that for any vector $r$, there exists a unique $R\in \mathrm{SO}(3)$ mapping $r/\|r\|$ to the $z$-axis. In this axis-aligned frame, all harmonics $R^{(\ell_f)}_{m_f}(r)$ vanish except at $m_f=0$, collapsing the CG contraction to a single $m_f$ component. The contraction thus reduces to a form:

$$(h^{(\ell_i)} \otimes R^{(\ell_f)}(r))^{(\ell_o)}_{m_o} = [P^{(\ell_i, \ell_f \rightarrow \ell_o)} h^{(\ell_i)}]_{m_o} \cdot \|r\|^{\ell_f} \cdot D^{(\ell_o)}(R^{-1})_{m_o, m_o'},$$

where $P^{(\ell_i, \ell_f \rightarrow \ell_o)}$ is a very sparse map corresponding to the reduced coupling, and $D^{(\ell_o)}(R^{-1})$ is the Wigner-D rotation returning to the global frame (Huang et al., 13 Feb 2026), Sec. 3.2.2.

This algebraic reduction trades dense $m$-index sums for (i) an axis alignment, (ii) a lookup and sparse matrix on the $m=0$ channel, and (iii) two Wigner-D multiplications.

2. Sparsity Patterns and Parity Re-Indexing

In the axis-aligned (pole) frame, the CG contraction simplifies:

$$\bigl(h\otimes\mathcal R(\vec r)\bigr)^{(\ell_o)}_{m_o} = \sum_{m_i} C^{(\ell_o,m_o)}_{(\ell_i,m_i),(\ell_f,0)}\, \tilde h^{(\ell_i)}_{m_i},$$

with $\tilde h^{(\ell_i)} = D^{(\ell_i)}(R) h^{(\ell_i)}$. Parity and selection rules further enforce at most one survivor per output $m_o$. Explicitly, for total degree $L_\Sigma = \ell_i+\ell_f+\ell_o$ and the parity index-mapping $\sigma_{L_\Sigma}$:

$$\sigma_{L_\Sigma}(m_o) = \begin{cases} m_o, & L_\Sigma \text{ even}, \ -m_o, & L_\Sigma \text{ odd}, \end{cases}$$

and sign factor $$\varepsilon_{L_\Sigma}(m_o) = \{1,\, -2(-1)^{m_o}\}$$, so

$$\bigl(\mathcal P(\tilde h)\bigr)_{m_o} = \varepsilon_{L_\Sigma}(m_o) C^{(\ell_o,m_o)}_{(\ell_i,\sigma_{L_\Sigma}(m_o)),(\ell_f,0)} \tilde h^{(\ell_i)}_{\sigma_{L_\Sigma}(m_o)}.$$

This deterministic sparsification enables both algebraic and hardware advantages, converting dense contractions into nearly free block-linear operations (Huang et al., 23 Jan 2026).

3. Implementation in Linear-Scaling Equivariant Transformers

EAAS is utilized as the primitive for value-path computation in equivariant attention layers of E2Former-V2. For an edge $i\leftarrow j$, the algorithm proceeds:

1. Compute $R$ mapping $r_{ij}/\|r_{ij}\| \to (0,0,1)$.
2. Rotate source feature blocks $h_j^{(\ell_i)}$ with $D^{(\ell_i)}(R^{-1})$.
3. Extract $m=0$ components, apply blockwise weights $W^{(\ell_i\to \ell_o)}$, and scale by radial functions.
4. Rotate outputs back using $D^{(\ell_o)}(R)$.

In multi-channel settings, block-matrix–Wigner-D products are fused for maximum efficiency. This pipeline wholly avoids materializing large edge-centric tensors, synergizing with custom fused kernels (e.g., Triton) for node-centric memory layouts (Huang et al., 13 Feb 2026), Alg. 1.

4. Equivariance and Invariance Properties

EAAS preserves exact $\mathrm{SO}(3)$-equivariance by construction:

• The initial axis-alignment is $\mathrm{SO}(3)$-covariant.
• The internal operation in the $z$-frame involves only the $m=0$ channel of the spherical harmonics, equivalent to coupling via an $\mathrm{SO}(2)$ scalar.
• The rotation back is the exact inverse of the alignment.

The layer satisfies, for all $g \in \mathrm{SO}(3)$:

$$D^{(\ell_o)}(g) \bigl[ (h \otimes R(r))^{(\ell_o)} \bigr] = ( D^{(\ell_i)}(g)h \otimes R(g\cdot r) )^{(\ell_o)}.$$

Translation invariance follows from message dependence only on relative vectors $r_{ij}$ (Huang et al., 13 Feb 2026).

5. Computational and Memory Efficiency

The transformation from dense CG to axis-aligned sparsification reduces per-coupling complexity from $\mathcal{O}((2\ell_i+1)(2\ell_f+1)(2\ell_o+1))$ to $\mathcal{O}(2\ell_i+1)(2\ell_o+1)$ plus Wigner-D multiplications; for practical angular degrees ($\ell \leq 3$), this yields $\sim$6× per-edge speedup (Huang et al., 13 Feb 2026).

End-to-end inference throughput improvements, as empirically reported on large biomolecular benchmarks (1 K–50 K atoms), include:

• $\sim$5–6× faster inference compared to dense CG-based E2Former-V1.
• $\sim$4× steps-per-second vs. strongest baseline (UMA-S).
• Transient activation memory savings of $\sim$10–20× due to elimination of edge-centric tensor materialization [(Huang et al., 13 Feb 2026), Table 7].

A table of throughput exemplars:

Atom count	E2Former-V1 (steps/s)	E2Former-V2 + EAAS (steps/s)	Speedup
1,000	12	61	5×
10,000	1.2	6.1	5×

The fused Triton attention kernel leveraging EAAS demonstrates up to 20× higher TFLOPS than naively implemented PyTorch edge-centric routines, achieving end-to-end GPU feasibility for systems up to 100 K atoms (Huang et al., 23 Jan 2026).

6. Empirical Validation and Impact

Empirical ablation (V1→V2) demonstrates that nearly all throughput gains are attributable to the EAAS algebra in combination with hardware-aware attention kernels. Predictive accuracy remains competitive or improved on large biomolecular and materials datasets (SPICE, OMol25), with ab initio-level fidelity on out-of-distribution systems up to $\sim$1,500 atoms and robust MD observable prediction (Huang et al., 13 Feb 2026).

Beyond raw performance, the method enables the deployment of large-scale, equivariant foundation models in computational biology and chemistry, reconciling quantum-accurate modeling with practical hardware constraints.

7. Relation to Broader Equivariant Network Techniques

EAAS generalizes and systematizes the sparsification of tensor products in $\mathrm{SO}(3)$-equivariant geometric deep learning. It builds on the Wigner-6j convolutional calculus by employing a judicious $\mathrm{SO}(3)\to \mathrm{SO}(2)$ basis change, allowing three-way CG contractions to be replaced by parity-based re-indexing and local block-lin-maps (Huang et al., 23 Jan 2026). The concept synergizes with node-centric memory layouts and streaming attention kernels, therefore is immediately applicable wherever dense $\mathrm{SO}(3)$ CG products constitute a speed or memory barrier in equivariant architectures.

Further implications likely extend to other physical modeling contexts where equivariant message passing is foundational, subject to the specific structure of coupling maps and the underlying symmetry groups.