Towards a universal model for spin-orbit coupled Wannier Hamiltonians

Abstract: While machine learning interatomic potentials (MLiPs) have matured to revolutionize material science, deep learning models for electronic structure are just beginning to emerge and restricted, almost exclusively, to non-orthogonal basis Hamiltonians. We introduce G(Wa)NN, the first deep-learning model capable of generating the electronic Hamiltonian of solid-state systems in an orthogonal Wannier basis. G(Wa)NN is trained on an unprecedented, diverse dataset of more than 111K Wannier Hamiltonians (150M+ hopping matrices) spanning 69 elements. The combination of optimized inference and linear-scaling methods for orthogonal Hamiltonians unlock transport simulations at massive scales (10K+ atoms). Crucially, the framework supports local finetuning, allowing users to adapt the base model to custom Wannier Hamiltonian datasets. To seamlessly translate these predictions into physical observables, we introduce Tailwater, a Python package providing an API interface to G(Wa)NN alongside a high performance post-processing library. Tailwater enables automated projection of the predicted Hamiltonian into an arbitrary low-energy subspace-directly mirroring familiar Wannier90 workflows-and includes a suite of Kernel Polynomial Method (KPM) functions that exploit the orthogonal basis to achieve strict linear scaling for spectral observables. The Tailwater ecosystem, with the G(Wa)NN model at its core, aims to help bridge the gap between deep learning and macro-scale quantum transport simulations.

• The paper introduces G(Wa)NN, an E(3)-equivariant graph neural network trained on 111,905 Wannier Hamiltonians across 69 elements and 210 space groups to generate sparse, orthogonal, SOC-resolved Hamiltonians for large-scale simulations.
• G(Wa)NN achieves a 0.080 eV band-gap MAE with R² = 0.95 and a median normalized DOS MAE of 0.035, while its momentum-space loss prioritizes Fermi-level accuracy, level spacings, and symmetry-protected band crossings.
• The Tailwater software validates practical applications through Bi₂Se₃ topology and transport calculations, supports subspace projection and fine-tuning, and enables efficient KPM-based simulations for systems exceeding 10,000 atoms.

G(Wa)NN is an $E(3)$-equivariant graph neural network that predicts spin-orbit coupled electronic Hamiltonians of crystalline solids directly in an orthogonal Wannier basis, accompanied by Tailwater, a Python package providing API access to the model and a post-processing library for transport observables. The work addresses a specific gap in machine learning for electronic structure: existing foundational Hamiltonian models such as Uni-HamGNN, DeepH-E3, and MACE-H operate in non-orthogonal LCAO bases, where the generalized eigenvalue problem $H\Psi = SE\Psi$ complicates downstream linear-scaling methods. The orthogonalization $H_{\text{orth}} = S^{-1/2}HS^{-1/2}$ destroys sparsity, so KPM-based $O(N)$ transport codes cannot be applied efficiently. By targeting a naturally orthogonal, maximally localized Wannier basis instead, G(Wa)NN produces strictly sparse Hamiltonians suitable for systems of 10K+ atoms.

Dataset construction

The training corpus comprises 111,905 Wannier Hamiltonians—over 150 million distinct hopping matrices $\hat{H}_{ij}^{\mathbf{R}}$—spanning 69 elements and 210 space groups, drawn from MP-Traj, C2DB, GNOME, and Alexandria. This is claimed to be the largest dataset of Wannier Hamiltonians assembled to date. The gauge covariance of Hamiltonian targets is emphasized as the reason ad-hoc data aggregation is impossible: unlike energies and forces in MLIP training, all data must be generated with systematic consistency.

The generation protocol uses Quantum ESPRESSO with fully relativistic PseudoDojo norm-conserving pseudopotentials, SOC included, a 60 Ry cutoff, and a uniform 6×6×6 Monkhorst-Pack grid. A fixed universal orbital basis is imposed by atomic number ($s$ only for $Z \leq 2$; $s,p$ up to $Z=18$; $s,p,d$ above), yielding 18×18 complex blocks per atom pair after doubling for spin. Disentanglement with a frozen window extending 1 eV above the Fermi level is retained deliberately: it improves accuracy near the Fermi energy and maximal localization acts as partial gauge fixing. Two acceptance criteria are enforced—a mean absolute error below 0.05 eV between DFT and Wannier band structures at random $H\Psi = SE\Psi$0-points, and Wannier center spreads below 1.5 times the largest lattice vector.

Notably, out-of-equilibrium structures from MP-Traj relaxation trajectories are included (89,331 structures from 34,564 materials), motivated by finite-temperature effects and surface/interface reconstruction. The paper concedes that disentanglement can inject noise into the conduction sector of the target data.

Model architecture

The network follows operator-learning formalism built on full O(3) Clebsch–Gordan tensor products via e3nn, rather than reduced SO(2) couplings. Radial weights are conditioned on both interatomic distance and chemical identities $H\Psi = SE\Psi$1. Spin is attached by coupling spatial tensors to a spin-scalar channel ($H\Psi = SE\Psi$2) and a spin-vector channel ($H\Psi = SE\Psi$3), with predicted blocks formed as covariant combinations of these basis tensors.

