Equivariant Deep Learning in Direct-Product Space for Electronic-Structure Hamiltonian Prediction

Shi Yin; Zujian Dai; Xinyang Pan; Lixin He

cs.LGNov 5, 2026

Equivariant Deep Learning in Direct-Product Space for Electronic-Structure Hamiltonian Prediction

Authors: Shi Yin, Zujian Dai, Xinyang Pan, Lixin He

Organizations: Institute of Artificial Intelligence, Hefei Comprehensive National Science Center · Laboratory of Quantum Information, University of Science and Technology of China

Abstract

Most existing equivariant architectures for Hamiltonian prediction propagate features in direct-sum spaces of irreducible spherical tensors, whereas the target Hamiltonian is naturally an operator block in a direct-product representation over local orbital and spin bases. This representation mismatch fragments contiguous Hamiltonian blocks, obscures direct manipulation of orbital couplings, and organizes features into irregular angular-momentum components that require numerous transformations and are poorly aligned with efficient hardware execution and memory management. To address both the representation mismatch and computational bottleneck, we propose DP-Learner, an equivariant learning framework that performs message passing directly on Hamiltonian-like operator blocks in the direct-product space. Theoretically, we propose Lie-generator-based directional embedding and invariantly weighted left--right products and contractions that preserve equivariance to SO(3), SU(2), and spatial inversions, without explicit Clebsch--Gordan decomposition or inverse reconstruction. Methodologically, we instantiate this principle as a Transformer architecture, in which Hamiltonian-like feature matrices or tensors evolve through equivariant operators in the direct-product space, while non-linear invariant scalar routing controls structured gating over matrix units, directional propagation, attention, and environment-embedded context interactions. We further construct MP-HAM-110K, a large-scale Hamiltonian dataset comprising more than 110,000 material structures covering 68 elements with both Hamiltonian and wavefunction data. Experiments on MP-HAM-110K and Materials-HAM-SOC demonstrate that DP-Learner dramatically improves Hamiltonian prediction accuracy and inference efficiency over the state-of-the-art method.

Figures & tables

Figure 1: Existing approaches often encode features in the direct-sum representations of irreducible spherical tensors, creating a mismatch with the Hamiltonian target and requiring irregular, transformation-heavy computations that are difficult to execute efficiently. In contrast, DP-Learner operates directly in the Hamiltonian’s native direct-product representation, preserving the target tensor structure, orbital and spin couplings, as well as the symmetry constraints throughout the network for accurate and efficient equivariant Hamiltonian prediction.

Figure 2: Overview of the neural network architecture of DP-Learner.

Method	Materials-HAM-SOC				MP-HAM-110K
Method	$N_{\theta}$	$\overline{\mathrm{MAE}}_{\mathrm{G}}$	$\overline{t}_{\mathrm{NN}}$	$\overline{M}_{\mathrm{GPU}}$	$N_{\theta}$	$\overline{\mathrm{MAE}}_{\mathrm{G}}$	$\overline{t}_{\mathrm{NN}}$	$\overline{M}_{\mathrm{GPU}}$
	(M)	(meV)	(s)	(GiB)	(M)	(meV)	(s)	(GiB)
NextHAM	50.285	1.417	3.014	20.398	50.263	2.394	2.032	6.135
DP-Learner	26.445	0.841	0.269	0.989	26.094	1.595	0.203	0.531

Table 1: Accuracy and computational efficiency comparison on the Materials-HAM-SOC benchmark ( Yin et al., 2026 ) and the proposed MP-HAM-110K benchmark.

N_{\theta}

denotes the total number of network parameters. As detailed in Appendix K ,

\overline{\mathrm{MAE}}_{\mathrm{G}}

denotes the Gauge MAE metric;

\overline{t}_{\mathrm{NN}}

denotes the mean per-sample runtime of the neural inference stage, including input feeding, neural-network execution, result format conversion and saving;

\overline{M}_{\mathrm{GPU}}

denotes the mean per-sample peak GPU memory reserved by PyTorch. Lower values indicate better performance for all metrics.

Figure 3: Band structure comparison for a representative testing sample from Materials-HAM-SOC. Red solid lines show the DFT reference bands used as ground truths in our dataset; blue dashed lines show the bands obtained from the Hamiltonians predicted by NextHAM (left) and DP-Learner (right). Additional DP-Learner results are shown in Figures 9 and 11 . The displayed energies are referenced to the Fermi energy $E_{F}$ .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Representation	Method	Feature size	Layer-update cost	Input/output conversion
Fixed channel multiplicity $C$
Direct-sum	$SO(3)$ tensor product	$\Theta(CL^{2})$	$\mathcal{O}(CL^{6}+C^{2}L^{4})$	–
Direct-sum	$SO(2)$ convolution (eSCN)	$\Theta(CL^{2})$	$\mathcal{O}(C^{2}L^{3}+CL^{3})$	–
Retaining all coefficients of a complete Hamiltonian block: $P=D_{a}D_{b}=\Theta(D^{2})$
Direct-sum	$SO(3)$ tensor product	$P$	$\Theta(P^{2})$	Up to $\mathcal{O}(P^{3/2})$
Direct-sum	$SO(2)$ convolution (eSCN)	$P$	$\Theta(P^{7/4})$	Up to $\mathcal{O}(P^{3/2})$

