PaperlayerView exact revision on arXiv

Source-linked AI summary

E(n) Equivariant Graph Neural Networks

Victor Garcia Satorras, Emiel Hoogeboom, Max Welling

arXiv:2102.09844v3cs.LGstat.ML

TL;DR

Existing geometric graph methods face costly higher-order representations and limited dimensionality. EGNNs address this with an E(n)-equivariant, permutation-equivariant message-passing architecture using coordinate differences and scalar edge weights, achieving strong results across three task families. The paper reports competitive or superior performance while retaining computational simplicity and broader dimensional scalability.

  • Problem

    Existing equivariant graph methods can require computationally expensive higher-order representations, while geometric equivariance has often been developed for three-dimensional spaces.

  • Method

    EGNNs update node embeddings and coordinates through distance-based message passing and radial coordinate updates, preserving translation, rotation, reflection, and permutation equivariance.

  • Results

    EGNNs achieve the best or very competitive performance across dynamical-system modelling, graph-autoencoder representation learning, and QM9 molecular-property prediction.

  • Takeaways & Limitations

    The architecture provides a computationally efficient and easy-to-implement route to equivariant graph learning across a wide range of tasks and dimensions.

  • Takeaways & Limitations

    The method assumes velocities transform differently from translations: they are unaffected by translations but rotate and permute equivariantly.

Abstract

from arXiv · show

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.

1. Introduction

The paper introduces EGNNs to enforce translation, rotation, reflection, and permutation equivariance without expensive higher-order representations, while extending naturally beyond three dimensions. Across dynamical systems, graph autoencoders, and QM9 molecular-property prediction, the method achieves best or highly competitive performance.

  • E(3) symmetries arise in point clouds, molecular structures, and particle simulations, where predictions are often required to be equivariant or invariant.
  • Existing equivariant methods often use higher-order intermediate representations whose transformations require expensive coefficients or approximations.
  • EGNNs jointly provide translation, rotation, reflection, and permutation equivariance without spherical harmonics, while scaling to higher-dimensional spaces without significant computational growth.
  • The method is evaluated on dynamical-system modelling, graph-autoencoder representation learning, and QM9 molecular-property prediction.
  • EGNNs report the best or very competitive performance in all three experiments.

2. Background

The background defines equivariance as compatibility between transformations of inputs and outputs, then reviews translation, rotation/reflection, and permutation equivariance alongside standard graph neural-network message passing.

  • 2.1. Equivariance: A function is equivariant when transforming its input and then applying the function matches applying the function first and transforming its output equivalently.
  • 2.1. Equivariance: Translation equivariance requires φ(x+g) = φ(x) + g for translated point sets.
  • 2.1. Equivariance: Rotation and reflection equivariance require Qy = φ(Qx) for any orthogonal matrix Q.
  • 2.1. Equivariance: Permutation equivariance requires permuting the input nodes to produce the same permutation of the output.
  • 2.1. Equivariance: Velocities are translation-invariant but transform equivariantly under rotations and permutations.
  • 2.2. Graph Neural Networks: A graph-convolution layer computes edge messages, aggregates neighboring messages, and applies a node update using edge and node operations.

3. Equivariant Graph Neural Networks

EGNN layers augment node features with coordinates and update both through distance-based message passing and radial coordinate changes. This construction preserves E(n) equivariance and supports optional momentum, local neighborhoods, and inferred edges.

  • Equivariant Graph Neural Networks: An EGNN layer takes node embeddings, n-dimensional coordinates, and edge information, then outputs updated embeddings and coordinates.
  • Equivariant Graph Neural Networks: Edge operations receive relative squared distances and edge attributes, allowing geometric information to enter message computation.
  • Equivariant Graph Neural Networks: Coordinates are updated by a normalized weighted sum of relative differences, with scalar weights produced from edge embeddings.The normalization constant is C = 1/(M −1).
  • Equivariant Graph Neural Networks: Aggregation and node updates follow standard GNN operations, using messages from all other nodes or a restricted neighborhood.
  • Equivariance analysis: Squared distances remain E(n)-invariant, while weighted relative-coordinate differences transform as vectors, so stacked EGCLs preserve invariant features and equivariant coordinates.
  • Momentum variant: A momentum variant tracks velocity explicitly and reduces to the original coordinate update when the initial velocity is zero.
  • Inferring the edges: Fully connected message exchange may not scale to large point clouds, motivating neighborhood-limited aggregation; missing edges can also be estimated from invariant messages.

