Equivariant Graph Neural Networks

Why Does This Exist?

When applying machine learning to physical systems—such as molecular drug discovery, protein folding, material crystal design, and n-body gravitational simulations—data is fundamentally embedded in 3D Euclidean space. Molecules do not possess an arbitrary coordinate frame: if an ethanol molecule rotates 45 degrees in space, its internal quantum potential energy (a scalar) remains strictly unchanged, while the atomic force vectors acting on its carbon atoms (3D vectors) must rotate by that exact same 45 degrees.

Standard Graph Neural Networks fail this physical requirement. If you feed raw Cartesian coordinates $(x, y, z)$ into standard MLPs or MPNNs, the network treats the coordinates as arbitrary numbers; rotating the molecule changes the predictions completely, forcing models to rely on data augmentation that requires millions of rotated training samples without ever guaranteeing physical consistency. Earlier equivariant models (such as Cormorant or TFN) solved this by projecting representations onto spherical harmonics and computing Clebsch-Gordan tensor products, but these required complex implementations and high computational overhead. Equivariant Graph Neural Networks (EGNN; Satorras et al., 2021) solved this by demonstrating that exact $E(n)$ equivariance can be achieved using simple scalar MLPs and radial vector differences.

Think of It Like This

How It Actually Works

Dual-Stream Invariant and Equivariant Message Passing

An Equivariant Graph Neural Network operates over an $n$-dimensional Euclidean group $E(n)$, which consists of all translations, rotations, and reflections in $\mathbb{R}^n$ (typically $n = 3$). A transformation $g \in E(n)$ acts on a coordinate $\mathbf{x} \in \mathbb{R}^3$ via:

$g \cdot \mathbf{x} = \mathbf{R} \mathbf{x} + \mathbf{t}$

where $\mathbf{R} \in \mathbb{R}^{3 \times 3}$ is an orthogonal matrix ($\mathbf{R}^\top \mathbf{R} = \mathbf{I}$, $\det(\mathbf{R}) = \pm 1$) and $\mathbf{t} \in \mathbb{R}^3$ is a translation vector.

Every node $i \in V$ carries two distinct streams of representations through layer $l$:

1. A scalar hidden state $\mathbf{h}_i^{(l)} \in \mathbb{R}^d$ (invariant).
2. A geometric coordinate vector $\mathbf{x}_i^{(l)} \in \mathbb{R}^3$ (equivariant).

An EGNN layer updates both representations across four steps:

1. Invariant Squared Distance Computation

For every pair of connected nodes $i$ and $j$, compute the Euclidean squared distance:

$d_{ij}^2 = \left\| \mathbf{x}_i^{(l)} - \mathbf{x}_j^{(l)} \right\|_2^2$

Because $\|\mathbf{R} \mathbf{x}_i + \mathbf{t} - (\mathbf{R} \mathbf{x}_j + \mathbf{t})\|^2 = \|\mathbf{R}(\mathbf{x}_i - \mathbf{x}_j)\|^2 = (\mathbf{x}_i - \mathbf{x}_j)^\top \mathbf{R}^\top \mathbf{R} (\mathbf{x}_i - \mathbf{x}_j) = d_{ij}^2$, the distance $d_{ij}^2$ is strictly invariant under $E(n)$.

2. Invariant Message Function $\phi_m$

Calculate edge messages using an MLP $\phi_m$ operating on invariant inputs:

$\mathbf{m}_{ij} = \phi_m\left(\mathbf{h}_i^{(l)}, \mathbf{h}_j^{(l)}, d_{ij}^2, a_{ij}\right)$

where $a_{ij}$ are optional edge attributes (like covalent bond types).

3. Equivariant Coordinate Update $\phi_x$

Update node positions by accumulating displacement vectors $(\mathbf{x}_i^{(l)} - \mathbf{x}_j^{(l)})$ scaled by a scalar MLP $\phi_x: \mathbb{R}^m \to \mathbb{R}$:

