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

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.