Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs

TL;DR

This work tackles the computational bottleneck of $SO(3)$-equivariant GNNs for 3D data by reducing to $SO(2)$-equivariant convolutions through edge-aligned irreps, yielding $O(L^3)$ complexity instead of $O(L^6)$. It introduces the Equivariant Spherical Channel Network (eSCN), which leverages edge rotations and circular harmonics to perform efficient, Clebsch-Gordan-free convolutions while preserving equivariance. Empirical evaluation on OC20 and OC22 demonstrates state-of-the-art performance in force prediction and relaxed-structure accuracy, with substantial speedups that enable higher angular resolutions (larger $L$) in practice. Overall, the approach provides a scalable pathway for accurate 3D equivariant GNNs with broad implications for chemistry and materials science predictions, by bridging $SO(3)$ and $SO(2)$ formalisms in a computationally efficient framework.

Abstract

Graph neural networks that model 3D data, such as point clouds or atoms, are typically desired to be $SO(3)$ equivariant, i.e., equivariant to 3D rotations. Unfortunately equivariant convolutions, which are a fundamental operation for equivariant networks, increase significantly in computational complexity as higher-order tensors are used. In this paper, we address this issue by reducing the $SO(3)$ convolutions or tensor products to mathematically equivalent convolutions in $SO(2)$ . This is accomplished by aligning the node embeddings' primary axis with the edge vectors, which sparsifies the tensor product and reduces the computational complexity from $O(L^6)$ to $O(L^3)$, where $L$ is the degree of the representation. We demonstrate the potential implications of this improvement by proposing the Equivariant Spherical Channel Network (eSCN), a graph neural network utilizing our novel approach to equivariant convolutions, which achieves state-of-the-art results on the large-scale OC-20 and OC-22 datasets.

Figures (10)

• Figure 1: We implement two mathematically equivalent equivariant GNNs. e3nn uses the e3nn PyTorch library, whereas eSCN uses our novel approach. Left: GPU memory allocated (%) at training time with resepct to the maximum degree of the spherical harmonics ($L$). Right: time (h) per epoch by $L$. We fix the number of channels to $C=64$ and remove any non-linearity.
• Figure 2: Visual representation of the tensor product $\mathbf{x}^{(2)} \otimes^{1}_{2, 1} \mathbf{Y}^{(l)}(\hat{\mathbf{r}})$ between $\mathbf{x} \in \mathbb{R}^5$ and $\mathbf{Y}^{(1)}(\hat{\mathbf{r}}) \in \mathbb{R}^3$. The computational cost of the tensor product can be lowered if the direction $\hat{\mathbf{r}}$ is aligned with the y-axis. The tensor product is then reduced to the multiplication of the green shaded entries. The orange shaded entries are the non-zero entries in the Clebsch-Gordan matrix.
• Figure 3: Visual representation of the Clebsch-Gordan matrices $\mathbf{C}^{(1, m_o)}_{(2, m_i), (1, m_f)} \in \mathbb{R}^{5 \times 3 \times 3}$ and $\mathbf{C}^{(2, m_o)}_{(2, m_i), (1, m_f)} \in \mathbb{R}^{5 \times 3 \times 5}$. We illustrate the sparsity of the Clebsch-Gordan matrices as stated in Proposition \ref{['prop:C-G']}.
• Figure 4: Illustration of how the spherical harmonics (top) are reduced to a set of circular harmonics (bottom) when $\theta$ is held constant.
• Figure 5: Block diagram of the message passing architecture containing the edge embedding block (yellow), the $SO(2)$ blocks (green) and the point-wise non-linearity (red).

Theorems & Definitions (3)

• Proposition 3.1
• Definition 2.1
• Theorem 2.2