Fast contracted Clebsch--Gordan tensor products for equivariant graph neural networks

TL;DR

This paper introduces an efficient $ ext{O}(L^3)$ algorithm for contracted Clebsch--Gordan tensor products in $ ext{O}(3)$-equivariant neural networks, improving computational scaling and enabling advanced angular momentum handling.

Contribution

The authors develop a novel structured grid approach that decouples radial and angular computations, recovering parity-odd channels and extending to parity-aware message passing in equivariant architectures.

Findings

Empirical wall-clock scaling as $L^2$ across practical $l_{max}$ range.

On-site contraction scales as $L^2$ pre-asymptotically and $L^3$ at large $l_{max}$.

Runtime is memory-bandwidth bound on $L^2$-sized grid tensors.

Abstract

We present an $\mathcal{O}(L^3)$ algorithm for evaluating contracted Clebsch--Gordan tensor products in $\mathrm{O}(3)$-equivariant machine learning potentials at fixed Canonical Polyadic (CP) rank. Mapping the angular integral to a structured Gauss--Legendre and Fourier tensor-product grid decouples the radial channel contractions from the angular transforms. The antisymmetric parity-odd Clebsch--Gordan channels, unreachable by the symmetric pointwise product on a scalar $S^2$ grid, are recovered through the surface-curl pairing $\hat r \cdot [\nabla_{S^2} A \times \nabla_{S^2} B]$, the spherical Poisson bracket, which supplies the $L=1$ angular momentum on the grid while preserving rotational equivariance. The construction extends to parity-aware equivariant message passing in atomic-cluster-expansion-style architectures and is verified by direct numerical quadrature. The full uncontracted Clebsch--Gordan tensor product remains subject to the $\mathcal{O}(L^4)$ output-size lower bound. A benchmark shows wall-clock scaling empirically as $L^2$ across the practical $l_{\max}$ range. For the on-site contraction this is pre-asymptotic, giving way to $L^3$ at large $l_{\max}$. For message passing it is structural and the runtime is memory-bandwidth bound on $L^2$-sized grid tensors.