Equivariant Atomic and Lattice Modeling Using Geometric Deep Learning for Crystal Structure Optimization

TL;DR

E3Relax tackles the bottleneck of crystal structure optimization by delivering an end-to-end, $SE(3)$-equivariant graph neural network that jointly updates atomic positions and lattice vectors. By promoting both atoms and lattice vectors to learnable nodes with dual scalar–vector features and employing layer-wise supervision, it mimics incremental convergence without iterative loops. Across four benchmark datasets (2D and 3D), E3Relax outperforms state-of-the-art iteration-free models in coordinate, shape, and volume metrics, and DFT validations show energetically favorable predictions that accelerate subsequent relaxations. The approach enables direct, physically consistent, and highly parallelizable structure relaxation, offering a practical starting point to speed up ab initio calculations like DFT.

Abstract

Structure optimization, which yields the relaxed structure (minimum-energy state), is essential for reliable materials property calculations, yet traditional ab initio approaches such as density-functional theory (DFT) are computationally intensive. Machine learning (ML) has emerged to alleviate this bottleneck but suffers from two major limitations: (i) existing models operate mainly on atoms, leaving lattice vectors implicit despite their critical role in structural optimization; and (ii) they often rely on multi-stage, non-end-to-end workflows that are prone to error accumulation. Here, we present E3Relax, an end-to-end equivariant graph neural network that maps an unrelaxed crystal directly to its relaxed structure. E3Relax promotes both atoms and lattice vectors to graph nodes endowed with dual scalar-vector features, enabling unified and symmetry-preserving modeling of atomic displacements and lattice deformations. A layer-wise supervision strategy forces every network depth to make a physically meaningful refinement, mimicking the incremental convergence of DFT while preserving a fully end-to-end pipeline. We evaluate E3Relax on four benchmark datasets and demonstrate that it achieves remarkable accuracy and efficiency. Through DFT validations, we show that the structures predicted by E3Relax are energetically favorable, making them suitable as high-quality initial configurations to accelerate DFT calculations.

Figures (11)

• Figure 1: Structural optimization using DFT and ML methods. (a) DFT-based structural optimization involves both inner and outer loops. (b) Iterative ML models surrogate replaces the DFT inner loop but still runs a geometry‑update outer loop until forces converge. (c) Existing iteration-free ML models initially estimate interatomic distances, subsequently identifying the structure that aligns with these predicted distances via an optimization process.
• Figure 2: An overview of E$^{3}$Relax. (a) The model jointly captures atomic displacements and lattice deformations during structural optimization by modeling both atoms and the lattice vectors. (b) Illustration of a multi-edge graph showing an atom connected to the same neighboring atoms in different translated unit cells. (c) Depiction of the three varieties of node features along with their interrelations.
• Figure 3: Two crystal structures from the X-Mn-O dataset optimized with E$^{3}$Relax: (a) $\rm Mn_4O_8Sr_4$ and (b) $\rm Ca_4Mn_4O_8$. $a$, $b$, and $c$ are lattice constants in angstroms ($\mathrm{\AA}$), and $\alpha$, $\beta$, and $\gamma$ are angles in degrees ($^\circ$).
• Figure 4: Comparison of coordinate MAE ($\rm \AA$) between E$^3$Relax and two ML-potential models using violin plots.
• Figure 5: Box plots of total energy distributions and residual ionic steps for E$^{3}$Relax predictions. (a) and (b) compare total energies of unrelaxed, DFT‐relaxed, and E$^{3}$Relax‐predicted structures on the X–Mn–O and C2DB datasets, respectively. (c) and (d) report the number of residual ionic steps required to converge DFT optimizations when initialized from unrelaxed and E$^{3}$Relax‐predicted structures for the layered vdW crystal dataset and the MoS2 defect dataset, respectively