PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Saro Passaro, C. Lawrence Zitnick

arXiv:2302.03655v2cs.LGphysics.chem-phphysics.comp-ph

TL;DR

SO(3)-equivariant convolutions become computationally expensive as representation degree increases. The paper aligns node irreps with edge directions to reduce these operations to equivalent SO(2) convolutions, yielding eSCN, which achieves state-of-the-art performance on atomic modeling tasks.

  • Problem

    Higher-degree SO(3)-equivariant convolutions are costly, limiting efficient use of equivariant GNNs for 3D data.

  • Method

    The method aligns irreps’ primary axes with edge directions, sparsifies tensor products, and implements the resulting operations as equivalent SO(2) convolutions in eSCN.

  • Results

    eSCN achieves state-of-the-art performance on atomic modeling tasks, including force prediction, on large-scale OC-20 and OC-22 datasets.

  • Takeaways & Limitations

    The approach reduces tensor-product complexity from O(L^6) to O(L^3), enabling equivariant models with substantially higher representation degrees.

  • Takeaways & Limitations

    The point-wise spherical nonlinearity is quasi-equivariant under discrete transformation and can be computationally expensive when numerical-precision equivariance is required.

Abstract

from arXiv · show

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.

1. Introduction

SO(3)-equivariant GNNs encode 3D rotational structure but become costly at higher representation degrees. The paper reduces this cost through an SO(2)-equivalent convolution approach and applies it in eSCN for atomic modeling.

  • SO(3)-equivariant GNNs use irreducible representations and equivariant tensor-product convolutions to process 3D geometric information.
  • O(L^6) complexity for tensor products up to degree L limits practical use of degrees above 2 or 3.
  • Aligning irreps’ primary axes with edge directions sparsifies tensor products, reducing complexity from O(L^6) to O(L^3) and removing Clebsch–Gordan coefficient computation.
  • With aligned irreps, message passing has one remaining rotational degree of freedom and can be expressed through mathematically equivalent SO(2) convolutions.
  • The Equivariant Spherical Channel Network applies the efficient convolutions to atomic-energy and force prediction on the large-scale OC-20 and OC-22 datasets.

2. Related work

The paper situates equivariant GNNs within group-theoretic symmetry modeling and describes spherical-harmonic irreps, tensor-product message passing, and the computational bottleneck addressed by the proposed implementation.

  • Equivariance requires outputs to transform under group actions in the same way as inputs, while invariance leaves outputs unchanged.
  • For atomic systems, SO(3) describes 3D rotations: energies are invariant, whereas forces transform equivariantly.
  • Spherical-harmonic coefficients form SO(3) irreps indexed by degree l and order m, with degrees ranging from 0 to L.
  • Spherical harmonics are steerable because 3D rotations act on their coefficients through Wigner D-matrices.
  • Invariant GNNs use distances or angular features, while newer equivariant models use geometric tensors and constrained SO(3)-equivariant operations.
  • In e3nn, message passing combines an input irrep with an edge-direction filter through generalized tensor products whose coefficients depend on distance and atomic numbers.
  • Retaining all nonzero tensor products through degree L becomes computationally unfeasible as L grows, while typical e3nn models therefore use low degrees.
  • Figure 1 compares mathematically equivalent e3nn and eSCN networks using GPU memory and epoch time as functions of L with 64 channels and no nonlinearity.

3. Efficient equivariant convolution

Aligning each edge direction to a fixed axis makes the spherical-harmonic tensor product sparse, enabling an equivalent but substantially cheaper equivariant convolution. The resulting formulation avoids precomputing Clebsch–Gordan coefficients and reduces complexity from O(L^6) to O(L^3).

  • The new equivariant convolution is mathematically equivalent to the original formulation through a linear bijection between its scalar coefficients and the original filter coefficients.This preserves the information represented by the convolution while changing its parameterization.
  • The aligned formulation retains only coefficients with m_f=0, reducing the tensor product to selected sparse entries.The corresponding Clebsch–Gordan matrices follow a particular sparsity pattern.
  • The efficient formulation no longer precomputes Clebsch–Gordan coefficients or sums over m_f, m_i, and l_f.These changes make the convolution substantially more efficient.
  • O(L^6) to O(L^3): aligning the edge direction to a fixed axis reduces the computational complexity of the tensor product.Choosing R so that R·r̂_st=(0,1,0) makes the spherical harmonics sparse and reduces the tensor-product computation.
  • At high m close to L, GPU utilization can decrease because fewer coefficients remain, although the authors report restricting |m| to 2 or 3 in practice.

