Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shi Yin, Xinyang Pan, Fengyan Wang, Lixin He

TL;DR

TraceGrad introduces a principled framework to reconcile strict SO(3)-equivariance with powerful non-linear expressiveness for electronic-structure Hamiltonian prediction. It uses an invariant trace $T = \mathrm{tr}(\mathbf{H} \mathbf{H}^\dagger)$ as a supervision signal to learn high-quality SO(3)-invariant features and then propagates non-linearity to SO(3)-equivariant representations via a gradient-based mechanism $\mathbf{v} = \frac{\partial z}{\partial \mathbf{f}}$, preserving dimensional consistency. The method combines an invariant decoding path for $T$ with an equivariant path for $\mathbf{H}$, trained jointly with a balanced loss, and extends to multi-degree representations through learnable CG decompositions. Empirically, TraceGrad achieves state-of-the-art accuracy across eight benchmarks, improves downstream quantities like occupied energies and wavefunctions, and accelerates convergence of traditional DFT methods. The work demonstrates that strict symmetry enforcement and non-linear expressiveness can be cohesively integrated for scalable, accurate electronic-structure predictions.

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.

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

TL;DR

TraceGrad introduces a principled framework to reconcile strict SO(3)-equivariance with powerful non-linear expressiveness for electronic-structure Hamiltonian prediction. It uses an invariant trace as a supervision signal to learn high-quality SO(3)-invariant features and then propagates non-linearity to SO(3)-equivariant representations via a gradient-based mechanism , preserving dimensional consistency. The method combines an invariant decoding path for with an equivariant path for , trained jointly with a balanced loss, and extends to multi-degree representations through learnable CG decompositions. Empirically, TraceGrad achieves state-of-the-art accuracy across eight benchmarks, improves downstream quantities like occupied energies and wavefunctions, and accelerates convergence of traditional DFT methods. The work demonstrates that strict symmetry enforcement and non-linear expressiveness can be cohesively integrated for scalable, accurate electronic-structure predictions.

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.
Paper Structure (26 sections, 2 theorems, 33 equations, 7 figures, 10 tables)

This paper contains 26 sections, 2 theorems, 33 equations, 7 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminary
  4. Theory
  5. Method
  6. Encoding and Decoding Framework
  7. Training
  8. Experiments
  9. Experimental Conditions
  10. Results and Analysis
  11. Summary of Appendices
  12. Conclusion
  13. Definition of Concepts
  14. Application Task Description: Electronic-structure Hamiltonian Calculation
  15. Formulation of Hamiltonian Prediction as an SO(3)-Equivariant Non-linear Representation Learning Problem
  16. ...and 11 more sections

Key Result

Theorem 4.1

The quantity $\mathbf{T}=tr(\mathbf{H} \cdot \mathbf{H}^\dagger)$ is SO(3)-invariant, where $\mathbf{H}$ is the simplified representation (without superscripts) of the basic block $\mathbf{H}^{l^i_p \otimes l^j_q}$ of Hamiltonian matrix defined in Appendix task, and $\dagger$ denotes the conjugate t

Figures (7)

  • Figure 1: The proposed framework for learning SO(3)-equivariant representations with strong non-linear expressiveness to regress the SO(3)-equivariant electronic-structure Hamiltonian.
  • Figure 2: Illustration of atomic pairwise Hamiltonian matrices partitioned from the complete Hamiltonian matrix of the whole system. Each Hamiltonian matrix contains multiple basic blocks, denoted as $\mathbf{H}^{l^i_p \otimes l^j_q}$.
  • Figure 3: Visualization of testing samples exhibiting non-rigid deformations due to thermal motions
  • Figure 4: Visualization of testing samples with interlayer twists.
  • Figure 5: Visualization of $MAE^H_{block}$ on each basic block of the Hamiltonian matrices for the Monolayer Graphene ($MG$) and Monolayer MoS2 ($MM$) databases.
  • ...and 2 more figures

Theorems & Definitions (17)

  • Theorem 4.1
  • Theorem 4.2
  • Remark 4.3
  • Definition 1
  • Definition 2
  • Definition 3
  • Definition 4
  • Definition 5
  • Definition 6
  • Definition 7
  • ...and 7 more