4. Related Work

Related methods achieve geometric equivariance through higher-order representations, spherical harmonics, or invariant molecular message passing. EGNN instead retains message-passing flexibility while using simpler, computationally cheaper operations.

  • Spherical-harmonic methods support transformations between higher-order representations but can require expensive computations and are limited to three-dimensional space.
  • Molecular message-passing methods differ in their symmetry coverage, with some permutation-equivariant, invariant, or SO(3)-equivariant approaches.
  • Table 1 compares Graph Neural Networks, Radial Field, Tensor Field Networks, SchNet, and EGNN using a common message-passing notation.
  • EGNN updates both node embeddings and coordinates, exchanging information between these variables in the edge operation.
  • EGNN retains GNN flexibility and E(n) equivariance without computing spherical harmonics.

5. Experiments

The experiments evaluate EGNNs on dynamical-system forecasting, graph autoencoding, and molecular-property prediction. Across these settings, EGNNs achieve strong performance while addressing symmetry-related limitations and supporting higher-dimensional embeddings.

  • 5.1. Modelling a dynamical system — N-body system: EGNN forecasts charged-particle trajectories under E(3)-equivariant dynamics using initial positions, velocities, and charges.The task estimates five particles’ positions after 1,000 timesteps from their initial states.
  • 5.1. Modelling a dynamical system — N-body system: 32% lower error than the second-best method was reported for EGNN, while the model remained efficient in forward-pass time.The comparison included equivariant and non-equivariant alternatives, including Radial Field, GNN, Tensor Field Networks, and the SE(3) Transformer.
  • 5.1. Modelling a dynamical system — N-body system: EGNN outperformed Radial Field and GNN across training sizes from 100 to 50,000 samples.The authors characterize this as combining the data efficiency of equivariant methods with the flexibility of GNNs in larger-data regimes.
  • 5.2. Graph Autoencoder: Featureless symmetric graphs can give identical node embeddings, preventing standard GNN autoencoders from reconstructing their adjacency matrices.A four-node cycle graph is presented as the clearest example of this symmetry problem.
  • 5.3. Molecular data — QM9: EGNN achieves very competitive QM9 molecular-property prediction while using only type-0 relative-distance representations instead of higher-order features.The paper reports that EGNN can outperform other equivariant networks using higher-order representations, angles, or spherical harmonics.

6. Conclusions

The paper presents an E(n)-equivariant graph architecture that is computationally efficient and improves over prior methods across diverse tasks. It targets applications involving molecules and other scientific data.

  • EGNNs provide a computationally efficient, easy-to-implement architecture for equivariant graph learning.
  • The method reportedly improves over the current state of the art across a wide range of tasks.
  • The authors identify drug discovery, protein folding, materials design, and 3D computer vision as potential application areas.

A. Equivariance Proof

The proof establishes that EGNN coordinate transformations preserve translation, rotation, and reflection equivariance while node features remain invariant under these transformations. The argument relies on invariant edge quantities and equivariant coordinate updates.

  • The proof targets translation equivariance for any g ∈ R^n and rotation/reflection equivariance for any orthogonal matrix Q ∈ R^n×n.
  • Assuming initial node features are invariant to transformations of coordinates, pairwise distances produce invariant edge operations.
  • The coordinate update is E(n)-equivariant because transformed inputs produce the same transformation at the output.
  • Rotating and translating x^l produces the same rotation and translation on x^(l+1).
  • The EGCL layer updates node features and coordinates while retaining E(n) equivariance.

