SE(3) Equivariant Graph Neural Networks with Complete Local Frames

TL;DR

This paper introduces a computationally efficient SE(3) equivariant graph neural network framework using local complete frames, improving geometric quantity approximation for physics and biology applications.

Contribution

It proposes a novel framework with local complete frames for SE(3) equivariant GNNs, balancing expressiveness and efficiency by leveraging differential geometry concepts.

Findings

Achieves state-of-the-art or competitive results on physics modeling tasks.

Demonstrates computational efficiency through frame construction using cross products.

Effective in molecular conformation generation and Newtonian mechanics modeling.

Abstract

Group equivariance (e.g. SE(3) equivariance) is a critical physical symmetry in science, from classical and quantum physics to computational biology. It enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which has been shown to improve the generalization performance and data efficiency for downstream tasks. Constructing an equivariant neural network generally brings high computational costs to ensure expressiveness. Therefore, how to better trade-off the expressiveness and computational efficiency plays a core role in the design of the equivariant deep learning models. In this paper, we propose a framework to construct SE(3) equivariant graph neural networks that can approximate the geometric quantities efficiently. Inspired by differential geometry and physics, we introduce equivariant local complete frames to graph neural networks, such that tensor information at given orders can be projected onto the frames. The local frame is constructed to form an orthonormal basis that avoids direction degeneration and ensure completeness. Since the frames are built only by cross product operations, our method is computationally efficient. We evaluate our method on two tasks: Newton mechanics modeling and equilibrium molecule conformation generation. Extensive experimental results demonstrate that our model achieves the best or competitive performance in two types of datasets.