Geometric and Physical Quantities Improve E(3) Equivariant Message Passing
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Molecular property prediction | QM9 (test) | mu23 | 174 | |
| Initial Structure to Relaxed Energy (IS2RE) | OC20 (Open Catalyst 2020) IS2RE (test) | Energy MAE (Avg)0.6101 | 30 | |
| Initial Structure to Relaxed Energy | OC20 IS2RE (val) | Energy MAE (ID)0.531 | 24 | |
| Initial Structure to Relaxed Energy | OC20 IS2RE (ID) | Energy MAE (eV)0.5327 | 13 | |
| Initial State to Relaxed Energy (IS2RE) | OC20 IS2RE OOD Adsorbate (test) | MAE (eV)0.6921 | 13 | |
| Initial Structure to Relaxed Energy | OC20 IS2RE (Out-of-Domain Both) | Energy MAE (eV)0.679 | 13 | |
| Initial Structure to Relaxed Energy | OC20 IS2RE Direct (test) | Energy MAE (ID)533 | 11 | |
| Future state prediction | M-complex Single System (5, 10) | MSE (x10^-2)15.01 | 10 | |
| Future state prediction | M-complex Single System (3, 3) | Prediction Error (MSE)0.1404 | 10 | |
| Particle position estimation | N-body system | MSE0.0043 | 9 |