Source-linked AI summary
Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics
Li Li, Stephan Hoyer, Ryan Pederson, Ruoxi Sun, Ekin D. Cubuk, Patrick Riley, Kieron Burke
TL;DR
Existing learned XC approximations have limited small-data generalization, while prior work did not achieve chemical accuracy or work in the dissociation limit. The method backpropagates through Kohn-Sham self-consistent iterations while training a neural exchange-correlation functional. KSR fits H2 and H4 well in the stretched limit and performs well on unseen H2H2 at R < 3 and R > 8, with slight overbinding around R = 5.
Problem
Existing learned XC approximations have limited small-data generalization, while prior work did not achieve chemical accuracy or work in the dissociation limit.
Method
The method backpropagates through Kohn-Sham self-consistent iterations while training a neural exchange-correlation functional.
Results
KSR fits H2 and H4 well in the stretched limit and performs well on unseen H2H2 at R < 3 and R > 8, with slight overbinding around R = 5.
Takeaways & Limitations
Kohn-Sham calculations in training, rather than only inference, are crucial to the reported generalization in the small-training-set regime.
Abstract
from arXiv · showhide
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 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.
Supplemental Material Building prior knowledge into machine-learned physics:
The supplemental material defines the one-dimensional model, XC functional, self-consistent calculations, automatic differentiation, and symmetry treatment used in the study.
- I. 1D MODEL SYSTEMS: The one-dimensional model uses exponential electron–electron interactions and solves the interacting Hamiltonian with DMRG for accurate ground-state energies and densities.The grid contains 513 points with spacing 0.08, and the DMRG energy convergence threshold is 10^-7 Ha.
- B. Initial density: Kohn-Sham calculations initialize the density from the non-interacting system and use linear density mixing to improve self-consistent convergence.The mixing factor follows α = 0.5 × 0.9^(k−1).
- C. XC functional: The XC functional maps a density on the grid to an XC energy-density array, and it may be conventional, such as LDA, or neural.The XC energy is obtained by discretizing the density-weighted integral over the grid.
- D. XC potential from automatic differentiation: Automatic differentiation computes the XC potential as the functional derivative of the XC energy with respect to density.The implementation uses JAX’s gradient of the discretized XC-energy function.
- E. Symmetry: Symmetry flips a grid function about the center and averages it with itself, stabilizing gradient information for stretched systems with vanishing KS gaps.The restriction does not change performance near equilibrium but supports convergence in stretched H2 and H4.
IV. TRAINING, VALIDATION AND TEST
Training uses exponentially weighted energy losses over KS iterations and fixes the iteration count separately for each molecular type based on LDA convergence estimates.
- IV. TRAINING, VALIDATION AND TEST: The energy-loss weights are defined as wk = 0.9^(K−k)H(k−10), emphasizing later iterations while excluding the first ten.K is the total number of iterations and H is the Heaviside function.
- IV. TRAINING, VALIDATION AND TEST: The KSR training runs use fixed iteration counts at least as large as LDA estimates for each molecule type, providing sufficient iterations for convergence.The LDA estimates are approximately 8, 25, 5, and 6 iterations for H2, H4, H2+, and H2H2, respectively.
C. Dataset for learning H2 dissociation from two molecules
The H2 dissociation experiment trains on exact data from one compressed and one stretched molecule, with an intermediate molecule used for validation.
- C. Dataset for learning H2 dissociation from two molecules: The H2 dissociation curves are trained from exact densities and energies at separations R = 1.28 and R = 3.84.The optimal checkpoint is selected using a validation molecule at R = 2.96.
1. Training molecules
The training molecules used for Figure 4 are specified by their nearby-atom distances in the supplemental dataset description.
- 1. Training molecules: The distances between nearby atoms for the Figure 4 training molecules are listed in Table II.The passage identifies the table as the source of the training-molecule geometries but does not provide the distance values.
2. Validation molecules
The validation molecules used in Figure 4 are specified by nearby-atom distances, with a fixed validation set for 4 ≤ Ntrain ≤ 20.
- Training molecules for Figure 4 are likewise characterized by distances between nearby atoms.
- The validation set for Figure 4 uses fixed molecular geometries specified by nearby-atom distances.The validation set remains fixed across calculations with 4 ≤ Ntrain ≤ 20.
3. Test molecules
The test molecules used in Figure 4 are specified by distances between nearby atoms, with a fixed test set across training-set sizes.
- The test set for Figure 4 is fixed for calculations with 4 ≤ Ntrain ≤ 20.Its molecular geometries are listed through distances between nearby atoms.
4. Test errors
The test-error analysis extends results across training-set sizes and reports dissociation-curve performance for molecules including unseen H2H2.
- Figure S1 extends test-error results to Ntrain = 20, with numerical values listed in Table V.
- The reported dissociation curves compare exact curves with uniform-gas LDA and KSR-global results.
- KSR fits H2 and H4 well even in the stretched limit and performs well on unseen H2H2 at small and large distances.For H2H2, the model slightly overbinds around R = 5.
- KSR perfectly predicts H+2 because of a self-interaction gate in the neural XC functional.
A. Architecture
The architectures combine neural XC models with Kohn-Sham computation, while local and semilocal variants restrict convolutional receptive fields to preserve their intended approximations.
- The direct ML model maps a density representation of molecular geometry through convolutional and dense layers to a scalar energy E.The dense layer is used only in the direct ML model.
- KSR-LDA uses a filter size of 1, while KSR-GGA uses three nearby density points to compute the center-point XC energy density.These receptive fields implement local and semilocal approximations, respectively.
- The KSR architectures differ from KSR-global by excluding global convolution for LDA and GGA, while retaining the remaining network structure except for the self-interaction gate.
- Global convolution uses multiple channels to capture interactions at different scales through trainable interaction-scale parameters.The 16-channel global convolution preserves the input density in one channel and bounds the remaining scales.
- The architecture comparison includes direct energy prediction, neural LDA and GGA with KSR, and a neural global functional.
- Training uses repeated random seeds, L-BFGS optimization, and validation checkpoints selected by average energy error per electron.Parameter checkpoints are saved every 10 steps until L-BFGS stops.
VI. TRAINING A NEURAL XC FUNCTIONAL WITHOUT KS REGULARIZATION
This section tests whether generalization comes from using Kohn-Sham calculations during inference rather than training. Comparing single-step and Kohn-Sham-regularized training shows that training through the Kohn-Sham iterations is crucial in the small-data regime.
- VI. TRAINING A NEURAL XC FUNCTIONAL WITHOUT KS REGULARIZATION: Single-step training fails to generalize in the small-training-set limit even when Kohn-Sham self-consistent calculations are used during inference.Figure S4 directly compares single-step and Kohn-Sham-regularized training under the same inference-stage calculation.
- VI. TRAINING A NEURAL XC FUNCTIONAL WITHOUT KS REGULARIZATION: The comparison is designed to isolate training-stage Kohn-Sham calculations as the crucial source of the observed generalization.The authors distinguish this effect from simply using self-consistent calculations at inference.
- VI. TRAINING A NEURAL XC FUNCTIONAL WITHOUT KS REGULARIZATION: A single-step model can work with a larger training set and exact vXC, so its failure is specific to the small-data comparison.The authors frame this as a limitation of the tested regime rather than a universal failure of single-step models.
VII. TRAINING A NEURAL XC FUNCTIONAL WITH “WEAKER” KOHN-SHAM REGULARIZATION
This section weakens Kohn-Sham regularization by restricting gradient flow while keeping the forward self-consistent calculation unchanged. The resulting experiments compare how much iteration history contributes to learning the dissociation curve.
- VII. TRAINING A NEURAL XC FUNCTIONAL WITH “WEAKER” KOHN-SHAM REGULARIZATION: Stopping gradients before iteration k* controls Kohn-Sham regularization strength without changing the forward Kohn-Sham calculations.Earlier iterations lose gradient access, while the number of self-consistent iterations K remains fixed.
- VII. TRAINING A NEURAL XC FUNCTIONAL WITH “WEAKER” KOHN-SHAM REGULARIZATION: The three settings expose gradients through the final iteration, the last two iterations, or all iterations.The no-stop-gradient graph matches the main-text computation, while the other settings progressively restrict training information.
- VII. TRAINING A NEURAL XC FUNCTIONAL WITH “WEAKER” KOHN-SHAM REGULARIZATION: The H2 dissociation experiments compare curves trained under these three gradient-flow settings, including the computational graph used in the main text.Figure S6 reports the corresponding dissociation curves and Figure S5 defines the graph variants.
- VII. TRAINING A NEURAL XC FUNCTIONAL WITH “WEAKER” KOHN-SHAM REGULARIZATION: The true dissociation curve is captured in the stop-gradient experiment.The supplied result passage states that the true curve is captured after describing the weaker-regularization comparison.