EquiformerV2: Equivariant Graph Transformers for Atomic Forces and Energies¶

EquiformerV2 is a Graph Transformer that predicts 3D atomic forces and total energy while mathematically guaranteeing rotational equivariance: rotate the input structure and the predicted force vectors rotate by the identical amount, with the predicted energy left unchanged. It avoids the brute-force cost of full 3D tensor algebra by re-using a clever alignment trick (eSCN) that turns expensive 3D rotation math into cheap 2D math.

0. Equivariance vs. Invariance — the Property Everything Below Protects¶

Formally, a function \(f\) is equivariant to a symmetry group \(G\) (here, 3D rotations \(SO(3)\)) if, for every rotation \(R\):

\[ f(R\cdot \mathbf{x}) = R \cdot f(\mathbf{x}) \]

i.e. rotate the input, and the output rotates the same way. This is what you need for forces — a vector quantity.

Invariance is the special case where the output doesn't rotate at all:

\[ E(R\cdot \mathbf{x}) = E(\mathbf{x}) \]

which is what you need for total energy — a scalar has to come out identical no matter which way the molecule was sitting when you looked at it. EquiformerV2's entire design is about building a network where every internal operation is provably equivariant, so the final energy comes out invariant and the final forces come out equivariant by construction, not by hoping the training data teaches it.

1. Representing 3D Shapes: Irreducible Representations (Irreps)¶

1.1 What the degree \(\ell\) actually encodes¶

Physical rotations of 3D functions decompose naturally into blocks called irreducible representations of \(SO(3)\), indexed by a non-negative integer degree \(\ell\). Each degree-\(\ell\) block has exactly \(2\ell+1\) components, indexed \(m = -\ell, \dots, \ell\), and transforms under a real spherical harmonic basis \(Y_\ell^m\):

\[ \mathbf{f} = \bigoplus_{\ell=0}^{L} \mathbf{f}^{(\ell)}, \qquad \dim \mathbf{f}^{(\ell)} = (2\ell+1)\times C_\ell \]

where \(C_\ell\) is the number of independent channels the network learns at that degree (analogous to "number of feature maps" in a CNN).

\(\ell\)	Components	Physical analogy
0	1	scalar (mass, charge, energy)
1	3	vector (position offset, velocity, force)
2	5	quadrupole-like directional tensor
\(\ell \geq 2\)	\(2\ell+1\)	increasingly fine angular detail (orbital-lobe-like shapes)

1.2 How rotation acts on each block¶

Under a rotation \(R\), each degree-\(\ell\) block transforms by its own \((2\ell+1)\times(2\ell+1)\) Wigner-D matrix \(D^{(\ell)}(R)\):

\[ \mathbf{f}^{(\ell)} \;\longrightarrow\; D^{(\ell)}(R)\,\mathbf{f}^{(\ell)} \]