Appendix

Table 2: Comparison of feature sizes and computational costs per directed edge. The upper part assumes

C

channels at each angular degree. The lower part compares methods that represent a complete Hamiltonian block with

P

coefficients, using the basis family in Eq. ( 119 ). The layer-update column reports the cost of one message-passing layer. Direct-sum methods additionally convert the input Hamiltonian into irreducible features and reconstruct the output Hamiltonian after message passing. The last column gives an upper bound on the cost of each conversion. DP-Learner operates directly on Hamiltonian blocks and requires neither conversion.

$Z$	Element	UPF	Pseudopotential valence	NAO basis	$r_{\mathrm{c}}$ (a.u.)
1	H	H.upf	$1s^{1}$	$2s\,1p$	7
2	He	He.upf	$1s^{2}$	$2s\,1p$	7
3	Li	Li.upf	$1s^{2}\,2s^{1}$	$4s\,1p$	7
4	Be	Be.upf	$1s^{2}\,2s^{2}$	$4s\,1p$	7
5	B	B.upf	$2s^{2}\,2p^{1}$	$2s\,2p\,1d$	7
6	C	C.upf	$2s^{2}\,2p^{2}$	$2s\,2p\,1d$	7

Appendix

Table 3: Element-resolved scalar-relativistic ONCV-PBE pseudopotentials and numerical atomic-orbital (NAO) basis sets used for the non-SOC dataset.

Panel of Figure 4 — Figure 4: Distributions of atomic counts and numbers of element types in the training, validation, and testing sets.

Panel of Figure 5 — Figure 5: Element-occurrence frequencies in the training, validation, and testing sets.

Figure 6: Sampled material structures from our MP-HAM-110K dataset.

Method	Gauge MAE (meV) $\downarrow$
	Materials-HAM-SOC	MP-HAM-110K
Only $H^{(0)}$	2.606	5.203
Without equivariant feature updates across layers	1.875	3.421
Without directional embedding	1.025	2.302
Without aggregation of equivariant context features	1.354	2.433
Without environment-dependent invariant coefficients	1.280	2.230

Appendix

Table 4: Results of ablation study on the Materials-HAM-SOC and MP-HAM-110K testing sets, using the SOC and non-SOC formulations, respectively. Gauge MAE is the gauge-corrected Hamiltonian mean absolute error in meV. The ablation term labeled as Only

H^{(0)}

reports the error obtained by directly using the zeroth-step Hamiltonian as the prediction, without any learned correction. Results for the full model and the five other ablation terms are averaged over three independent training runs. Lower values indicate better performance.

Figure 7: Band structures obtained using only $H^{(0)}$ and using DP-Learner for three representative structures from the MP-HAM-110K dataset. Each row shows the crystal structure (left), the bands obtained from $H^{(0)}$ (middle), and the DP-Learner bands (right). Red solid lines show the DFT reference bands used as ground truths in our dataset, while blue dashed lines show the results of the corresponding method. Energies are referenced to the Fermi energy $E_{F}$ .

Figure 8: Gauge MAE (meV) for element-defined subsets of the Materials-HAM-SOC testing set. Each subset contains all test structures with the indicated element. Colors are clipped to the displayed range; numerical labels give the actual errors.

Figure 9: More band structure results for representative Materials-HAM-SOC testing samples. Red solid lines show the DFT reference bands used as ground truths in our dataset; blue dashed lines show DP-Learner predictions. Energies are referenced to the Fermi energy $E_{F}$ .

Figure 10: Gauge MAE (meV) values for element-defined subsets of the MP-HAM-110K testing set. Each subset contains all test structures with the indicated element. Colors are clipped to the displayed range; numerical labels give the actual errors.

Figure 11: Additional band-structure results for representative test structures from MP-HAM-110K. Red solid lines show the DFT reference bands used as ground truths in our dataset; blue dashed lines show DP-Learner predictions. Energies are referenced to the Fermi energy $E_{F}$ .

Equivariant Neural Networks

Equivariant neural networks encode transformations such as rotations, translations, and permutations so model outputs change consistently with their inputs.

See all 165 papers

Jun 12, 2026

A Fixed-Point Neural Operator for Size- and Functional-Transferable Hamiltonian Prediction