Training uses a composite loss combining real-space matrix-element MAE with a momentum-space term containing three components: a band-edge-anchored eigenvalue error that down-weights states far from the Fermi level, an offset-invariant level-spacing term, and a crossing-protection term concentrated on near-degenerate pairs to avoid artificially gapping symmetry-enforced crossings. The design intent is accuracy concentrated where gaps, effective masses, and transport properties are determined.

Benchmarks

The authors argue matrix-element losses are misleading under gauge covariance and instead report gauge-invariant observables on a 2000-material validation set from GNOME:

Metric	Result
Matrix-element MAE	10 meV
Normalized DOS MAE (median)	0.035
Band-gap MAE / $H\Psi = SE\Psi$4	0.080 eV / 0.95

For comparison, models trained identically on the same data yield: CGCNN 0.303 eV ($H\Psi = SE\Psi$5), ALIGNN 0.092 eV (0.93), coGN 0.111 eV (0.91). Since coGN and ALIGNN rank first and fourth on the MatBench band-gap leaderboard, the authors claim G(Wa)NN is at or approaching state-of-the-art for band-gap prediction. Two caveats apply: direct comparison to the DOS results of prior multimodal and PET approaches is impossible due to differing train/test splits, and the validation set has a low mean gap of 0.47 eV, making the 0.08 eV MAE nontrivial but also reflecting a distribution where small errors matter proportionally more.

Subspace projection and topological validation

Tailwater automates projection of the predicted full-basis Hamiltonian into a user-defined energetic window via Gaussian-weighted fitting, emulating Wannier90 workflows at negligible local computational cost. The Bi$H\Psi = SE\Psi$6Se$H\Psi = SE\Psi$7 demonstration validates the full pipeline: agreement across all 90 bands against plane-wave DFT, correct $H\Psi = SE\Psi$8 invariant from Z2Pack in both full and projected (30-band) representations, resolution of the gapless surface Dirac cone via KPM, and multi-orders-of-magnitude runtime reduction for spin-Hall conductivity computed with WannierBerri as $H\Psi = SE\Psi$9-grids densify from 4³ to 128³. Because Bi$H_{\text{orth}} = S^{-1/2}HS^{-1/2}$0Se$H_{\text{orth}} = S^{-1/2}HS^{-1/2}$1's topology depends entirely on SOC-driven band inversion at Γ, this case directly tests the model's relativistic fidelity—and it passes.

Fine-tuning and scale

Local fine-tuning requires only the model head (~61K parameters) and accepts Wannier Hamiltonians from any DFT code or basis set, with automatic orbital masking when the fine-tuning basis differs from the training basis. Applied to stacked-and-shifted WSe$H_{\text{orth}} = S^{-1/2}HS^{-1/2}$2 bilayers—including DFT-D3 dispersion corrections absent from the base training set—the baseline already shows reasonable accuracy, and short fine-tuning substantially improves eigenvalue agreement, producing band structures for θ = 32.2° twisted bilayer WSe$H_{\text{orth}} = S^{-1/2}HS^{-1/2}$3 in close agreement with direct QE calculations. Inference benchmarks on an Apple M3 Max show large speed advantages over Uni-HamGNN's pipeline and orders-of-magnitude speedup over QE on an AMD EPYC 7713 cluster, supporting the stated goal of real-time inference usable by AI agents during high-throughput screening.

Limitations and open questions

Several restrictions are explicit. Magnetic systems, lanthanides, and elements with $H_{\text{orth}} = S^{-1/2}HS^{-1/2}$4 are excluded from training, so spin-ordered and f-electron chemistry remain outside the model's coverage. Hybrid functionals (HSE06) are not incorporated, leaving the well-known DFT band-gap underestimate uncorrected except through fine-tuning. Electron-phonon coupling lacks a native interface. The disentanglement procedure introduces noise in conduction bands, and the uniform 6×6×6 k-grid prioritizes Wannier quality over convergence guarantees for larger cells. Finally, the DOS comparison to prior work is qualitative because test sets differ—an unavoidable but unresolved benchmarking ambiguity in this subfield.

Conclusion

This paper delivers two coupled contributions: a foundational equivariant model generating sparse, orthogonal, SOC-resolved Wannier Hamiltonians trained on an unprecedented corpus, and an open-source ecosystem translating those predictions into transport-relevant observables with linear-scaling KPM routines. The band-gap benchmark (0.080 eV MAE, $H_{\text{orth}} = S^{-1/2}HS^{-1/2}$5) and the preserved topological classification through inference and subspace reduction constitute the strongest quantitative evidence that deep-generated Wannier Hamiltonians can be physically trustworthy. Whether the framework sustains its accuracy across the planned expansion to 500K+ Hamiltonians, magnetic order, and hybrid functionals remains the central open question for establishing it as a general-purpose replacement for first-principles electronic structure at device scales.