4. SO(2) formulation

Aligning each node embedding’s primary axis with the message direction leaves only roll rotation, reducing the SO(3) operation to equivalent SO(2) convolutions. This viewpoint also explains why the generalized tensor product becomes sparse and computationally cheaper.

  • SO(3) to SO(2) reduction: Aligning the primary axis with the message direction leaves only one rotational degree of freedom: roll about that direction.An SO(2)-equivariant message-passing function is therefore also SO(3)-equivariant.
  • Circular-harmonic representation: After alignment, spherical-harmonic coefficients can be expressed through SO(2) circular harmonics whose rotation is parameterized by the roll angle.The colatitude is fixed by alignment, while the longitude represents rotation about the aligned axis.
  • SO(2) convolution: SO(2)-equivariant convolution over the roll angle can be implemented as a point-wise product in the Fourier-domain circular-harmonic representation.Circular harmonics form a Fourier series, enabling this spectral simplification.
  • Equivalence: The generalized tensor product can therefore be viewed as a set of generalized SO(2) convolutions after appropriately rotating the spherical-harmonic coefficients.The proposed SO(2) expression coincides with the tensor product under the stated coefficient identification.

5. Architecture

The architecture updates atom embeddings through edge-conditioned message passing built from alignment, SO(2) convolutions, spherical nonlinearities, aggregation, and output heads for energies and forces.

  • Atom embeddings: Each atom is represented by irreps indexed by degree, order, and channel, initialized from atomic number and updated through repeated message passing.The embedding component x_i^(l)_mc is the m-th component of channel c at degree l.
  • Message passing: Each message-passing layer combines an edge embedding block, generalized SO(2) convolution blocks, and a point-wise spherical nonlinearity.The edge block uses invariant edge information; the SO(2) blocks process aligned embeddings before the nonlinear operation.
  • Message construction: The point-wise spherical nonlinearity applies SiLU across the sphere, using a discrete approximation of the integral before rotating the message back to the original frame.The operation is equivariant because the same function is applied at every point on the sphere.
  • Message aggregation: Messages from neighboring atoms are summed, combined with the target embedding through another spherical nonlinear function, and added residually to produce the updated embedding.The aggregation function is a three-layer network with SiLU activations.
  • Outputs: Final atom embeddings produce per-atom forces and total energy through spherical integrations, approximated by sampling 128 weighted spherical-Fibonacci points.The energy and force heads use separate three-layer neural networks, with the force integral weighted by the unit direction vector.

6. Experiments

The experiments evaluate eSCN on large OC20 and OC22 benchmarks, showing strong performance on direction-sensitive force tasks and substantial computational savings. Ablations indicate that degree benefits force prediction, while depth benefits energy prediction.

  • Benchmark evaluation: 130M, 38M, and 8M examples are used for OC20 All, OC20 All+MD, and OC22 evaluations, respectively.The experiments cover S2EF, IS2RS, and IS2RE tasks, with relaxation-based evaluation for the latter two.
  • Main results: 9% and 21% force MAE improvements are reported for OC20 All+MD and OC22 ID, respectively.These gains occur on tasks requiring high-fidelity directional information.
  • Main results: Over 15% higher AFbT accuracy is achieved on OC20 IS2RS.AFbT measures how often relaxed structures are found according to DFT using model force estimates.
  • Main results: Energy prediction is on par with SCN and GemNet-OC on OC20 MD+All and OC22.The paper links energy prediction to longer-range reasoning between nodes.
  • Efficiency: An order-of-magnitude computational reduction is observed as the maximum representation degree L increases.The comparison uses mathematically equivalent implementations and measures GPU memory and time per epoch.
  • Ablations: Force MAE is more sensitive to degree L, whereas energy MAE is more sensitive to layer count K.Increasing M improves results up to M = 2, while removing the point-wise activation lowers accuracy.
  • Efficiency: Runtime is similar to SCN across the reported model variations.The authors note that smaller models could run faster with larger batch sizes.

7. Discussion

