Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

EGNN: E(n)-Equivariant Graph Neural Networks

4 minute read

Published:

TL;DR: EGNN (Satorras et al., 2021) tracks both feature embeddings h_i and coordinates r_i per node. Coordinates are updated via weighted sums of relative position vectors — equivariant by construction. Features are updated via distances — invariant by construction. No spherical harmonics, no expensive tensor products. Simple and effective.

The EGNN Approach

EGNN maintains two quantities per node:

  • h_i ∈ ℝ^d: invariant features
  • r_i ∈ ℝ^n: equivariant coordinates (n=3 for 3D)

These are updated jointly across layers, always maintaining E(n) equivariance.

The EGNN Update Rule

Edge messages (invariant quantities only):

m_{ij} = φ_e( h_i, h_j, ||r_i - r_j||², a_{ij} )
The message uses the squared distance r_i - r_j ² — an E(n)-invariant scalar. Optional edge attribute a_{ij} is also invariant.

Coordinate update (equivariant):

r_i ← r_i + C Σ_{j ≠ i} (r_i - r_j) φ_x( m_{ij} )

Relative positions (r_i - r_j) are translation-invariant and rotation-equivariant. A weighted sum of equivariant vectors is equivariant: rotating all coordinates by R gives R applied to the entire sum. So this update is automatically E(n)-equivariant.

Feature update (invariant):

m_i = AGG_{j ∈ N(i)}( m_{ij} ) h_i ← φ_h( h_i, m_i )

Features are updated only using invariant messages, so h_i remains invariant.

Why this works: The two key operations are: (1) distances — translation and rotation invariant; (2) relative position vectors (r_i - r_j) — translation invariant, rotation equivariant. By using only these in the update, EGNN never "accidentally" breaks symmetry. There is no way to write the update rule that violates E(n) equivariance, because no absolute coordinates appear.

Comparison to TFN / SE(3)-Transformers

PropertyEGNNTFN / SE(3)-Trans
Symmetry groupE(n)SE(3)
Geometric featuresDistances + relative positionsSpherical harmonics (irreps)
Computational costO(N E d)O(N E d L²) where L = max degree
Output typesScalars + vectors (l=0,1)Arbitrary tensors (l=0,1,…,L)
ComplexitySimple MLPsClebsch-Gordan tensor products
Practical speedFastSlow for high L

EGNN gives up expressive power beyond l=1 features (it has no l=2 or higher tensor outputs) in exchange for dramatically simpler and faster computation.

Applications

Molecular dynamics: predict energies and forces. Forces are the negative gradient of energy, so if energy is invariant, forces are automatically equivariant. EGNN can directly output equivariant force predictions.

N-body simulation: predict trajectories of charged particles. EGNN models the interaction forces and propagates positions forward.

Protein structure: predict residue positions given contact maps. Equivariance ensures predictions rotate consistently with the input.

Point cloud processing: EGNN applied to LiDAR point clouds maintains rotational equivariance for 3D object detection.

EGNN vs SchNet

SchNet (Schütt et al., 2017) is an earlier distance-based model: messages depend on r_i - r_j , so it is rotationally invariant. But SchNet does not update coordinates — it only updates scalar features. EGNN updates both features and coordinates, enabling it to predict vector quantities (forces, displacements) directly.

Training EGNN

For energy prediction:

  • Target: energy E (scalar) — use invariant readout: E = Σ_i φ_out(h_i)
  • Loss: MSE(E_pred, E_true)

For force prediction:

  • Forces: F_i = -∂E/∂r_i (autograd through the EGNN coordinate update)
  • Or: directly predict F_i using equivariant output head

Simultaneous energy and force training (combined loss) improves generalisation — forces carry information about the energy surface curvature.

Summary

ComponentEquivarianceHow
m_{ij} (messages)InvariantUses only distances
r_i updateEquivariantWeighted sum of (r_i - r_j)
h_i updateInvariantUses only invariant messages
Graph readoutInvariantSum of invariant node features

EGNN is the practical choice when E(n) equivariance is needed and the task requires only scalar (energy) or vector (force) outputs. For higher-order outputs (tensors, multipoles), TFN or MACE are needed.

References

You May Also Enjoy

Open Problems and Future Directions in Sheaf Neural Networks

5 minute read

Published:

Sheaf Neural Networks are a young field with many open research problems. This post surveys the most important unsolved questions: scalability to million-node graphs, stochastic sheaves for uncertainty quantification, sheaf autoencoders for generative models, connections to persistent homology, and unification with other topological deep learning frameworks.

Sheaf GNNs for Molecular Property Prediction

7 minute read

Published:

Molecular graphs are naturally heterophilic — atoms of different types bonded together — and have rich geometric structure (bond angles, torsion angles). Sheaf GNNs can encode this structure via restriction maps that represent the relational geometry between bonded atoms. This post explores how sheaf theory applies to molecular property prediction.

Sheaf GNNs for Heterophilic Node Classification: Benchmarks and Results

6 minute read

Published:

A systematic comparison of sheaf GNNs against baselines on standard heterophilic node classification benchmarks. Which datasets benefit most from sheaf structure? How does stalk dimension d affect performance? When do sheaves win — and when don’t they?

Sheaves Meet Attention: Transformer-Inspired Sheaf Models

6 minute read

Published:

Graph Transformers use global attention to capture long-range dependencies that local message passing misses. Sheaf Transformers combine sheaf-based local aggregation (with learnable relational geometry) with global attention (for long-range dependencies). This post surveys the design space of sheaf-augmented attention models.