A Complete Guide to Spherical Equivariant Graph Transformers

Spherical equivariant graph neural networks (EGNNs) provide a principled framework for learning on three-dimensional molecular and biomolecular systems, where predictions must respect the rotational symmetries inherent in physics. These models extend traditional message-passing GNNs and Transformers by representing node and edge features as spherical tensors that transform under irreducible representations of the rotation group SO(3), ensuring that predictions change in physically meaningful ways under rotations of the input. This guide develops a complete, intuitive foundation for spherical equivariant modeling - from group representations and spherical harmonics, to tensor products, Clebsch-Gordan decomposition, and the construction of SO(3)-equivariant kernels. Building on this foundation, we construct the Tensor Field Network and SE(3)-Transformer architectures and explain how they perform equivariant message-passing and attention on geometric graphs. Through clear mathematical derivations and annotated code excerpts, this guide serves as a self-contained introduction for researchers and learners seeking to understand or implement spherical EGNNs for applications in chemistry, molecular property prediction, protein structure modeling, and generative modeling.

💡 Research Summary

This paper presents a comprehensive tutorial on spherical equivariant graph neural networks (EGNNs), focusing on their mathematical foundations, kernel construction, and practical implementation in the Tensor Field Network (TFN) and SE(3)-Transformer architectures. The motivation stems from the observation that conventional graph neural networks and Transformers, which operate on fixed reference frames, fail to respect the rotational symmetries inherent in three‑dimensional molecular and biomolecular data. Consequently, models trained on one orientation may not generalize to rotated versions, leading to inefficiencies and physically inconsistent predictions.

Section 2 introduces the necessary group‑theoretic background. The rotation group SO(3) is defined, and its representations ρ(g) are described as linear operators acting on feature spaces. Features are organized as spherical tensors of degree l, each transforming under an irreducible representation of SO(3). Spherical harmonics Yₗᵐ constitute an orthonormal basis for these representations, encoding angular dependence. The tensor product of two spherical tensors of degrees k and l is decomposed into a direct sum of irreducible components of degree J, where |k − l| ≤ J ≤ k + l, using Clebsch–Gordan coefficients. This decomposition is the cornerstone for building rotation‑equivariant kernels.

Section 3 derives the SO(3)‑equivariant kernel. The equivariance constraint ρₗ(g) K_{lk} = K_{lk} ρ_k(g) forces the kernel to be expressed as a linear combination of basis functions formed by the product of a learnable radial function f_J(r) (depending only on inter‑node distance) and spherical harmonics Y_J(Ω) (encoding relative orientation). The kernel thus takes the form K_{lk}(r,Ω) = ∑J C{lk}^J f_J(r) Y_J(Ω), where C_{lk}^J are learnable coefficients. The radial part can be parameterized by multilayer perceptrons, Gaussian radial basis functions, or other differentiable families, allowing the network to adapt distance‑dependent interactions during training.

Section 4 builds the Tensor Field Network using the equivariant kernel. Each node i holds a set of features h_i^{(l)} for several degrees l. For each neighbor j, a message is computed as m_{ij}^{(l)} = ∑k K{lk}(r_{ij},Ω_{ij}) · h_j^{(k)}. Messages are summed over the neighborhood, then passed through a linear self‑interaction (a learned linear map that preserves degree) and a channel‑mixing step that combines information across different degrees. The implementation leverages the Deep Graph Library (DGL) with user‑defined functions to efficiently perform message passing while explicitly tracking tensor degrees and channel multiplicities.

Section 5 extends TFN to an attention‑based architecture, the SE(3)-Transformer. Queries Q_i, keys K_j, and values V_j are constructed as spherical tensors of possibly different degrees (l_q, l_k, l_v). Attention scores are defined as equivariant inner products ⟨Q_i, K_j⟩, which are computed using the same Clebsch–Gordan machinery to guarantee rotation invariance of the scalar scores. After a softmax normalization, the scores weight the corresponding value tensors, and the weighted sum yields the updated node representation. Multi‑head attention is realized by stacking several independent sets of (Q,K,V) with distinct radial functions and degree assignments, dramatically increasing expressive power. Additional components such as norm‑based nonlinearity and edge‑feature incorporation are described to stabilize training and enrich relational modeling.

Section 6 demonstrates the full pipeline on the QM9 molecular property prediction benchmark. Molecules are converted into geometric graphs where atoms are nodes and inter‑atomic vectors become edge attributes. Node embeddings encode atomic number and electronic configuration; edge embeddings encode distances and directions via spherical tensors. The model stacks three to five equivariant layers (TFN or SE(3)-Transformer blocks) and is trained end‑to‑end without any data augmentation. Results show that the equivariant models achieve state‑of‑the‑art accuracy on multiple QM9 targets while converging faster than non‑equivariant baselines, confirming the practical benefits of built‑in rotational symmetry.

The conclusion emphasizes that spherical equivariant architectures embed physical symmetries directly into the learning process, yielding data‑efficient, physically consistent, and highly expressive models. By pairing rigorous mathematical derivations with annotated PyTorch/DGL code, the guide equips researchers and engineers with a ready‑to‑use toolkit for implementing, extending, and applying EGNNs to a broad range of 3D learning problems, from quantum chemistry to protein design and beyond.