An Efficient High-Degree, High-Order Equivariant Graph Neural Network for Direct Crystal Structure Optimization

Abstract

Crystal structure optimization is fundamental to materials modeling but remains computationally expensive when performed with density-functional theory (DFT). Machine-learning (ML) approaches offer substantial acceleration, yet existing methods face three key limitations: (i) most models operate solely on atoms and treat lattice vectors implicitly, despite their central role in structural optimization; (ii) they lack efficient mechanisms to capture high-degree angular information and higher-order geometric correlations simultaneously, which are essential for distinguishing subtle structural differences; and (iii) many pipelines are multi-stage or iterative rather than truly end-to-end, making them prone to error accumulation and limiting scalability. Here we present E$^{3}$Relax-H$^{2}$, an end-to-end high-degree, high-order equivariant graph neural network that maps an initial crystal directly to its relaxed structure. The key idea is to promote both atoms and lattice vectors to graph nodes, enabling a unified and symmetry-consistent representation of structural degrees of freedom. Building on this formulation, E$^{3}$Relax-H$^{2}$ introduces two message-passing mechanisms: (i) a high-degree, high-order message-passing module that efficiently captures high-degree angular representations and high-order many-body correlations; and (ii) a lattice-atom message-passing module that explicitly models the bidirectional coupling between lattice deformation and atomic displacement. In addition, we propose a differentiable periodicity-aware Cartesian displacement loss tailored for one-shot structure prediction under periodic boundary conditions.

Figures (9)

• Figure 1: Comparison of structural optimization using DFT and machine learning. (a) Conventional DFT-based structural relaxation, which involves two nested loops: an inner self-consistent electronic optimization and an outer geometry update loop. (b) Iterative ML-based relaxation methods, where a learned force field replaces the DFT inner loop, but an explicit geometry-update loop is still required. (c) Existing iteration-free ML approaches that bypass iterative relaxation by first predicting relaxed interatomic distances and subsequently reconstructing atomic coordinates. (d) Our iteration-free, end-to-end structure optimization framework, which directly predicts atomic and lattice displacements under SE(3)-equivariance, enabling a single forward pass from the initial structure to the relaxed configuration without intermediate optimization steps.
• Figure 2: Background concepts for GNN-based crystal structure modeling. (a) Periodic crystal graph illustrating an atom connected to symmetry-equivalent neighbors through different translated unit cells. (b) Example of a rotation- and translation-invariant GNN: the predicted energy remains unchanged under global rigid transformations of the input structure. (c) Example of a rotation-equivariant GNN: predicted atomic forces transform consistently with rotations of the input. (d) Real spherical harmonics $Y_{\ell m}(\theta,\phi)$ for degrees $\ell=0,1,2$, visualized on the unit sphere. Each row corresponds to a fixed degree $\ell$ and contains $(2\ell+1)$ basis functions with orders $m\in\{-\ell,\ldots,\ell\}$.
• Figure 3: Overview of the proposed high-degree, high-order message passing (H$^2$-MP). (a) For a central atom $i$ (node $0$), spherical harmonics $Y_{\ell m}(\hat{\mathbf{r}}_{ji})$ up to degree $\ell_{\max}$ are evaluated on the directions from neighboring atoms $j \in \mathcal{N}(i)$, providing high-degree angular descriptors for each edge. (b) A local reference direction $\hat{\mathbf{r}}_{\mathrm{ref},i}$ is constructed by aggregating all neighbor directions, and degree-wise bilinear invariants $z_{ji}^\ell$ are obtained by contracting spherical harmonics of $\hat{\mathbf{r}}_{ji}$ and $\hat{\mathbf{r}}_{\mathrm{ref},i}$. (c) The geometric invariants $\mathbf{z}_{ji}$ are mapped to invariant filters that gate scalar and vector edge messages without maintaining persistent high-degree equivariant features. (d) Edge-wise scalar and vector messages are aggregated over neighbors to update node features.
• Figure 4: Distributions of structural deviations between unrelaxed and DFT-relaxed structures across the six benchmark datasets. The left panel shows atomic coordinate MAE, and the right panel shows cell-shape deviation.
• Figure 5: Distribution of coordinate MAE ($\mathrm{\AA}$) between predicted and DFT-relaxed structures for Dummy and E$^{3}$Relax-H$^{2}$ on MP and X-Mn-O. (a) MP. (b) X-Mn-O. (c) Example prediction on MP ($\mathrm{Bi_8Cr_8Mg_8O_{40}}$). (d) Example prediction on X-Mn-O ($\mathrm{Ca_4Mn_4O_8}$). Lattice constants $a$, $b$, and $c$ are reported in angstroms (Å), and lattice angles $\alpha$, $\beta$, and $\gamma$ in degrees ($^\circ$).

Theorems & Definitions (2)

• proof
• proof