\(D^{(\ell)}(R)\) is unitary, so it rotates the block without stretching or shrinking it — exactly what "rotating a physical quantity" should mean mathematically. \(D^{(0)}(R) = 1\) (a scalar doesn't change), \(D^{(1)}(R) = R\) itself (a vector rotates like the coordinates do).

Simple explanation: instead of one flat list of numbers per atom, EquiformerV2 keeps a scalar "shelf" (\(\ell=0\)), a vector "shelf" (\(\ell=1\), like little 3D arrows), a quadrupole "shelf" (\(\ell=2\)), and so on — each shelf has its own well-defined rule for how it spins when you rotate the molecule.

2. Replacing 3D Multiplication: eSCN Convolutions¶

2.1 Why the naive approach is slow¶

Mixing information between two irreps of degree \(\ell_1\) and \(\ell_2\) to produce a new degree-\(\ell_3\) feature uses the Clebsch–Gordan tensor product:

\[ \left(\mathbf{f}^{(\ell_1)} \otimes \mathbf{g}^{(\ell_2)}\right)^{(\ell_3)}_{m_3} = \sum_{m_1=-\ell_1}^{\ell_1}\sum_{m_2=-\ell_2}^{\ell_2} C^{\ell_3 m_3}_{\ell_1 m_1,\, \ell_2 m_2}\; f^{(\ell_1)}_{m_1}\, g^{(\ell_2)}_{m_2} \]

The \(C^{\ell_3 m_3}_{\ell_1 m_1, \ell_2 m_2}\) are fixed Clebsch–Gordan coefficients (standard angular-momentum coupling constants from quantum mechanics). Computing this for every pair of degrees up to a maximum \(L_{\max}\) costs \(O(L_{\max}^6)\) operations — it blows up fast as you add angular resolution.

2.2 The eSCN trick: align, then do 2D math¶

Compute the alignment rotation. From the bond vector between atoms \(i\) and \(j\), \(\hat{\mathbf{r}}_{ij} = (\mathbf{R}_j-\mathbf{R}_i)/\lVert \mathbf{R}_j-\mathbf{R}_i\rVert\), build the Wigner-D matrix \(\mathbf{D}(\mathbf{R}_{ij})\) that would rotate \(\hat{\mathbf{r}}_{ij}\) onto the \(z\)-axis.
Rotate both atoms' features into that frame:

\[ \mathbf{f}_{\text{aligned}} = \mathbf{D}(\mathbf{R}_{ij}) \cdot \mathbf{f} \]

Exploit the alignment. Spherical harmonics evaluated exactly on the \(z\)-axis vanish for every \(m \neq 0\). This means once the bond is aligned to \(z\), the tensor product's dependence on the bond direction collapses to a small set of \(m\)-indexed channels, and the coupling across those channels behaves like independent 2D (planar) rotations — the group reduces from \(SO(3)\) to \(SO(2)\) per channel. So a learnable per-\((\ell,m)\) linear layer (\(SO(2)\)-linear layer, weights \(W_{\ell,m}\)) can replace the full Clebsch-Gordan contraction:

\[ \mathbf{h}_{\ell,m} = W_{\ell,m}\, \mathbf{f}_{\text{aligned},\,\ell,m} \]

Rotate back with the inverse (= transpose, since \(\mathbf{D}\) is unitary) rotation: \(\mathbf{h} = \mathbf{D}(\mathbf{R}_{ij})^{\top}\, \mathbf{h}_{\text{aligned}}\).

This drops the cost from \(O(L_{\max}^6)\) to roughly \(O(L_{\max}^3)\), which is what lets EquiformerV2 push to higher \(L_{\max}\) (more angular detail) than older equivariant networks at the same compute budget.

Simple explanation: instead of multiplying two spinning globes at odd angles to each other, twist both globes until their poles line up on the same axis, do simple flat (2D) arithmetic, then twist everything back to where it started.

3. Stabilizing the Network: Separable Layer Normalization (SLN)¶

Ordinary layer norm forces a feature to have mean 0, variance 1 — but shifting or rescaling the direction of a physical vector by an arbitrary learned offset would break equivariance (a force vector's direction is physically meaningful; you can't zero-center it). SLN normalizes magnitude only, never direction, and treats scalars and higher-degree tensors separately:

Scalars (\(\ell=0\)) — ordinary layer norm is safe, since a scalar has no direction to preserve:

\[ \hat{\mathbf{f}}^{(0)} = \frac{\mathbf{f}^{(0)}-\mu}{\sqrt{\sigma_0^2+\epsilon}} \]

Vectors/tensors (\(\ell \geq 1\)) — compute a single collective variance from the norms of all higher-degree blocks (never touching their directional components), then divide every block by that shared scale:

\[ \sigma_{\text{vec}}^2 = \frac{1}{L}\sum_{\ell=1}^{L} \big\lVert \mathbf{f}^{(\ell)} \big\rVert^2, \qquad \hat{\mathbf{f}}^{(\ell)} = \frac{\mathbf{f}^{(\ell)}}{\sqrt{\sigma_{\text{vec}}^2+\epsilon}} \]

Because every degree-\(\ell\) block is divided by the same scalar \(\sqrt{\sigma_{\text{vec}}^2+\epsilon}\), and scaling a vector by a scalar doesn't change its direction, equivariance is preserved exactly.

Simple explanation: it reins in how big the directional "arrows" are allowed to get during training, without ever nudging which way they point.

4. Vector Non-Linearity: Separable \(S^2\) Activation¶

You can't apply ReLU coordinate-by-coordinate to a vector (zeroing just the negative \(x\)-component while leaving \(y,z\) alone physically rotates the vector's meaning, breaking equivariance). EquiformerV2's fix:

Forward transform to the sphere. Using the inverse spherical harmonic transform, evaluate the irrep features as a smooth scalar function sampled at grid points \((\theta,\phi)\) on a virtual unit sphere:

\[ v(\theta,\phi) = \sum_{\ell=0}^{L}\sum_{m=-\ell}^{\ell} f_{\ell m}\, Y_\ell^m(\theta,\phi) \]

Apply an ordinary pointwise activation (Swish/SiLU or ReLU) — safe now, because each grid value is just a plain number, not a directional coordinate:

