Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

TL;DR: Every Kohn–Sham DFT calculation spends most of its time repeatedly solving the same kind of eigenvalue problem: given an effective potential, find the electron density of non-interacting electrons in it. This paper argues that this map itself — potential in, density out — is the right thing to learn with a neural network, rather than the two targets the field has tried before (a kinetic-energy functional, or the final converged answer). A single SE(3)-equivariant Fourier neural operator trained on 8,504 molecules and crystals then drives self-consistent field loops to convergence across organic molecules, insulators, and metals, and — after fine-tuning on small magnesium cells — converges an 8,250-atom, 82,500-electron dislocation on one GPU with empirical O(N^{1.03}) scaling, versus O(N^{3.37}) for Quantum ESPRESSO. The interesting part isn't just the scale; it's the clean argument, backed by controlled experiments, for why this learning target works where the others fail.

The bottleneck, and why the obvious fixes miss

Kohn–Sham DFT is arguably the most-run quantum simulation on Earth (the paper cites ~30% of NERSC's 2018 workload). Its structure is a fixed-point loop. You guess an electron density n(\mathbf{r}), build an effective potential v_{\mathrm{KS}}[n] from it (external + Hartree + exchange-correlation terms, all cheap and explicit), then solve a non-interacting Schrödinger equation in that potential:

where the \phi_p are orbitals, \varepsilon_p their energies, and f_p occupations. Feed n^{\mathrm{out}} back in, repeat until n^{\mathrm{out}} = n. The eigensolve is the expensive part — O(N^3) in electron count, done at every iteration, at every k-point for solids.

Two ML strategies have attacked this, and the paper's framing puts them at opposite extremes of a spectrum:

Learn a kinetic-energy functional (T_s[n], the "orbital-free DFT" dream). Then you can minimize energy over the density directly. But used variationally, what you actually need is the functional derivative \delta T_s/\delta n, which by the Euler equation equals \mu - v_{\mathrm{KS}} — i.e., you've implicitly learned the density-to-potential map. This is the inverse of the physical solve, and as we'll see, it's catastrophically ill-conditioned.

Learn the endpoint directly (atomic structure → converged density, or → converged Hamiltonian). This asks the model to compress an arbitrarily long, nonlinear, XC-functional-specific SCF trajectory — including long-range charge redistribution — into one forward pass. It works in-distribution and falls apart on bigger systems.

The aha: learn the middle step, not the ends

The insight is that a "just right" target has been sitting inside every KS calculation all along: the forward Kohn–Sham map,

the answer to a single non-interacting eigenproblem. Learn that, keep everything else in the SCF loop exactly as it was — Hartree and XC potential construction, density mixing, convergence checks — and you've surgically removed only the O(N^3) piece:

Three properties make this target unusually good, and each gets its own experiment in the paper.

It's XC-agnostic and universal. The non-interacting solve doesn't care how the input potential was constructed. Chemical composition, geometry, and choice of exchange-correlation functional all change the input, not the operator. So potential–density pairs from different XC functionals are samples of one operator and can be pooled in training — and a PBE-trained model can run PBEsol SCFs without retraining (they test this; more below).

It's densely supervised for free. Every iteration of every conventional SCF run yields an exact (potential, density) training pair — not just the converged endpoint. The 8,504 training structures yield 59,500 labels, and even unconverged reference trajectories are usable data.

Errors get corrected instead of accumulated. A mistake in one update perturbs the next input potential, and subsequent iterations — plus density mixing — can fix it. This is the same logic as chain-of-thought over one-shot answers (the authors make this analogy explicitly): keep the intermediate computation, don't compress it away.

Why direction matters: the conditioning argument

The most satisfying part of the paper is Section 2.1, which explains why learned kinetic-energy functionals have historically blown up in self-consistent use — not with hand-waving, but with the spectrum of a physical operator.

Linearize the forward map: a small potential change produces a density change \delta n = \chi_s\, \delta v_{\mathrm{KS}}, where \chi_s is the non-interacting density response. Linearizing the Euler equation gives the inverse relation: the learned kinetic potential's Jacobian is -\chi_s^{-1}. The authors computed \chi_s's spectrum across 1,100 SCF states of 100 QM9 molecules and found what physics predicts: many modes with tiny eigenvalues |\lambda_j(\chi_s)| — directions in which the density barely responds to the potential.

In the forward direction, these weak-response modes suppress model errors. Inverted, they amplify errors by |\lambda_j|^{-1} — several orders of magnitude. And there's a second, more structural problem: the forward map is defined for any potential, while the inverse is defined only for "non-interacting v-representable" densities, a domain with no known complete characterization. An optimizer wandering through density space will happily query the learned inverse where no exact answer even exists.

The controlled experiment is the killer detail: same FNO backbone, same training data, only the mapping direction flipped. The forward model converges every test SCF, driving density errors from 10–18% (initial guess) to below 1%. The inverse model fails after a few iterations on essentially everything. For calibration, prior work (Remme et al.) got stable learned-functional SCFs on QM9 only by training on ~2.25 million labels including specially constructed perturbed densities — and still saw a 28% failure rate without them. This paper uses 59,500 ordinary SCF labels.

The architecture, briefly

The map is between spatial fields on a grid, so a Fourier neural operator is the natural choice: spectral convolutions capture the global nonlocality of quantum mechanics at O(N_g \log N_g) cost in grid points N_g — no cutoff radius (unlike GNNs), no quadratic attention (unlike transformers). Two modifications matter. First, domain invariance: standard FNO spectral weights are indexed by integer mode number, tying them to one box size; here, radially factorized filters are functions of physical reciprocal-space radius |\mathbf{k}|, so the same learned filters can be sampled on the k-grid of any cell — one model for a small molecule box and a huge crystal supercell. Second, radial filters plus spherical mode truncation make the spectral convolution rotation-equivariant, giving full SE(3) equivariance. For periodic systems there's a bonus: the loop evolves only the lattice-periodic density on the unit cell, with no explicit k-point sampling — one FNO call replaces N_{\mathbf{k}} diagonalizations per iteration.

The evidence

Extrapolation vs. direct prediction (same backbone, same data, only the target changed). In-distribution on held-out QM9 molecules, both are fine (0.625% vs. 0.662% density error). Out of distribution — QMugs drug-like molecules, larger and containing S/Cl/P in molecular environments never seen in the molecular training data — the gap opens dramatically:

Dipole errors on QMugs drop from 0.237 to 0.026 D per electron, electrostatic-potential errors from 40.8 to 3.28 mHa. The one-shot model must extrapolate an entire converged solution; the iterative model just keeps applying the same local operator and lets feedback do the rest.

Solids, spectra, and XC transfer. On 200 held-out MC3D crystals, all but one SCF converges within 80 iterations; mean density errors are 0.75% (semiconductors) and 1.46% (metals) — metals being the historical graveyard of orbital-free methods. For metallic Sr₃SnO, a single fixed-density diagonalization on the FNO density reproduces the PBE band structure and DOS, with a Harris total energy off by 2.29 meV/atom. Equation-of-state curves match to 0.23% in equilibrium volume and 1.1% in bulk modulus. And the operator-universality claim survives its test: swapping PBE→PBEsol in the potential construction at inference only, with zero retraining, gives V_0 within 0.61% and B_0 within 7.7% of self-consistent PBEsol.

The 82,500-electron test. The Mg \langle c+a\rangle screw-dislocation problem is a deliberate flex: the 2019 Gordon Bell finalist DFT-FE study needed 1,300 Summit nodes (7,800 V100s) for a 6,164-atom ground state. Here, the base model diverged on these cells — and, importantly, announced its own failure through a runaway fixed-point residual, a built-in reliability diagnostic no one-shot predictor has. After fine-tuning on 1,203 Mg structures of at most 364 atoms (bulk, surfaces, stacking faults, core fragments), every dislocation cell converges to the standard 10^{-3} residual, up to 8,250 atoms on a single NVIDIA B300. Fitted scaling exponents: 1.03 for the FNO pipeline vs. 3.37 for Quantum ESPRESSO. On core-centered crops small enough to verify against reference PBE (≤528 atoms), density errors sit at 0.33–0.35% with no growth in system size.

What to be skeptical of

The paper is refreshingly honest about its scope, and the caveats matter:

Only half the map is learned. The full orbital-free target is (n^{\mathrm{out}}, T_s); this work learns only the density. Total energies and band structures still require one post-SCF orbital diagonalization at fixed density. That's a huge win over diagonalizing at every iteration and k-point, but the 82,500-electron result is a converged density, not yet a converged energy — the pyramidal-I vs. pyramidal-II core energy difference that motivates the Mg problem is not actually computed. The kinetic-energy component is promised for future work, and its conditioning may not be as forgiving.

The headline timing is not hardware-matched. FNO on one B300 GPU vs. Quantum ESPRESSO CPU timings linearly rescaled by core count to a hypothetical 192-core node. The scaling exponents are the trustworthy comparison; the absolute wall-times are not.

Accuracy at the largest scale is indirectly verified. No PBE reference exists for the 8,250-atom cells; validation is on ≤528-atom crops plus a qualitative electrostatic-potential map. The residual diagnostic catches divergence but, as the authors note, "convergence alone does not certify agreement" with the true ground state.

Generalization has edges. The base model failed outright on dislocations — defect-scale physics required targeted fine-tuning, so this is not yet a foundation model you point at arbitrary matter. The PBEsol bulk-modulus transfer error (7.7%) is noticeably worse than in-functional results, and residual spikes on very large grids hint at numerical fragility. Everything is norm-conserving pseudopotentials at the GGA level; ultrasoft/PAW, spin polarization, and hybrids are untested.

Why it matters, and where to read

If this holds up, the practical consequence is that near-linear-scaling DFT stops depending on density-matrix sparsity assumptions (which fail for metals) and starts depending on a learned operator plus commodity GPU memory — with a self-diagnosing failure mode and an unusually cheap supervision signal, since every legacy SCF log in every DFT archive is training data for the same universal operator. The conceptual contribution may outlive the specific model: "replace the dominant repeated operation, keep the physical iteration that constructs and validates the answer" is a design principle that plausibly transfers to other iterative solvers in scientific computing.

Section 2.1 is the part to read carefully — the forward/inverse controlled comparison and the \chi_s spectral analysis (Fig. 2) are a genuinely clarifying explanation of a decade of orbital-free ML failures, and they stand on their own even if the large-scale claims need independent replication.