[Paper Review] E(n) Equivariant Graph Neural Networks

Victor García Satorras, Emiel Hoogeboom+1 more|arXiv (Cornell University)|Feb 19, 2021

Machine Learning in Materials Science37 references105 citations

TL;DR

EGNN introduces an E(n)-equivariant graph neural network that updates coordinates and features in a way that preserves translation, rotation, reflection, and permutation equivariance without relying on spherical harmonics, and it scales to higher dimensions.

ABSTRACT

This paper introduces a new model to learn graph neural networks equivariant to rotations, translations, reflections and permutations called E(n)-Equivariant Graph Neural Networks (EGNNs). In contrast with existing methods, our work does not require computationally expensive higher-order representations in intermediate layers while it still achieves competitive or better performance. In addition, whereas existing methods are limited to equivariance on 3 dimensional spaces, our model is easily scaled to higher-dimensional spaces. We demonstrate the effectiveness of our method on dynamical systems modelling, representation learning in graph autoencoders and predicting molecular properties.

Motivation & Objective

• Motivate and design a graph neural network that is equivariant to E(n) transformations (rotations, translations, reflections) and permutations.
• Avoid expensive higher-order representations while maintaining competitive performance.
• Demonstrate effectiveness across dynamical systems, graph autoencoders, and molecular property prediction in QM9.
• Show scalability to spaces with dimensionality n > 3 and compare against existing equivariant methods.

Proposed method

• Introduce an Equivariant Graph Convolutional Layer (EGCL) that processes node features h^l, coordinates x^l, and edge information E.
• Edge updates use m_{ij} = φ_e(h_i^l, h_j^l, ||x_i^l - x_j^l||^2, a_{ij}).
• Coordinate updates are x_i^{l+1} = x_i^l + C sum_{j≠i} (x_i^l - x_j^l) φ_x(m_{ij}) with C = 1/(M-1).
• Aggregate messages to form m_i = sum_j m_{ij} and update node features h_i^{l+1} = φ_h(h_i^l, m_i).
• Optionally extend with particle momentum: v_i^{l+1} = φ_v(h_i^l) v_i^{init} +… and x_i^{l+1} = x_i^l + v_i^{l+1}.
• Include a mechanism to infer edges when adjacency is not provided via φ_inf to produce soft edge values.

Experimental results

Research questions

• RQ1 Can EGNNs achieve E(n) (including SE(3) with reflections) equivariance without spherical harmonics or high-order representations?
• RQ2 Do EGNNs maintain or improve performance relative to non-equivariant GNNs and existing E(3)-equivariant models across dynamic, autoencoding, and molecular tasks?
• RQ3 Is the method scalable to higher dimensional spaces beyond 3D while remaining computationally efficient?
• RQ4 How does integrating coordinates and feature messages within the edge and node updates affect data efficiency and predictive accuracy?
• RQ5 Can edges be inferred on-the-fly without explicit graphs and still preserve equivariance?

Key findings

• EGNN achieves competitive or superior performance compared to both non-equivariant GNNs and some equivariant baselines across dynamical N-body tasks, graph autoencoding, and QM9 molecular prediction.
• On a 3D charged-particle N-body dataset, EGNN attains the lowest MSE (0.0071) among listed methods and is faster than several high-cost equivariant models.
• The model remains E(n) equivariant with respect to translations, rotations, and reflections, and scalable to higher dimensions without relying on spherical harmonics.
• In a graph autoencoder setup, EGNN substantially improves edge-reconstruction metrics over GNN variants, and handles symmetry-breaking noise in a translation/rotation-equivariant manner.
• On the QM9 benchmark, EGNN achieves competitive mean absolute errors for molecular properties, outperforming several baselines in key metrics (e.g., MAE ≈ 0.071 for α).
• EGNN supports an edge-inference mechanism to learn graph connectivity when adjacency is not provided, while preserving E(n) equivariance.