Source-linked AI summary
Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory
Danish Khan, Maurice D. Hanisch, Nikolai Argatoff, Evan Xie, Sandeep Sharma, Anima Anandkumar
TL;DR
Repeated orbital diagonalizations make KS-DFT scale cubically, motivating orbital-free methods that avoid auxiliary orbitals without sacrificing KS accuracy. This work learns the forward KS potential-to-density map with a domain-invariant SE(3)-equivariant FNO, replacing orbital solves inside SCFs. The method converges molecular, insulating, and metallic systems and reaches dislocation cells with up to 82,500 valence electrons on one GPU.
Problem
Repeated orbital diagonalizations give KS-DFT cubic scaling, while prior orbital-free and direct ground-state learning approaches have not provided accurate, transferable density-only calculations at larger scales.
Method
A domain-invariant, SE(3)-equivariant Fourier neural operator learns the KS map’s density component from a KS potential and replaces the eigensolver within self-consistent iterations.
Results
The same method converges SCFs across organic molecules, semiconductors, and metals, while producing KS-DFT-quality densities, electronic observables, spectra, and equations of state.
Takeaways & Limitations
After fine-tuning on cells with at most 364 atoms, the model converges magnesium dislocation densities with up to 8,250 atoms and 82,500 valence electrons on one GPU.
Abstract
from arXiv · showhide
Kohn--Sham density functional theory (DFT) underpins electronic-structure simulations, but repeated orbital diagonalizations lead to cubic scaling, restricting quantum calculations to modest scales only. Eliminating these auxiliary orbitals while retaining Kohn--Sham accuracy is the central goal of orbital-free DFT, but both analytical and machine-learning methods have so far fallen short. Prior learning approaches either try to learn the variational kinetic-energy functionals, which are ill-conditioned, or directly predict the ground state, which extrapolate poorly to larger systems. Instead, we identify the Kohn--Sham map as the right learning target for orbital-free DFT. It maps a Kohn--Sham potential directly to the corresponding density and noninteracting kinetic energy, quantities otherwise obtained through an orbital diagonalization. Focusing on the density component in this work, a domain-invariant $\mathrm{SE}(3)$-equivariant Fourier neural operator learns to predict it from the potential as input on real-space grids, enabling stable quasi-linear scaling SCFs. Trained jointly on 8,504 molecules and solids, a single model generalizes to out-of-distribution organic molecules, insulators, and metals. For the first time, the same method converges SCFs across these systems without explicitly constructing Kohn--Sham orbitals, while reproducing densities, electronic spectra, and structural observables at Kohn--Sham DFT accuracy. Linear-scaling SCFs additionally allow converging magnesium dislocation densities containing up to 82,500 valence electrons on a single GPU.
1 Introduction
Kohn–Sham DFT is accurate and widely used, but repeated orbital solutions create cubic scaling that limits first-principles simulations at larger scales. Orbital-free approaches seek density-only calculations, yet existing analytical and learned kinetic-energy functionals and direct ground-state predictors have important limitations.
- O(N3) orbital solving and orthogonalization at every SCF iteration restrict the length scales accessible to extended-defect, interface, and electrochemical simulations.
- 30% of the workload at NERSC in 2018 was attributed to DFT calculations.
- Existing near-linear-scaling methods extend KS-DFT, but their performance depends on density-matrix decay, sparsity, dimensionality, or stochastic error.
- No general framework currently provides near-linear-scaling DFT with KS accuracy across molecules, metals, insulators, and chemically heterogeneous systems.
- Orbital-free DFT targets density-only energy minimization, but analytical KEDFs lack needed accuracy and transferability while learned KEDFs require difficult functional derivatives for stable SCFs.
- Direct ground-state models learn the endpoint of an arbitrarily long, XC-specific trajectory instead of retaining KS-DFT’s intermediate SCF computation.
A Self-consistent field (SCF) cycle
The proposed Kohn–Sham map replaces each orbital-solution step with a neural operator mapping an SCF potential to its output density. A domain-invariant, SE(3)-equivariant FNO enables quasi-linear updates while preserving self-consistent feedback and explicit potential construction.
- Kohn–Sham map: At each SCF iteration, the density defines a KS potential whose non-interacting solution and occupied eigenstates produce the output density.
- Kohn–Sham map: The learned density map vKS[n](r) 7−→ nout(r) replaces the orbital solve for stable orbital-free SCFs, while energies and spectra use fixed-density post-SCF diagonalization.
- Kohn–Sham map: Compared with inverse OF-DFT and direct prediction, the forward KS map is defined by a single non-interacting Hamiltonian solution and retains intermediate SCF computation.
- Neural-operator architecture: The FNO captures global nonlocality with quasi-linear O(Ng log Ng) scaling on the real-space grid, making the complete update scale as O(Ng log Ng).
- Neural-operator architecture: Domain-invariant radial filters can be sampled on reciprocal-space grids for differently sized domains, while spherical truncation and translation equivariance provide full SE(3) equivariance.
2 Results & Discussion
The forward Kohn–Sham map provides a stable, self-consistent learning target that improves extrapolation and enables orbital-free SCFs across molecules, solids, and large metallic defects. The method reproduces key observables while achieving near-linear empirical scaling on large dislocation cells.
- Stable optimization: 59,500 ordinary SCF labels from 8,504 structures support stable convergence across molecular and periodic domains, compared with 2.25 million labels used in a prior energy-functional study.The authors caution that the architectures are not directly comparable, but emphasize the contrast in training requirements and convergence scope.
- Stable optimization: The forward KS map avoids inverse-response amplification by learning potential-to-density prediction rather than the ill-conditioned density-to-potential relation.Weak-response modes suppress forward potential errors but amplify density errors under inversion; the instability persists along SCF trajectories.
- Improved extrapolation: 2.23% versus 9.97% density error on larger QMugs molecules shows better extrapolation for self-consistent Kohn–Sham FNO than direct ground-state prediction.The QMugs set is larger and includes molecular bonding environments involving S, Cl, and P absent from the molecular training subset.
- Periodic systems: All but one of 200 held-out crystals converged within 80 iterations, with mean density errors of 0.75% for semiconductors and 1.46% for metals.The fixed-point residual decreased smoothly across the periodic test set.
- Periodic systems: 0.23% and 1.1% are the PBE FNO–DFT differences in fitted equilibrium volume and bulk modulus for Sr3SnO across seven displaced volumes.The equation-of-state test probes geometries beyond the equilibrium structures used for training.
- Magnesium dislocation: p = 1.03 for Kohn–Sham FNO versus p = 3.37 for Quantum ESPRESSO demonstrates near-linear empirical SCF scaling on magnesium dislocation cells.The largest converged cell contains 8,250 Mg atoms and 82,500 valence electrons on one NVIDIA B300 GPU; adaptive density mixing handles occasional residual spikes.
3 Conclusion
The Kohn–Sham map is proposed as an effective orbital-free DFT learning target, with its density component replacing repeated orbital solutions during SCF while preserving the physical iteration. The method generalizes across molecular, semiconducting, and metallic systems, but currently learns only density output and still requires a final orbital solve for some observables.
- The model replaces one non-interacting orbital solution at an arbitrary SCF iteration rather than predicting the complete ground state in one step.This factorization targets the repeated operation inside the physical SCF procedure.
- Controlled comparisons show more stable density optimization than the inverse kinetic-potential map and more reliable extrapolation than direct ground-state prediction.
- A single Kohn–Sham FNO drives molecular, semiconducting, and metallic SCFs and transfers across training settings including PBE to PBEsol without retraining.
- The present model learns only the density component of the joint Kohn–Sham operator.
- Orbital-resolved observables and total energies use one fixed-density post-SCF calculation, while eliminating that solve requires learning kinetic-energy or free-energy output.