$\mathbf{x}_i^{(l+1)} = \mathbf{x}_i^{(l)} + \sum_{j \in \mathcal{N}(i)} \frac{1}{|\mathcal{N}(i)|} (\mathbf{x}_i^{(l)} - \mathbf{x}_j^{(l)}) \phi_x(\mathbf{m}_{ij})$

Under transformation $\mathbf{x}' = \mathbf{R}\mathbf{x} + \mathbf{t}$, the updated coordinate transforms as: $\mathbf{x}_i'^{(l+1)} = \mathbf{R}\mathbf{x}_i^{(l)} + \mathbf{t} + \sum_{j} \mathbf{R}(\mathbf{x}_i^{(l)} - \mathbf{x}_j^{(l)}) \phi_x(\mathbf{m}_{ij}) = \mathbf{R} \mathbf{x}_i^{(l+1)} + \mathbf{t}$ Equivariance holds without using spherical harmonics or group representation theory.

4. Invariant Node Feature Update $\phi_h$

Aggregate messages and update the scalar representations via MLP $\phi_h$:

$\mathbf{m}_i = \sum_{j \in \mathcal{N}(i)} \mathbf{m}_{ij}, \qquad \mathbf{h}_i^{(l+1)} = \phi_h\left(\mathbf{h}_i^{(l)}, \mathbf{m}_i\right)$

Worked Example

Consider a 2-node physical bond:

• Node 1 at $\mathbf{x}_1 = [0.0, 0.0, 0.0]^\top$, Node 2 at $\mathbf{x}_2 = [2.0, 0.0, 0.0]^\top$.
• Scalar features: $h_1 = 1.0, h_2 = 1.0$.
• Scalar coordinate function: $\phi_x(m) = 0.1$ (constant force coefficient).

1. Original Frame Update:

   • Relative displacement vector: $\mathbf{x}_1 - \mathbf{x}_2 = \begin{bmatrix} 0.0 - 2.0 \\ 0.0 - 0.0 \\ 0.0 - 0.0 \end{bmatrix} = \begin{bmatrix} -2.0 \\ 0.0 \\ 0.0 \end{bmatrix}$
   • Updated coordinate for node 1: $\mathbf{x}_1^{(1)} = \mathbf{x}_1^{(0)} + (\mathbf{x}_1 - \mathbf{x}_2) \phi_x = \begin{bmatrix} 0.0 \\ 0.0 \\ 0.0 \end{bmatrix} + \begin{bmatrix} -2.0 \\ 0.0 \\ 0.0 \end{bmatrix} \times 0.1 = \begin{bmatrix} -0.2 \\ 0.0 \\ 0.0 \end{bmatrix}$

2. Rotated Frame Verification: Now rotate the entire system $90^\circ$ counterclockwise around the $z$-axis ($\mathbf{R} = \begin{bmatrix} 0 & -1 & 0 \\ 1 & 0 & 0 \\ 0 & 0 & 1 \end{bmatrix}$): $\mathbf{x}_1^{\text{rot}} = \begin{bmatrix} 0.0 \\ 0.0 \\ 0.0 \end{bmatrix}, \quad \mathbf{x}_2^{\text{rot}} = \begin{bmatrix} 0.0 \\ 2.0 \\ 0.0 \end{bmatrix}$

   • Rotated displacement vector: $\mathbf{x}_1^{\text{rot}} - \mathbf{x}_2^{\text{rot}} = \begin{bmatrix} 0.0 \\ -2.0 \\ 0.0 \end{bmatrix}$
   • Updated coordinate under EGNN rule: $\mathbf{x}_1^{\text{rot},(1)} = \begin{bmatrix} 0.0 \\ 0.0 \\ 0.0 \end{bmatrix} + \begin{bmatrix} 0.0 \\ -2.0 \\ 0.0 \end{bmatrix} \times 0.1 = \begin{bmatrix} 0.0 \\ -0.2 \\ 0.0 \end{bmatrix}$
   • Checking $\mathbf{R} \mathbf{x}_1^{(1)}$: $\mathbf{R} \begin{bmatrix} -0.2 \\ 0.0 \\ 0.0 \end{bmatrix} = \begin{bmatrix} 0.0 \\ -0.2 \\ 0.0 \end{bmatrix} = \mathbf{x}_1^{\text{rot},(1)}$ The rotation commutes with the neural update.