B.1. Equivariance proof for velocity type inputs

The velocity-input variant extends EGNN coordinate updates with velocity features and remains E(n)-equivariant. The appendix also specifies shared architectures and training settings used across experiments and baselines.

  • Equivariance proof for velocity type inputs: The velocity-input formulation requires equivariance under translations by g and orthogonal transformations by Q.
  • Equivariance proof for velocity type inputs: Only the coordinate update is modified from the original EGNN, so the proof focuses on the rewritten velocity-inclusive equation.
  • Equivariance proof for velocity type inputs: The proof verifies equivariance separately for the first and second lines of the velocity-inclusive update.
  • Equivariance proof for velocity type inputs: The velocity-input EGCL maps transformed coordinates and velocities to equivalently transformed outputs.
  • Implementation details: The shared EGNN components use two-layer MLPs for edge, coordinate, and node functions, with the same edge and node functions used in GNN comparisons.
  • Implementation details: The dynamical-systems experiment uses a 3D five-particle N-body system with 3,000 training, 2,000 validation, and 2,000 testing trajectories.

Dataset

The graph experiments use Community Small and Erdős–Rényi datasets, including standard train/validation/test partitions and a small overfitting partition. Models share controlled architectural and optimization settings.

  • The experiments use Community Small and Erdős–Rényi generated graph datasets.
  • Community Small contains 5,000 training graphs, 500 validation graphs, and 500 testing graphs.
  • The Erdős–Rényi dataset uses edge probability p = 0.25 and graph sizes ranging from 7 to 16 nodes.
  • The graph autoencoding overfitting experiment uses 100 Erdős–Rényi graphs for each p value from 0.1 to 0.9.
  • Graph models use four layers, 64 hidden features, and Swish activations, with EGNN and GNN sharing edge and node functions.
  • Training uses learning rate 10^-4, batch size 1, Adam, weight decay 10^-16, and early stopping for the 5,000-sample datasets.

C.3. Implementation details for QM9

The QM9 EGNN implementation uses seven layers and property-specific training with Adam, normalization, and learning-rate schedules.

  • QM9 properties are normalized by subtracting the mean and dividing by the Mean Absolute Deviation.
  • The EGNN has 7 layers, with inferred edges produced by a linear layer followed by sigmoid.
  • Each property is trained individually for 1,000 epochs using Adam with batch size 96 and weight decay 10^-16.
  • The learning rate uses cosine decay from 5 · 10^-4, except for Homo, Lumo, and Gap, which start at 10^-3.

D.1. Graph Autoencoder

The graph autoencoder evaluates reconstruction across embedding sizes and datasets, finding that EGNN performance improves with larger embeddings and surpasses the comparison models.

  • Table 4 reports wrong-edge percentages and F1 scores across embedding sizes for GNN, Noise-GNN, and EGNN.
  • For small embedding sizes (n = 4), all methods perform poorly, but EGNN significantly outperforms the others as embedding size grows.
  • The analysis covers the Community Small and Erdos&Renyi datasets.

E. Sometimes invariant features are all you need.

The paper argues that pairwise distance norms are invariant under E(n) transformations and uniquely determine point-set geometry up to an E(n) transformation.

  • The section defines pairwise ℓ2 distance norms as a unique identifier of point geometry up to E(n) transformations.
  • An E(n) transformation has the form x 7→ Qx + t, with orthogonal Q and translation t.
  • Pairwise ℓ2 distances remain unchanged because translations cancel and orthogonal transformations satisfy Q^TQ = I.
  • When two point collections have identical pairwise distance norms, the proof constructs a linear map matching corresponding points and shows it is orthogonal.
  • The constructed map is orthogonal on the span of the points and can be extended to the whole space using an orthogonal complement.
Loading 2102.09844v3…