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

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.

Efficient Equivariant Graph Neural Networks through SO(3) to SO(2) Reduction

This paper presents a novel methodology to enhance the computational efficiency of graph neural networks (GNNs) that are equivariant to the group of 3D rotations, SO(3), which is a critical feature for modeling 3D data such as point clouds or atomic systems. Traditional approaches to equivariant convolutions in these networks encounter significant computational complexity as they involve operations with higher-order tensors. The authors propose a solution by reducing these SO(3) convolutions to two-dimensional SO(2) convolutions while maintaining the mathematical equivalence.

Core Contribution

The pivotal contribution of the paper is the reduction of the computational complexity of equivariant convolutions from $O(L^6)$ to $O(L^3)$, where $L$ is the degree of representation. This reduction is achieved by aligning the primary axis of node embeddings with edge vectors, leading to a sparser tensor product configuration. The novel method was empirically validated through the development of the Equivariant Spherical Channel Network (eSCN), which implemented the efficient SO(2) convolution-based message passing.

Numerical Results and Claims

The eSCN model is evaluated on the OC-20 and OC-22 datasets, large-scale benchmarks designed to model atomic energies and forces. The results demonstrate that eSCN achieves state-of-the-art performance, with notable improvements on tasks requiring high directional fidelity such as force predictions. For instance, eSCN provides advancements of up to 21% in force MAE over other leading models, highlighting its efficacy in capturing critical structural details.

Implications and Future Directions

The theoretical and practical implications of this research extend to any domain that involves the processing of 3D geometric data with inherent symmetries. By leveraging the reduction to SO(2) convolutions, the proposed method significantly lowers the computational barriers associated with using higher-degree representations in equivariant GNNs. This could spur further innovation in designing deep learning architectures that are efficient yet maintain geometric fidelity, facilitating new applications in material science, chemistry, and beyond.

Looking ahead, the paper opens new avenues for exploring equivariant networks with even higher degrees and more complex symmetries. It also sets a precedent for the potential reductions in other symmetries beyond SO(3) that could benefit from analogous transformations. The intersection of computational efficiency and geometric deep learning will likely remain a fertile ground for advancing AI capabilities in understanding and manipulating complex 3D systems.