NeuralSCF: Neural network self-consistent fields for density functional theory

TL;DR

NeuralSCF introduces a mechanics-based machine learning framework that learns the Kohn-Sham density map as a self-consistent field, using a SE(3)-equivariant graph transformer to operate on density coefficients. By training explicitly on KS-SCF trajectories and then applying implicit differentiation for fine-tuning, it achieves state-of-the-art accuracy in predicting self-consistent densities and derived electronic properties, with strong zero-shot generalization to out-of-distribution systems. The approach highlights the value of embedding KS mechanics into ML objectives, enabling accurate, transferable predictions and potential acceleration of large-scale or high-throughput DFT calculations. This work points toward universal electronic-structure surrogates that leverage the intrinsic physics of KS-DFT to improve extrapolation and efficiency across diverse chemical spaces.

Abstract

Kohn-Sham density functional theory (KS-DFT) has found widespread application in accurate electronic structure calculations. However, it can be computationally demanding especially for large-scale simulations, motivating recent efforts toward its machine-learning (ML) acceleration. We propose a neural network self-consistent fields (NeuralSCF) framework that establishes the Kohn-Sham density map as a deep learning objective, which encodes the mechanics of the Kohn-Sham equations. Modeling this map with an SE(3)-equivariant graph transformer, NeuralSCF emulates the Kohn-Sham self-consistent iterations to obtain electron densities, from which other properties can be derived. NeuralSCF achieves state-of-the-art accuracy in electron density prediction and derived properties, featuring exceptional zero-shot generalization to a remarkable range of out-of-distribution systems. NeuralSCF reveals that learning from KS-DFT's intrinsic mechanics significantly enhances the model's accuracy and transferability, offering a promising stepping stone for accelerating electronic structure calculations through mechanics learning.

Figures (4)

• Figure 1: Overview of NeuralSCF.(a) A comparison of the workflow of standard Kohn-Sham DFT and NeuralSCF. NeuralSCF models the Kohn-Sham density map using an equivariant graph transformer. NeuralSCF's prediction of the ground-state electron density is defined by its fixed point, solved through self-consistent iterations aided by a density mixing scheme. Finally, ground state can be obtained from the predicted density with an extra Kohn-Sham step. (b) The electron density is represented by the expansion coefficients under a set of atom-centered Gaussian basis functions, which can be decomposed into atom-wise spherical tensors. (c) The two-stage training strategy of NeuralSCF. The explicit pre-training stage learns the Kohn-Sham density map from SCF trajectory data, while the implicit fine-tuning stage further aligns the model's fixed point with the self-consistent electron density via implicit differentiation.
• Figure 2: Network architecture.(a) Network architecture of the proposed NeuralSCF density map $\hat{\mathbf{d}}_\text{out} = f_{\theta}(\mathbf{d}_\text{in};\mathcal{X})$. (b) Network architecture of DensFormer, an end-to-end baseline model $\hat{\mathbf{d}} = g_{\theta}(\mathcal{X})$ that directly predicts self-consistent density from atomic configurations, sharing the same architecture as the NeuralSCF density map except for the input layer. (c) The density encoder transforms input atom-wise density coefficients, atom type, and the local environment into homogeneous input node features $\mathbf{x}_i^{(0)}$. Here, "FFN" stands for feed forward network, "$\oplus$" denotes concatenation, and "$+$" denotes element-wise addition. (d) The density decoder transforms output node features $\mathbf{x}_i^{(N)}$ into output atom-wise density coefficients.
• Figure 3: Results on QM9 and MD datasets.(a-d) Histograms of absolute errors for self-consistent electron density, dipole moment, derived total energy, and derived HOMO-LUMO gap on QM9. The top 5 error outliers for NeuralSCF and its end-to-end density predictor counterpart DensFormer are marked with their respective colors. Vertical dotted lines indicate the mean absolute error (MAE) for each quantity. Property learning SOTAs are directly taken from a recent benchmark liao2023equiformerv2 as a reference. (e-g) Violin plots of absolute errors for self-consistent electron density, derived total energy and forces on ethanol (MD17) and Ac-Ala3-NHMe (MD22), with outliers shown as jittered scatter points. MAEs are indicated for each distribution.
• Figure 4: Zero-shot generalization to a wide range of datasets.(a-c) Violin plots of absolute errors for self-consistent electron density, derived total energy and forces on MD17's ethanol and malonaldehyde. (d-f) Violin plots of absolute errors for self-consistent electron density, derived total energy and dipole moment on BFDb-SSI. (g) Visualization of predicted density errors on the Glu-/Lys+ system, a challenging example from the BFDb-SSI dataset featuring significant charge transfer. (h) The true (dark grey) and predicted torsion energy profiles of 4-dimethylamino-$4^\prime$-nitrostilbene. Shaded regions represent error within chemical accuracy 1kmol.