Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI ๐Ÿ‘จ๐Ÿปโ€๐Ÿ’ป๐Ÿค–

EGNN: E(n)-Equivariant Graph Neural Networks

6 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.
EGNN equivariant message passing
E(n) Equivariant Graph Neural Network architecture (Satorras et al., 2021)

The EGNN Approach

Intuition First: EGNN is like a choreography that works in any room orientation. Each dancer (atom) tracks two things: their personal identity (invariant features h_i โ€” atomic number, charge) and their position on stage (equivariant coordinates r_i). When two dancers communicate, they whisper only about the distance between them (invariant), then each adjusts their stage position by moving slightly toward or away from the other. Because they only use relative positions โ€” never absolute coordinates โ€” the whole dance looks the same whether the room is rotated 30ยฐ or 300ยฐ.

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.

Worked Example: One EGNN Layer on 3 Atoms

Setup: 3 atoms โ€” H at rโ‚=(0,0,0), O at rโ‚‚=(1,0,0), H at rโ‚ƒ=(1,1,0). Features hโ‚=hโ‚ƒ=1 (hydrogen), hโ‚‚=2 (oxygen). ฯ†_e and ฯ†_x are simple identity functions for illustration.

Step 1 โ€” Edge messages (squared distances):

  • mโ‚โ‚‚ = ฯ†_e(hโ‚, hโ‚‚, โ€–rโ‚โˆ’rโ‚‚โ€–ยฒ) = ฯ†_e(1, 2, 1.0) โ†’ scalar sโ‚โ‚‚
  • mโ‚‚โ‚ƒ = ฯ†_e(hโ‚‚, hโ‚ƒ, โ€–rโ‚‚โˆ’rโ‚ƒโ€–ยฒ) = ฯ†_e(2, 1, 1.0) โ†’ scalar sโ‚‚โ‚ƒ
  • mโ‚โ‚ƒ = ฯ†_e(hโ‚, hโ‚ƒ, โ€–rโ‚โˆ’rโ‚ƒโ€–ยฒ) = ฯ†_e(1, 1, 2.0) โ†’ scalar sโ‚โ‚ƒ (longer bond)

Step 2 โ€” Coordinate update for O (atom 2), C=1:

  • rโ‚‚ โ† rโ‚‚ + (rโ‚‚โˆ’rโ‚)ยทฯ†_x(sโ‚โ‚‚) + (rโ‚‚โˆ’rโ‚ƒ)ยทฯ†_x(sโ‚‚โ‚ƒ)
  • (rโ‚‚โˆ’rโ‚) = (1,0,0), (rโ‚‚โˆ’rโ‚ƒ) = (0,โˆ’1,0)
  • If ฯ†_x(s) = 0.1: rโ‚‚ โ† (1,0,0) + 0.1ยท(1,0,0) + 0.1ยท(0,โˆ’1,0) = (1.1, โˆ’0.1, 0)

Equivariance check: rotate everything 90ยฐ around z-axis: rโ‚=(0,0,0)โ†’(0,0,0), rโ‚‚=(1,0,0)โ†’(0,1,0), rโ‚ƒ=(1,1,0)โ†’(โˆ’1,1,0).

  • (rโ‚‚โ€™โˆ’rโ‚โ€™) = (0,1,0), (rโ‚‚โ€™โˆ’rโ‚ƒโ€™) = (1,0,0)
  • rโ‚‚โ€™ โ† (0,1,0) + 0.1ยท(0,1,0) + 0.1ยท(1,0,0) = (0.1, 1.1, 0) โ€” exactly the 90ยฐ rotation of (1.1,โˆ’0.1,0) โœ“
Key Insight: EGNN achieves E(n)-equivariance without any spherical harmonics or Clebsch-Gordan coefficients โ€” just distances and relative position vectors. The trade-off is that it can only represent scalar (l=0) and vector (l=1) quantities. For polarisability tensors or octupole moments (l=2, l=3), you need TFN or MACE. For energy and forces in most molecular dynamics simulations, EGNN is sufficient and far faster.

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

Modern Activation Functions: GELU, SiLU, Mish, and Smooth Gating

11 minute read

Published:

Once ReLU became the default, researchers started asking a better question: can we keep the easy optimization while making the activation smoother, softer, and more expressive? This chapter covers the modern answers.

Activation Functions in Neural Networks: Why Non-Linearity Matters

11 minute read

Published:

Activation functions are the reason neural networks can model curved decision boundaries instead of collapsing into one giant linear map. This chapter builds the intuition first, then walks through the classical functions that shaped deep learning.

FoPE: Fourier Position Embedding for Length Generalization

6 minute read

Published:

FoPE rethinks long-context positional encoding from a frequency-domain perspective. Instead of only stretching RoPE heuristically, it explicitly improves attentionโ€™s periodic extension so Transformers generalize more gracefully to longer sequences.