Code

1from typing import Tuple2import numpy as np3
4
5class SimpleEGNNLayer:6    """Minimal Equivariant Graph Neural Network layer in NumPy."""7
8    def __init__(self, feat_dim: int) -> None:9        rng = np.random.default_rng(seed=42)10        # phi_m: predicts scalar message from (h_i, h_j, dist_sq)11        self.w_m: np.ndarray = rng.standard_normal((2 * feat_dim + 1, feat_dim)) * 0.112        # phi_x: predicts scalar multiplier for radial displacement vector13        self.w_x: np.ndarray = rng.standard_normal((feat_dim, 1)) * 0.114        # phi_h: updates scalar features15        self.w_h: np.ndarray = rng.standard_normal((2 * feat_dim, feat_dim)) * 0.116
17    def forward(18        self,19        h: np.ndarray,  # (N, feat_dim) scalar features20        x: np.ndarray,  # (N, 3) 3D Cartesian coordinates21        edges: list[Tuple[int, int]],22    ) -> Tuple[np.ndarray, np.ndarray]:23        num_nodes = h.shape[0]24        coord_updates = np.zeros_like(x)25        msg_sums = np.zeros_like(h)26
27        for i, j in edges:28            # 1. Invariant squared distance29            diff = x[i] - x[j]30            dist_sq = np.sum(diff ** 2, keepdims=True)31
32            # 2. Invariant message33            msg_input = np.concatenate([h[i], h[j], dist_sq], axis=0)34            m_ij = np.maximum(0.0, msg_input @ self.w_m)  # ReLU35            msg_sums[i] += m_ij36
37            # 3. Equivariant coordinate displacement38            phi_x_val = m_ij @ self.w_x  # Scalar weight39            coord_updates[i] += diff * phi_x_val40
41        # Update coordinates and scalar features42        new_x = x + coord_updates43        h_update_input = np.concatenate([h, msg_sums], axis=1)44        new_h = np.maximum(0.0, h_update_input @ self.w_h)45        return new_h, new_x46
47
48# Verification: Rotate coordinates and verify equivariance49layer = SimpleEGNNLayer(feat_dim=2)50h_in = np.array([[1.0, 0.5], [1.0, 0.5]])51x_orig = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]])52edges = [(0, 1), (1, 0)]53
54# 90-degree z-rotation matrix R55rot_matrix = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])56x_rot = x_orig @ rot_matrix.T57
58_, out_x_orig = layer.forward(h_in, x_orig, edges)59_, out_x_rot = layer.forward(h_in, x_rot, edges)60
61# Commutative check: R @ out_x_orig vs out_x_rot62expected_rot_out = out_x_orig @ rot_matrix.T63diff = np.max(np.abs(expected_rot_out - out_x_rot))64print("Max Equivariance Discrepancy:", np.round(diff, 8))65# -> Max Equivariance Discrepancy: 0.0 (exact numerical equivariance!)

Watch Out For

The Quick Version

• Equivariant GNNs guarantee that rotating or translating 3D input coordinates transforms predicted vector outputs by that exact same geometric rotation or translation.
• Scalar node features (like mass or charge) remain strictly invariant, while coordinate updates are scaled along relative displacement vectors.
• EGNN achieves exact $E(n)$ equivariance with standard MLPs and $\mathcal{O}(|E|)$ complexity, avoiding computationally heavy spherical harmonics.

What to Read Next

• Graph Transformers
• Spatiotemporal Graph Neural Networks
• Decision Trees