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

arXiv: https://arxiv.org/abs/2405.05722 · HTML: https://mlsanity.com/paper/2405.05722/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 33 references that are in the map.
- First posted on arXiv 9 May 2024.

## Builds on (most cited first, 14 of 14)

- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data](https://mlsanity.com/paper/1807.02547/) (2018)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [e3nn: Euclidean Neural Networks](https://mlsanity.com/paper/2207.09453/) (2022)
- [Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems](https://mlsanity.com/paper/2307.08423/) (2023)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [Unifying machine learning and quantum chemistry -- a deep neural network for molecular wavefunctions](https://mlsanity.com/paper/1906.10033/) (2019)
- [Spherical Channels for Modeling Atomic Interactions](https://mlsanity.com/paper/2206.14331/) (2022)
- [Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian](https://mlsanity.com/paper/2306.04922/) (2023)
- [QH9: A Quantum Hamiltonian Prediction Benchmark for QM9 Molecules](https://mlsanity.com/paper/2306.09549/) (2023)
- [Self-Consistency Training for Density-Functional-Theory Hamiltonian Prediction](https://mlsanity.com/paper/2403.09560/) (2024)
- [Towards Harmonization of SO(3)-Equivariance and Expressiveness: a Hybrid Deep Learning Framework for Electronic-Structure Hamiltonian Prediction](https://mlsanity.com/paper/2401.00744/) (2024)

## Built on by (most cited first, 4 of 4)

- [Learning local equivariant representations for quantum operators](https://mlsanity.com/paper/2407.06053/) (2024)
- [Efficient Prediction of SO(3)-Equivariant Hamiltonian Matrices via SO(2) Local Frames](https://mlsanity.com/paper/2506.09398/) (2025)
- [Advancing Universal Deep Learning for Electronic-Structure Hamiltonian Prediction of Materials](https://mlsanity.com/paper/2509.19877/) (2025)
- [NextCrystal: a Symmetry-Driven Generative Framework for Crystal Structure Prediction](https://mlsanity.com/paper/2602.17176/) (2026)

## Abstract (arXiv)

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.