Physical structure as an organizing principle for universal Hamiltonian learning
Authors: Mingkang Liu, Zhanghao Zhouyin, Guanghui Wang, Daye Zheng, Lei Shen, Qiangqiang Gu
Organizations: 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 · School of Artificial Intelligence and Data Science, University of Science and Technology of China, Hefei 230026, China · AI for Science Institute, Beijing 100080, China · National University of Singapore (Chongqing) Research Institute, Chongqing 401120, China · National University of Singapore Guangzhou Research Translation and Innovation Institute, Guangzhou, Guangdong 510700, China
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(NlogN) 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.
Figures & tables
Figure 1: Architecture of UniTB. a , UniTB uses the prior P=H(0) and UniTB-Flash uses P=HGF . Each model trains an onsite and a hopping network of the same architecture (stacked frames). The prior enters along the orange paths: it initializes the features ( ), its rotation invariants d(Pij) feed the router ( ), and it is the baseline of the residual sum in Eq. ( 1 ) ( ). Three interaction layers update edge and node features with PDQ-MoE SO(2) maps (green) and element-dependent tensor products (TP). Only UniTB applies the spectral projection, shown removing a spurious in-gap level (inset). b , In each PDQ-MoE SO(2) layer, the router selects two of four expert cores De . The cores act between the shared bases QT and P inside the bond-frame rotations, and Ws is always active. c , In the charge-equilibration head, readouts of hi give electronegativities χi and hardnesses Ji . Charge equilibration at zero total charge (Eq. ( 7 )) gives latent charges qi , and their site potentials vi scale the overlap Sij . Values are schematic.
Figure 2: Assembly and cost of the two prior Hamiltonians and their effect on training. a , The grid-free prior HGF of UniTB-Flash and the zeroth-step Hamiltonian H(0) of UniTB are assembled in five steps (Methods). They differ in the two-center step and in step 4, where HGF sums finite-range atomic terms and H(0) evaluates its potential on a grid that fills the cell. b , Assembly time of each prior for a 1,024-atom PbTe supercell. c , Inference time for PbTe supercells on one H200 GPU. Full inference runs from the structure to H and S and includes the prior assembly and the network. The labels ∝N and ∝NlogN give the derived cost forms, and dashed lines mark the prior times of b . d , Onsite and hopping losses, 41 (MAE + RMSE), of three dense networks that are trained identically on MPHam-12K and differ only in the prior. Each trajectory runs from the first checkpoint (open circles) to the last (filled circles).
Matrix MAE (meV)
Band fw10 (meV)
Model
Onsite
Hopping
Combined
Mean
Median
Published (Yin et al.) 12,000 structures, 100 epochs
NextHAM
2.834
Ours, MPHam-12K 11,953 structures, 100 epochs
LEM
7.925
2.098
2.160
88.11
62.08
SLEM
11.319
1.278
1.386
80.62
48.53
Table 1: Matrix-element and band errors. Matrix errors are evaluated on the 3,000 test structures and band errors on 200 of them drawn at random. Onsite, Hopping and Combined are mean absolute errors of the real spin-up–spin-up block. The band error fw10 is computed after the spectral projection with b=10 eV, or for UniTB-Flash after canonical orthogonalization (Methods). Indented rows add components cumulatively; UniTB also includes site channels and prior perturbations. UniTB-Flash uses the grid-free prior. LEM and SLEM 13 are trained with the H(0) prior on the same data and budget as the MPHam-12K models. NextHAM reports no aggregate band errors 17 ; our labels are recomputed with its DFT settings (Supplementary Note Supplementary Note 4: Datasets and labels ).
Figure 3: Matrix-element errors across the periodic table. Values are for UniTB (MPHam-29K) on the 2,431 nonmagnetic test structures. Each tile averages the matrix-element MAE over the structures that contain the element, so a structure counts once for each of its elements. Dashes mark elements not in the test set. Logarithmic color scale. Values beyond the color range take the end color.
Figure 4: Site-potential errors and the charge-equilibration head. a , A dense model without the charge head is trained for this diagnostic on its own split of 3,063 training and 256 test structures. Bars give the RMS of the component of its matrix error that has the overlap-weighted site-potential form of SCC-DFTB. b , Change in combined matrix MAE when the charge head is added to PDQ-MoE or to the dense model, for the 3,000 test structures binned by atoms per cell. Changes are paired over the same structures, with 95% intervals. Percentages give the share of structures on which the head lowers the error of PDQ-MoE. The PDQ-MoE-versus-dense comparison is made without the head (Supplementary Note Supplementary Note 7: Charge-head diagnostics ).
Figure 5: Spectral projection relative to H(0) and its effect on the bands. a , The predicted residual Δ=H−H(0) is whitened in the overlap metric and diagonalized. Eigenvalues outside [−b,b] (orange) are clipped, and adding the result back to H(0) gives Hb , whose bands obey the bound shown (Methods). For a near-singular S , overlap eigenvalues below 10−5 are dropped first (gray branch). b , Bands of LiZrPd 2 (mp-1185423), a test structure in which a strict spurious state is detected before projection. In its primitive cell, Li is pink, Zr green and Pd grey. The projection uses b=10 eV, and each spectrum is aligned to its own valence-band top.
Figure 6: Properties of amorphous materials from models trained on crystals. a , The amorphous Ge 2 Sb 2 Te 5 cell used in b – e , with Ge in purple, Sb in orange and Te in olive. b – d , Optical constants and conductivity of this cell, computed with the same Kubo–Greenwood protocol from the DFT, UniTB and UniTB-Flash Hamiltonians. The strips show model minus DFT, and energies in d are relative to the DFT midgap. e , Valence-band photoemission spectra computed for this cell and for a 216-atom amorphous Si network. Shape errors compare area-normalized spectra. f , Seven cells, one for each composition on the GeTe–Sb 2 Te 3 line. ρS is the Spearman correlation with DFT across the compositions, equal for the two models. Protocols are in Methods.
Figure 7: Model comparison and library screening. a , Two-stage screen of a library of public crystal structures. The pools keep structures that match nonmagnetic Materials Project entries near the convex hull (rules in Methods), and translucent cubes are structures outside the pools. Blue cells count structures with valid UniTB-Flash values, and the pools overlap. Orange cells are the UniTB-Flash top 200, which UniTB evaluates. Median ∣Δ∣ is the median of ∣ model − DFT ∣ /DFT over the UniTB top 20 with a DFT value. b , Comparison of the two models on nonmagnetic test crystals in six domains: photonics ( n ), photovoltaics (SLME), semiconductors ( Eg and the window), mobility ( mc ), thermoelectrics (PF, n- and p-type combined) and interconnects ( ρ0λ ). The radius is a model’s score or inverse error relative to the better model. Screening is scored by the mean top-20 recall over K=5 –60 (maximum 0.893), or by F1 for the band-gap window. Filled and open markers show whether the paired 95% interval of the gain lies above zero or includes it. c , The final 20 n -type thermoelectric candidates in UniTB order, with values on the final k -point meshes. PF/τ is the power factor per relaxation time at 1020 cm -3 and 300 K. The reference median and Mg 2 Si are the median and the strongest of the reference materials; grey marks the better side of these lines. Stars mark large power factors measured in pure or doped samples 57 , 58 , 59 .
Figure 8: SO2CUDA architecture and performance. a , Data flow of SO2CUDA for one edge of the UniTB edge update, with four experts and top-2 routing; depth shows channels. Rotate-as-pack rotates the features into the bond frame and writes expert-sorted copies. One grouped matrix multiplication per copy applies the expert maps, and rotate-as-gather rotates back and writes each output once. The backward pass reuses both kernels. The plain PyTorch band shows the same layer evaluated in separate loops over orders and experts. b , Forward-plus-backward time of the SO(2) linear layer relative to SO2CUDA. Points are 14 uniform and 12 nonuniform configurations, and bars are medians. c , Median time of one training step of the onsite network on real crystal batches, for UniTB and the dense model with three backends. Columns give the time, its ratio to SO2CUDA and the peak memory. All timings use one H200 GPU (Methods).
Setting
Value (onsite and hopping models unless stated)
Backbone
Interaction layers
3
Channels at angular degrees 0–6
128, 24, 16, 16, 32, 24, 48
Latent hidden widths; output width
200, 128; 128
Radial basis size; readout channels
128; 32
Element conditioning
Full tensor products in interaction layers and the output readout
Table 20: Module settings of the trained models. Settings of UniTB at both training budgets and of UniTB-Flash at MPHam-29K. Each independently trained model has its own response head. The onsite distance mask is d<10−6 Å. Hopping retains graph edges with d≥10−6 Å and has no upper cutoff beyond the stored graph (Supplementary Note Supplementary Note 5: Gauge alignment and graph conventions ).
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.
Density Functional Theory (DFT) is widely used for first-principles simulations in chemistry and materials science, but its computational cost remains a key limitation for large systems. Motivated by recent advances in ML-based exchange-correlation (XC) functionals, this paper introduces a differentiable framework that integrates machine learning models into density functional theory (DFT) for solids and other periodic systems. The framework defines a clean API for neural network models that can act as drop in replacements for conventional exchange-correlation (XC) functionals and enables gradients to flow through the full self-consistent DFT workflow. The framework is implemented in Python using a PyTorch backend, making it fully differentiable and easy to use with standard deep learning tools. We integrate the implementation with the DeepChem library to promote the reuse of established models and to lower the barrier for experimentation. In initial benchmarks against established electronic structure packages (GPAW and PySCF), our models achieve relative errors on the order of 5-10%.
Rakshit Kumar Singh, Aryan Amit Barsainyan, Bharath Ramsundar
Machine-learning electronic Hamiltonians achieve orders-of-magnitude speedups over density-functional theory, yet current models omit long-range Coulomb interactions that govern physics in polar crystals and heterostructures. We derive closed-form long-range Hamiltonian matrix elements in a nonorthogonal atomic-orbital basis through variational decomposition of the electrostatic energy, deriving a variationally consistent mapping from the electron density matrix to effective atomic charges. We implement this framework in HamGNN-LR, a dual-channel architecture combining E(3)-equivariant message passing with reciprocal-space Ewald summation. Benchmarks demonstrate that physics-based long-range corrections are essential: purely data-driven attention mechanisms fail to capture macroscopic electrostatic potentials. Benchmarks on polar ZnO slabs, CdSe/ZnS heterostructures, and GaN/AlN superlattices show two- to threefold error reductions and robust transferability to systems far beyond training sizes, eliminating the characteristic staircase artifacts that plague short-range models in the presence of built-in electric fields.
Yang Zhong, Xiwen Li, Xingao Gong +1
Key Laboratory of Computational Physical Sciences (Ministry of Education), Institute of Computational Physical Sciences, State Key Laboratory of Surface Physics, and Department of Physics, Fudan University, Shanghai, 200433, China.