The discussion connects eSCN to the broader importance of atomic modeling while identifying computational and equivariance-related boundaries. The paper concludes that the SO(3)-to-SO(2) relationship enables efficient, high-performing force prediction, but spherical nonlinearities remain costly for numerical precision.

  • Broader context: Atomic modeling is presented as relevant to climate-change challenges, while the paper cautions that chemistry advances can have harmful effects.The Haber–Bosch process is given as an example of both societal benefit and environmental cost.
  • Limitations: Point-wise spherical nonlinearities are quasi-equivariant because irreps are discretely transformed to the sphere.Achieving equivariance to numerical precision can make this operation computationally expensive.
  • Conclusion: The conclusion links SO(3) and SO(2) convolutions to dramatic computational-efficiency improvements.The approach is demonstrated with eSCN on atomic modeling tasks including force prediction.

A.1. Proof of Proposition 3.1

The appendix establishes the representation-theoretic basis for the sparse tensor product: aligning the edge direction with the y-axis retains only selected Clebsch–Gordan components. It then relates the resulting decomposition to SO(2) irreducible representations through an invertible projection.

  • Sparse tensor products: Aligning the edge direction to the y-axis allows the tensor product to retain only coefficients with m_f = 0.This produces sparse Clebsch–Gordan matrices and reduces the computational structure of the product.
  • Clebsch–Gordan relations: SU(2) Clebsch–Gordan coefficients are used as an equivalent basis because SO(3) coefficients are related to them by an invertible linear transformation.The appendix therefore transfers coefficient properties from SU(2) to SO(3).
  • Selection rules: The relevant SU(2) coefficients are nonzero only when m_o = m_i + m_f, and with m_f = 0 this becomes m_o = m_i.These selection rules explain the sparsity after alignment.
  • SO(2) decomposition: The projection groups SO(3) components into one-dimensional and paired subspaces that transform as SO(2) irreducible representations.The paired components correspond to circular harmonics of degree k, while the projection remains invertible.
  • Selection rules: The SO(3) coefficients satisfy parity- and index-based zero conditions that further constrain the aligned tensor product.The appendix states conditions involving l_i + l_f + l_o and the magnitudes of magnetic indices.
  • Representation framework: The representation framework defines SO(3) irreps in dimensions 2l + 1 and SO(2) irreps as one-dimensional for k = 0 or two-dimensional for k > 0.These classifications provide the spaces used by the projection argument.

C. Computational comparison with e3nn tensor product

The computational comparison analyzes mathematically equivalent GNNs implemented with e3nn and with the proposed SO(2)-based tensor product. The theoretical scaling and GPU measurements support lower cost as the representation degree grows.

  • Complexity analysis: The naive tensor-product implementation requires O(C · L^6) operations for its 3D matrix multiplications.It performs O(C · L^3) tensor-product operations before linear self-interactions.
  • Complexity analysis: The SO(2) formulation requires O(C^2 · L^3) operations for the tensor-product computation.It computes O(L) two-dimensional matrix multiplications and adds O(C · L^3) operations for edge-direction alignment.
  • Memory scaling: Training memory scales as O(C^2 · L^3) for the SO(2) formulation versus O(C · L^6 + C^2 · L^3) naively.The comparison concerns GPU memory allocated during training.
  • Empirical validation: The empirical comparison trains mathematically equivalent GNNs for 1 ≤ L ≤ 5 while measuring time per epoch and GPU memory.The hidden-channel count is fixed at 64 and nonlinearities are removed.

D. Analysis on the quasi-equivariance of the spherical activation function

The spherical activation is only approximately equivariant because discrete sampling can introduce aliasing, but the measured error remains small under the paper’s settings and can be reduced with finer grids.

  • Discrete sampling causes small per-layer equivariance loss because nonlinear activations introduce higher frequencies that can alias.The spherical nonlinearity samples uniformly spaced angular grids, whereas a linear activation would preserve perfect equivariance at those resolutions.
  • 1.5% relative error occurs with grid size 14, SiLU, and AMP in the paper’s settings.The reported error is measured between two rotated messages.
  • ReLU may require higher resolution than SiLU when numerically perfect equivariance is desired.The paper attributes this requirement to the higher frequencies introduced by the activation function.

E. Sample efficiency of eSCN

eSCN preserves the sample efficiency of SCN on OC-20, based on force-error curves across training epochs.

  • eSCN preserves SCN’s sample efficiency on the OC-20 2M dataset.The comparison uses Force MAE plotted against training epochs and includes GemNet-OC as a state-of-the-art reference.
Loading 2302.03655v2…