Predicting the Kohn-Sham Hamiltonian with machine learning can accelerate density functional theory while retaining access to molecular orbitals, energy levels, and electronic-structure observables that energy-only surrogates cannot resolve. Yet element-wise agreement with the converged Hamiltonian, an implicit fixed point of the self-consistent field iteration, does not determine the occupied subspace that governs orbital energies and densities. Here we present HamEvo, a neural operator that learns the single-step self-consistent update and returns the converged Hamiltonian as its fixed point. HamEvo is pre-trained on intermediate self-consistent trajectories and calibrated at equilibrium with density-matrix supervision. Across benchmarks from MD17 to drug-like QMugs, HamEvo lowers Hamiltonian errors by 35-49% over direct-regression and deep-equilibrium baselines, and predicts QMugs HOMO and LUMO energies with mean absolute errors of 0.036 and 0.053 eV, near the 1 kcal/mol chemical-accuracy scale. Few-shot fine-tuning with only 20 reference conformations extends HamEvo to molecules of up to 122 atoms, well beyond the size range covered by pre-training. With thermal molecular-dynamics sampling, HamEvo captures temperature-dependent HOMO-LUMO gap renormalization beyond the harmonic approximation. Inference is up to 242 times faster than conventional DFT.

Yunhong Lou, Xihang Yue, Xinran Wei +2

1Zhejiang University, Hangzhou, China. · 2Zhongguancun Academy, Beijing, China. · 3Zhongguancun Institute of Artificial Intelligence, Beijing, China.

Quantum Chemistry Molecular Property Prediction Operator Learning Neural Operators Density Functional Theory

Ask PDF

Oct 26, 2026

Physical structure as an organizing principle for universal Hamiltonian learning

Density functional theory (DFT) gives the electronic structure of materials, but its cost limits large-supercell calculations and materials screening. Machine-learned Hamiltonians bypass the self-consistent cycle but are usually system-specific; universal ones could cover new compositions and disordered solids. Such models must generalize across chemistry and keep downstream properties physical, while staying inexpensive for large cells. Here we introduce UniTB, a universal E(3) equivariant Hamiltonian representation framework that integrates physical priors rooted in density-functional tight-binding theory with deep learning. UniTB's network, which corrects a prior Hamiltonian, generalizes across chemistry through a mixture of experts routed by that prior and a charge-equilibration head as in self-consistent-charge tight binding. A projection bounds band shifts from the prior, suppressing spurious states and keeping downstream properties physical. We provide the fastest SO(2) tensor-product GPU kernels for the network. The variant UniTB-Flash assembles its prior from two-center tables at

O(N)

cost for

N

atoms, against

O(N\log N)

on a grid. Trained on 29,242 crystals, UniTB reaches a 0.48 meV matrix-element error on nonmagnetic test crystals spanning 68 elements. Without amorphous training data, both models agree with DFT optics and conductivity of amorphous Ge-Sb-Te. UniTB-Flash ranks candidates with no significant difference from UniTB, whose property values are more accurate; together they screen 16,195 crystals for five application domains of industrial relevance.

Mingkang Liu, Zhanghao Zhouyin, Guanghui Wang +3

Suzhou Institute for Advanced Research, University of Science and Technology of China, Suzhou 215123, China · Department of Mechanical Engineering, National University of Singapore, 9 Engineering Drive 1, Singapore 117575, Singapore · Department of Physics, McGill University, Montreal, Quebec H3A 2T8, Canada +4

Computational Materials Science Physics-Informed ML Quantum Chemistry Materials Property Prediction Equivariant Representation Learning

Ask PDF

Jun 27, 2026

MALOQ: Massively Accelerated Learning of Operators for Quantum Transport

Machine-learned (ML) operator models can be trained to predict density functional theory (DFT) Hamiltonian/density matrices at significantly reduced computational cost, thus extending electronic-structure calculations to previously unfeasible scales. Here, we introduce MALOQ (Massively Accelerated Learning of Operators for Quantum Transport), an application built to train on and predict electronic-structure matrices for systems made of few to 100k atoms, described by large basis sets, and covering a wide range of atomic elements. Based on a state-of-the-art, SO(2)-equivariant backbone architecture, MALOQ provides (i) custom data-processing kernels to handle high-rank Hamiltonian matrix data and (ii) a scalable edge-wise distribution of atomic graph(s). Trained on the largest molecular Hamiltonian datasets available today, it reduces time-per-epoch by over 30% compared to a molecule-wise-distributed framework, and enables inference on material graphs of arbitrary size. We demonstrate scalable training and inference for 3,000-12,000 atoms on the Alps supercomputer, up to 192 GPUs and 256 GPUs, respectively.

Manasa Kaniselvan, Alexander Maeder, Denghui Lu +2

D-ITET, ETH Zurich, Zurich, Switzerland

Quantum Chemistry High-Performance Computing Distributed Training Computational Materials Science Computational Quantum Chemistry

Ask PDF

Abstract

Equivariant Neural Networks

Explore similar work

A Fixed-Point Neural Operator for Size- and Functional-Transferable Hamiltonian Prediction

Physical structure as an organizing principle for universal Hamiltonian learning

MALOQ: Massively Accelerated Learning of Operators for Quantum Transport