Geometric and Physical Quantities Improve E(3) Equivariant Message Passing

TL;DR

The paper tackles the limitation of existing equivariant graph networks that rely on invariant node/edge features by introducing Steerable E(3) Equivariant Graph Neural Networks (SEGNNs). SEGNNs employ steerable MLPs, CG tensor products, and spherical-harmonic embeddings to allow covariant information (vectors, tensors) to flow through both messages and updates, achieving non-linear, equivariant convolutions. The authors provide a unifying framework for non-linear, steerable convolutions and demonstrate substantial improvements on N-body toy datasets, along with competitive results on QM9 and state-of-the-art performance on OC20 IS2RE, all supported by comprehensive ablation studies. This approach offers a principled way to inject geometric and physical quantities into graph-based models, with broad implications for efficient and accurate physics/chemistry simulations. Overall, SEGNNs advance the capacity of equivariant models to leverage rich geometric information, yielding practical benefits across molecular and material tasks.

Abstract

Including covariant information, such as position, force, velocity or spin is important in many tasks in computational physics and chemistry. We introduce Steerable E(3) Equivariant Graph Neural Networks (SEGNNs) that generalise equivariant graph networks, such that node and edge attributes are not restricted to invariant scalars, but can contain covariant information, such as vectors or tensors. This model, composed of steerable MLPs, is able to incorporate geometric and physical information in both the message and update functions. Through the definition of steerable node attributes, the MLPs provide a new class of activation functions for general use with steerable feature fields. We discuss ours and related work through the lens of equivariant non-linear convolutions, which further allows us to pin-point the successful components of SEGNNs: non-linear message aggregation improves upon classic linear (steerable) point convolutions; steerable messages improve upon recent equivariant graph networks that send invariant messages. We demonstrate the effectiveness of our method on several tasks in computational physics and chemistry and provide extensive ablation studies.

Figures (4)

• Figure 1: Commutation diagram for an equivariant operator $\phi$ applied to a 3D molecular graph with steerable node features (visualised as spherical functions); As the molecule rotates, so do the node features. The use of steerable vectors allows neural networks to exploit, embed, or learn geometric cues such as force and velocity vectors.
• Figure 2: Left: representation of an O($3$) steerable vector $\tilde{\mathbf{h}} \in V_L = V_0 \oplus V_1 \oplus V_2$, a spherical harmonic embedding of vector $\mathbf{x}$, e.g. relative orientation, velocity or force. Right: for each subspace the embedding with the basis functions $Y_m^{(l)}$ is shown (a). The transformation of $\tilde{\mathbf{h}}$ via $\mathbf{D}^{(l)}(g)$ acts on each subspace of $V_l$ separately (b).
• Figure C.1: Trajectory between $t=4$ and $t=5$ of 100 particles under gravitational interaction. Marker shows final position at end of simulation, with opacity decreasing over time.
• Figure C.2: Mean number of messages as a function of cutoff radius for the QM9 train partition. Standard deviation is indicated by the shaded region. The dashed black line indicates the transition from disconnected to connected graphs and may be thought of as a minimum cutoff radius.