\[ v'(\theta,\phi) = \text{Swish}\big(v(\theta,\phi)\big) \]

Project back via the forward spherical harmonic transform (numerically, a weighted sum over the sampled grid points):

\[ f'_{\ell m} = \int v'(\theta,\phi)\, Y_\ell^m(\theta,\phi)\, d\Omega \]

Since rotating the original irreps rotates the sampling grid on the sphere but does not change the values sampled at each physical point, and the activation acts identically at every point, the whole round-trip stays equivariant.

Simple explanation: unfold every directional feature onto the surface of a globe, squash the numbers at each point on the globe with an ordinary activation function, then fold the globe back into directional form.

5. The Transformer Part: Equivariant Graph Attention¶

EquiformerV2 is a Transformer, so messages between atoms are weighted by learned attention, not just summed uniformly. The trick for keeping this equivariant: attention weights are always plain scalars (computed only from invariant, \(\ell=0\) information plus the bond distance), and a scalar multiplying an equivariant tensor is still equivariant.

\[ \alpha_{ij} = \text{softmax}_j\Big(\mathbf{a}^{\top}\, \text{LeakyReLU}\big(W_Q \mathbf{f}_i^{(0)} \,\Vert\, W_K \mathbf{f}_j^{(0)} \,\Vert\, \text{RBF}(\lVert \mathbf{r}_{ij}\rVert)\big)\Big) \]

\[ \mathbf{m}_{ij} = \alpha_{ij}\cdot \big(W_V\, \mathbf{f}_j\big), \qquad \mathbf{f}_i \leftarrow \mathbf{f}_i + \sum_{j\in\mathcal{N}(i)} \mathbf{m}_{ij} \]

\(\text{RBF}(\lVert\mathbf{r}_{ij}\rVert)\): a Gaussian radial-basis expansion of the bond length, e.g. \(\sum_k \exp\!\big(-(r-\mu_k)^2/2\sigma^2\big)\), giving the attention score explicit access to how far apart the atoms are, independent of the angular (irrep) channels.
\(W_V \mathbf{f}_j\): the "value" is the full multi-degree feature of neighbor \(j\) (scalars through degree \(L\)), passed through the eSCN machinery described in Section 2.
Because \(\alpha_{ij}\) is a plain scalar and \(W_V\mathbf{f}_j\) is equivariant, their product \(\mathbf{m}_{ij}\) is still equivariant — attention adds learned importance weighting on top of the geometry without ever touching direction directly.

Simple explanation: each atom asks its neighbors "how relevant are you to me right now?" using only distance and rotation-blind (scalar) information, then scales each neighbor's full directional message by that relevance score before adding it in.

6. Summary of the Data Flow¶

[3D Graph of Atoms]
       │
       ▼
[Convert to Irreps Vectors f]  ──────────────►  Split into Scalars (ℓ=0), Vectors (ℓ=1), Higher tensors (ℓ≥2)
       │
       ▼
[Equivariant Graph Attention] ───────────────►  Score neighbors via scalar+distance info (softmax weights)
       │
       ▼
[eSCN Convolutions]        ──────────────────►  Twist bond to Z-axis → fast SO(2) math → twist back
       │
       ▼
[Separable Layer Normalization] ─────────────►  Normalize magnitudes only, directions untouched
       │
       ▼
[Separable S² Activation]   ─────────────────►  Map to sphere → apply Swish/ReLU → map back
       │
       ▼
[Final Predictions] ─────────────────────────►  Invariant total Energy (ℓ=0) + Equivariant Forces (ℓ=1)

7. Why This Matters for DFT-Adjacent Workflows¶

EquiformerV2-style models are trained on datasets like the Open Catalyst Project (OC20/OC22) — large sets of DFT-relaxed adsorbate–slab structures with energies and forces — precisely to act as a fast surrogate for expensive VASP/QE single points. In practice this shows up as:

High-throughput catalyst/adsorption screening: instead of running a full DFT relaxation for every candidate surface + adsorbate configuration in a Sabatier-volcano-style scan, the model gives near-DFT-accuracy energies/forces in milliseconds, letting you pre-filter candidates before committing VASP/QE time to the promising ones.
Faster geometry relaxation / MD: the model's predicted forces can drive structure optimization or short MD trajectories directly, with DFT reserved for final validation of the converged structures.
Force-constant-adjacent uses: while not a substitute for a proper FORCE_CONSTANTS/fc3.hdf5 calculation, finite-difference or autodiff Hessians from a trained equivariant model can flag which candidate structures are worth the cost of a full phono3py third-order force-constant run.