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

2024-05-09

AI Summary

This work addresses the high-accuracy prediction of electronic structure Hamiltonians. We propose a trace-guided SO(3)-equivariant learning framework: leveraging SO(3)-invariant traces derived from the Hamiltonian as supervision signals to learn physically invariant features, and designing a gradient-driven equivariant encoding mechanism that strictly enforces SO(3) equivariance while incorporating strong nonlinear representational capacity. Our approach introduces the novel “trace supervision + gradient enhancement” paradigm, unifying theoretical rigor with modeling flexibility. Evaluated on eight Hamiltonian benchmarks, it achieves state-of-the-art (SOTA) accuracy, significantly improving predictive fidelity for downstream physical properties—including energy gaps and dipole moments—and enabling acceleration of density functional theory (DFT) computations.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Learning to Search
• Game Theory and Economic Paradigms: Adversarial Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

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 gradient-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 consistency of physical dimensions to the regression targets, while theoretically preserving equivariant properties, establishing a strong foundation for predicting Hamiltonian. Our method achieves state-of-the-art performance in prediction accuracy across eight challenging benchmark databases on Hamiltonian prediction. Experimental results demonstrate that this approach not only improves the accuracy of Hamiltonian prediction but also significantly enhances the prediction for downstream physical quantities, and also markedly improves the acceleration performance for the traditional Density Functional Theory algorithms.

Problem

• Electronic Structure
• Hamiltonian Prediction
• Atomic Electron Behavior

Innovation

• TraceGrad
• SO(3)-Equivariant Representations
• Hamiltonian Prediction