TraceGrad: a Framework Learning Expressive SO(3)-equivariant Non-linear Representations for Electronic-Structure Hamiltonian Prediction

Shi Yin, Xinyang Pan, Fengyan Wang, Lixin He

Abstract: We propose a framework to combine strong non-linear expressiveness with strict SO(3)-equivariance in prediction of the electronic-structure Hamiltonian, by exploring the mathematical relationships between SO(3)-invariant and SO(3)-equivariant quantities and their representations. The proposed framework, called **TraceGrad**, first constructs theoretical SO(3)-invariant **trace** quantities derived from the Hamiltonian targets, and use these invariant quantities as supervisory labels to guide the learning of high-quality SO(3)-invariant features. Given that SO(3)-invariance is preserved under non-linear operations, the learning of invariant features can extensively utilize non-linear mappings, thereby fully capturing the non-linear patterns inherent in physical systems. Building on this, we propose a **grad**ient-based mechanism to induce SO(3)-equivariant encodings of various degrees from the learned SO(3)-invariant features. This mechanism can incorporate powerful non-linear expressive capabilities into SO(3)-equivariant features with correspondence of physical dimensions to the regression targets, while theoretically preserving equivariant properties, establishing a strong foundation for predicting electronic-structure Hamiltonian. Experimental results on eight challenging benchmark databases demonstrate that our method achieves state-of-the-art performance in Hamiltonian prediction.

Lay Summary: Electronic structure calculations are essential for understanding electron behavior in condensed matter systems and for predicting properties such as conductivity, magnetism, optical response, and chemical reactivity. They are widely used in the fields of materials science, chemistry, and energy research. At their core lies the solution of the Hamiltonian matrices, which yield key physical quantities such as orbital energies, band structures, and electronic wavefunctions. We propose **TraceGrad**, a novel deep learning methodology that combines strict 3D rotational equivariance with strong non-linear expressiveness. It achieves state-of-the-art accuracy in predicting electronic-structure Hamiltonians and related properties, offering great potential to advance research on computational physics and chemistry.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: Atomic System Modeling; Electronic-structure Hamiltonian prediction; SO(3)-equivariant representation learning; Non-linear expressiveness

Submission Number: 14417