Kohn-Sham Equations as Regularizer: Building Prior Knowledge into Machine-Learned Physics

Including prior knowledge is important for effective machine learning models in physics and is usually achieved by explicitly adding loss terms or constraints on model architectures. Prior knowledge embedded in the physics computation itself rarely draws attention. We show that solving the Kohn-Sham equations when training neural networks for the exchange-correlation functional provides an implicit regularization that greatly improves generalization. Two separations suffice for learning the entire one-dimensional ${\mathrm{H}}_2$ dissociation curve within chemical accuracy, including the strongly correlated region. Our models also generalize to unseen types of molecules and overcome self-interaction error.

Figure 1

One-dimensional ${\mathrm{H}}_2$ dissociation curves for several ML models trained from two molecules (red diamonds) with optimal models (highlighted in color) selected by the validation molecule at $R=3$ (black triangles). The top panel shows energy (with $E_{\mathrm{NN}}$, the nucleus-nucleus repulsion energy) with exact values shown by the black dashed line. The bottom panel shows the difference from the exact curves with chemical accuracy in gray shadow. (a) directly predicts $E$ from geometries and clearly fails to capture the physics from very limited data. (b)–(d) show our method (KSR) with different inputs to the model to align with the first two rungs of Jacob’s ladder [47] (LDA and GGA) and then global (a fully nonlocal functional). Uniform gas LDA [46] is shown in brown. Gray lines denote 15 sampled functionals during training, with darker lines denoting later samples. Atomic units used throughout.

Figure 2

KS-DFT as a differentiable program. Black arrows are the conventional computation flow. The gradients flow along red dashed arrows to minimize the energy loss $L_E$ and density loss $L_n$. (a) The high-level KS self-consistent calculations with linear density mixing (purple diamonds). (b) A single KS iteration produces $v_{\mathrm{XC},\theta }[n]$ and $E_{\mathrm{XC},\theta }[n]$ by invoking the XC energy calculation twice, once directly and once calculating a derivative using automatic differentiation. (c) The XC energy calculation using the global XC functional.

Figure 3

(a) t-distributed stochastic neighbor embedding (t-SNE) visualization [78] of density trajectories (gray dots) sampled by KSR during training for $R=3.84$ from initial guess (cross) to exact density (red diamond). Darker trajectories denote later optimization steps $t$. Note t-SNE projection does not perfectly preserve the distance between densities. The light red ellipse illustrates the vicinity of the exact density within ${\log }_{10}[\int dx(n_{\mathrm{KS}}-n_{\mathrm{DMRG}}{)}^2/N_e]\le -4.25$. Densities from each KS iteration in trajectories are plotted in the corresponding highlighted colors for (b) $t=0$ untrained, (c) $t=220$ optimal in Fig. 1, and 1 $t=560$ overfitting to training with bad generalization on validation.

Figure 4

Test generalization of models as a function of the total number of training examples $N_{\text{train}}$: full KSR (blue), energy only KSR (pink), and direct ML (orange) on (a) holdout ${\mathrm{H}}_2$ and ${\mathrm{H}}_4$, and unseen types of molecules (b) ${\mathrm{H}}_2^{+}$, (c) ${\mathrm{H}}_2{\mathrm{H}}_2$. Black dashed lines show chemical accuracy. See the Supplemental Material [48] for training details.

Figure 5

Density and KS potential of ${\mathrm{H}}_4$ with $R=4.32$ from neural XC functionals trained with (a) full KSR (blue) and (b) energy only KSR (pink) on training set of size $N_{\text{train}}=20$. Exact curves are in gray. $v_s$ are shifted by a constant for better comparison.