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

physics.chem-phAug 24, 2026

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

Abstract

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.

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). $E_{g}$ 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 $v_{\rm ext}^{\rm loc}$ and $n_{\rm SAD}$ . Initialize $n_{0}=n_{\rm SAD}$ and the mixing factor $\alpha_{0}$ .
2	for $i=0,\ldots,I_{\max}-1$ do
3	Construct $v_{\rm KS}=v_{\rm ext}^{\rm loc}+v_{\rm H}[n_{i}]+v_{\rm xc}[n_{i}]$ .
4	Predict $\widetilde{n}_{i}^{\rm FNO}=n_{\rm SAD}+F_{\theta}[v_{\rm KS},v_{\rm ext}^{\rm loc}]$ and project $n_{i}^{\rm FNO}=\mathcal{P}_{N_{e}}[\widetilde{n}_{i}^{\rm FNO}]$ .
5	Evaluate the relative fixed-point residual $R_{i}^{\rm rel}=\frac{1}{N_{e}}\!\int\|n_{i}^{\rm FNO}(\mathbf{r})-n_{i}(\mathbf{r})\|\,\text{d}\mathbf{r}$ .
6	if $R_{i}^{\rm rel}<\varepsilon$ then return $n^{\star}=n_{i}^{\rm FNO}$ .
7	Mix to obtain a new density $n_{i+1}=\mathcal{P}_{N_{e}}[(1-\alpha_{i})n_{i}+\alpha_{i}n_{i}^{\rm FNO}]$ .

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 $\chi_{s}$ (left) and inverse response $-\chi_{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

L^{1}

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

/N_{e}

are normalized per valence electron. (C) Density error versus number of heavy atoms. The QM9 bin at

\leq

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 ; $Pm\overline{3}m$ , 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 $V_{0}$ and bulk moduli $B_{0}$ . 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 $N_{g}$ on logarithmic axes. Symbols show Quantum ESPRESSO (black squares) and Kohn–Sham FNO (blue circles) timings; lines are empirical power-law fits $t\propto N_{g}^{p}$ , 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 $v_{\mathrm{ext}}^{\mathrm{loc}}+v_{\mathrm{H}}[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.

Density Functional Theory

See all 25 papers

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory
Apr 22, 2026Roman Remme, Fred A. HamprechtSurrogate Modeling Quantum Chemistry
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University
MALOQ: Massively Accelerated Learning of Operators for Quantum Transport
Jun 27, 2026Manasa Kaniselvan, Alexander Maeder, Denghui Lu +2Quantum Chemistry High-Performance Computing
D-ITET, ETH Zurich, Zurich, Switzerland
A Fixed-Point Neural Operator for Size- and Functional-Transferable Hamiltonian Prediction
Jun 12, 2026Yunhong Lou, Xihang Yue, Xinran Wei +2Quantum Chemistry Molecular Property Prediction
1Zhejiang University, Hangzhou, China. · 2Zhongguancun Academy, Beijing, China. · 3Zhongguancun Institute of Artificial Intelligence, Beijing, China.

Apr 22, 2026J/K move · Enter open · S save

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory

Roman Remme, Fred A. Hamprecht

Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University

Surrogate Modeling Quantum Chemistry Adaptive Sampling Computational Materials Science Density Functional Theory Computational Quantum Chemistry

Open paper AskPDF arXiv

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(N^3)

orthononormalization step required by prior work, yielding improved runtime scaling for larger systems.

Abstract

Density Functional Theory

Explore similar work

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory

MALOQ: Massively Accelerated Learning of Operators for Quantum Transport

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