Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory
Authors: Danish Khan, Maurice D. Hanisch, Nikolai Argatoff, Evan Xie, Sandeep Sharma, Anima Anandkumar
Organizations: Department of Computing and Mathematical Sciences, California Institute of Technology · Department of Mathematics, ETH Zürich · Division of Chemistry and Chemical Engineering, California Institute of Technology · Marcus Center for Theoretical Chemistry, Pasadena, CA 91125, USA
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 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.
Figures & tables
Figure 1: Summary of the Kohn–Sham Fourier Neural Operator (Kohn-Sham FNO) based self-consistent field (SCF) cycles for density optimization. (A) Conventional DFT and Kohn-Sham FNO workflows. The model replaces the orbital-solution step (see SI) while retaining the rest of a standard DFT SCF cycle. Our SCF implementations for both molecules and solids closely follow Quantum ESPRESSO (QE) 7.5 [ 27 , 28 ] with minor modifications for retaining quasi-linear scaling. All DFT calculations are also performed using QE. See SI. (B) Schematic of the model mapping an input potential to the density produced by solving the corresponding Kohn-Sham equations. See SI. (C) Band structures for tetragonal perovskite CaTiO 3 (MC3D [ 29 ] ID mc3d-69107 ; I4/mcm , No. 140) obtained using the Kohn-Sham FNO (blue) converged densities and regular DFT (black). Eg values denote the corresponding band gaps. See SI. (D) Density error versus molecular size for self-consistent and direct ground-state prediction on small-organic QM9 [ 30 ] and larger drug-like QMugs molecules [ 31 ] . See Fig. 3 . (E) Measured total SCF cycle wall time versus electron count for QE and Kohn-Sham FNO. See SI.
1
Construct vextloc and nSAD . Initialize n0=nSAD and the mixing factor α0 .
2
for i=0,…,Imax−1 do
3
Construct vKS=vextloc+vH[ni]+vxc[ni] .
4
Predict niFNO=nSAD+Fθ[vKS,vextloc] and project niFNO=PNe[niFNO] .
5
Evaluate the relative fixed-point residual Rirel=Ne1∫∣niFNO(r)−ni(r)∣dr .
6
if Rirel<ε then return n⋆=niFNO .
7
Mix to obtain a new density ni+1=PNe[(1−αi)ni+αiniFNO] .
Algorithm 1 Self-consistent field cycle with Kohn–Sham FNO.
Figure 2: Learning the forward orbital-free KS map avoids inverse-response instability. (A) Density error relative to PBE [ 36 ] ground-states for 100 held-out small-organic molecules from the QM9 dataset [ 30 ] . See SI. Kohn–Sham FNO (blue) maps the KS potential to its output density and is iterated to a fixed-point. The kinetic-potential FNO (red) uses the same domain-invariant FNO model as the backbone and learns the density to kinetic-potential relation and is used to solve the Euler equation based on the implementation in DFTpy [ 40 ] . See SI. Both models use the same training data (Fig. 3 A) and FNO backbone. Thin curves show individual molecules, bold curves show their mean, and circles and crosses mark converged and failed calculations, respectively. (B) Binned eigenspectra of the projected forward response χs (left) and inverse response −χs−1 (right) over 1,100 PBE SCF states for the same molecules. Modes are ordered from strongest to weakest forward response; points mark bin medians and surface color gives the number of eigenvalues per bin. The responses were computed using all-electron calculations in Gaussian basis sets with the PySCF [ 41 ] and KS-PIES [ 42 ] packages. See SI.
Figure 3: Improved extrapolation via the self-consistent Kohn–Sham map. (A) Elemental composition of the shared training set of 2,004 QM9 molecules [ 30 ] and 6,500 MC3D crystals [ 29 ] . Each value is the percentage of structures containing that element; gray denotes absence from training. (B) Mean errors of the self-consistent Kohn–Sham FNO (blue) densities and direct ground-state predicted densities on 100 held-out QM9 molecules containing the elements C, H, N, O, F and 200 larger drug-like molecules from the QMugs dataset [ 31 ] containing the elements C, H, N, O, F, S, Cl, P. From left to right, the metrics are the normalized L1 density error, Euclidean norm of the dipole-vector error, Frobenius norm of the quadrupole-tensor error, electrostatic-potential error, Coulomb-weighted density-error metric, and absolute XC-energy error; quantities labeled /Ne are normalized per valence electron. (C) Density error versus number of heavy atoms. The QM9 bin at ≤ 9 heavy atoms consists of the same 100 test molecules from Fig. 2.1 . The QMugs bins from 20−35 heavy atoms each consist of 50 randomly sampled ground-state conformers. All errors are relative to converged Quantum ESPRESSO PBE calculations. Data construction, the direct model, and the density metric: see SI.
Figure 4: Generalization across periodic systems, SCF trajectories, geometries and XC approximation. (A) Fixed-point trajectories and final density errors for a held-out MC3D test set composed of 100 semiconductors and 100 metals as characterized by converged PBE calculations. Chemical composition of these crystals spans the same subspace of the periodic table as in Fig. 3 A. Thin curves denote individual crystals, dots and crosses mark converged and failed calculations, and bars summarize the mean errors for semiconductors and metals. (B) Results for metallic cubic Sr 3 SnO (MC3D ID mc3d-54098 ; Pm3m , No. 221). Left, band structure and density of states from one fixed-density post-SCF diagonalization using the FNO density (blue) and from self-consistent PBE (black dashed), each referenced to its Fermi energy. Right, clamped-ion PBE (squares) and PBEsol [ 37 ] (circles) equations of state; solid blue and dashed black curves denote FNO and DFT, respectively, and the table reports fitted equilibrium volumes V0 and bulk moduli B0 . The PBEsol test uses the PBE-trained model without fine-tuning and replaces only the XC potential in the KS potential construction during inference. See SI. Training, post-SCF observables, and EOS fitting: see SI.
Figure 5: Near linear-scaling density optimization for large metallic dislocation cells. (A) Time per complete SCF iteration t versus the number of real-space grid points Ng on logarithmic axes. Symbols show Quantum ESPRESSO (black squares) and Kohn–Sham FNO (blue circles) timings; lines are empirical power-law fits t∝Ngp , where p is the fitted scaling exponent. Quantum ESPRESSO calculations were run on 192-core CPU nodes, whereas the FNO used one NVIDIA B300 GPU. See SI. (B) Fixed-point residuals for pyramidal-I and pyramidal-II dislocation cells using the base model (dashed), trained on the QM9 and MC3D training set in Fig. 3 A, and the model fine-tuned on Mg-defect data (solid; see SI); color denotes valence-electron count. (C) Difference in the electrostatic potential vextloc+vH[n] between the 8,250-atom pyramidal-I and pyramidal-II cells using Kohn–Sham FNO converged densities. Grey dots denote the Mg atom locations from the pyramidal-I cell. The heat map shows the 99.5th percentile difference. (D) The dislocation core at full resolution. The region boxed in (C) contains 600 atoms and is magnified roughly 3.9x. Training, fine-tuning, and inference data generation, convergence, and timing calculations specific to these results: see SI.
We introduce surrogate functionals: machine-learned energy functionals for orbital-free density functional theory (OF-DFT) which are defined not by universal fidelity to a physical reference, but merely by the requirement that density optimization with a fixed procedure yields the true ground-state density. Helpfully, training surrogate functionals requires only ground-state densities, no energies or gradients away from the ground state. We here propose a gradient-descent-improvement loss that guarantees exponential convergence of the density to the ground state, and combine it with an adaptive sampling scheme that concentrates learning around the optimization trajectories actually visited during inference. On the QM9 and QMugs benchmarks, surrogate functionals achieve density errors competitive with or improving upon the state of the art for fully supervised machine-learned OF-DFT, while eliminating the need for the O(N3) orthononormalization step required by prior work, yielding improved runtime scaling for larger systems.
Roman Remme, Fred A. Hamprecht
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University
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
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.