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

TL;DR

This paper introduces E3Relax-H2, an end-to-end equivariant graph neural network that efficiently models high-degree angular and higher-order correlations for direct crystal structure optimization, improving accuracy and scalability.

Contribution

It presents a novel unified graph representation including atoms and lattice vectors, with advanced message-passing modules for high-order correlations, enabling direct and accurate crystal relaxation.

Findings

Achieves state-of-the-art accuracy in crystal structure optimization

Reduces computational cost compared to DFT methods

Demonstrates robustness across diverse crystal systems

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.