[Paper Review] Geometric and Physical Quantities Improve E(3) Equivariant Message Passing

J. Brandstetter, Rob Hesselink+3 more|arXiv (Cornell University)|Oct 6, 2021

Advanced Graph Neural NetworksComputer Science41 citations

TL;DR

The paper introduces SEGNNs, a general E(3) equivariant GNN framework that uses steerable vector spaces and Clebsch-Gordan products to inject geometric and physical quantities into both message and update steps, achieving state-of-the-art or competitive results on several physics/chemistry benchmarks.

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.

Motivation & Objective

• Generalize equivariant GNNs to allow node and edge attributes beyond scalars to include covariant information (vectors, tensors).
• Introduce steerable MLPs and steerable activation functions to inject geometric/physical quantities into message passing.
• Provide a unifying view of equivariant GNNs as non-linear convolutions and identify key components that boost performance.
• Demonstrate effectiveness and offer extensive ablations on N-body, QM9, and OC20 datasets.

Proposed method

• Use steerable vector spaces V_L = V0 ⊕ V1 ⊕ ... ⊕ VL and steerable MLPs whose linear maps are conditioned by geometric information via Clebsch-Gordan products.
• Represent node/edge attributes as steerable vectors (including spherical harmonic embeddings of relative positions).
• Define non-linear steerable message passing with messages that depend on steerable edge attributes and squared distances, and updates steered by node attributes.
• Incorporate spherical harmonics embeddings to convert vectors into type-l steerable vectors for CG tensor products.
• Propose steerable activation functions via steerable node attributes that enable non-linear, equivariant processing of covariant information.
• Relate SEGNNs to non-linear convolutions and compare against linear equivariant/convolutional baselines.

Experimental results

Research questions

• RQ1 Can node and edge attributes in GNNs be extended beyond scalars to covariant quantities (vectors, tensors) while preserving E(3) equivariance?
• RQ2 Do non-linear, steerable message passing layers improve performance over linear, invariant/steerable message passing in physical/chemistry tasks?
• RQ3 How do geometric and physical quantities (like velocity, force) affect SEGNN performance on molecular and N-body problems?
• RQ4 What constitutes the essential components of effective E(3) equivariant GNNs in this framework?

Key findings

• SEGNNs with l_f = 1 and l_a = 1 achieve state-of-the-art or competitive results across benchmarks, outperforming several baselines.
• Inclusion of both geometric and physical quantities (G+P) significantly boosts performance over geometry-only (G) or invariant methods.
• Non-linear steerable convolutions provide advantages over linear steerable message passing and over invariant messages.
• Using spherical harmonic embeddings for relative positions enables effective steerable CG-based message transformations.
• SEGNNs deliver strong results on N-body tasks, QM9, and OC20 IS2RE while maintaining rotational/equivariant consistency.