Machine Learning the Physical Non-Local Exchange-Correlation Functional of Density-Functional Theory

TL;DR

This paper introduces a neural network-based exchange-correlation functional for density-functional theory that is highly non-local, computationally efficient, and consistent, aiming to improve accuracy in strongly-correlated systems.

Contribution

It develops a fully non-local, machine-learned functional that reproduces both energy and potential, maintaining efficiency and addressing delocalization issues in DFT.

Findings

Successfully models one-dimensional two-electron systems

Maintains computational efficiency similar to local approximations

Shows potential to improve DFT accuracy in strongly-correlated systems

Abstract

We train a neural network as the universal exchange-correlation functional of density-functional theory that simultaneously reproduces both the exact exchange-correlation energy and potential. This functional is extremely non-local, but retains the computational scaling of traditional local or semi-local approximations. It therefore holds the promise of solving some of the delocalization problems that plague density-functional theory, while maintaining the computational efficiency that characterizes the Kohn-Sham equations. Furthermore, by using automatic differentiation, a capability present in modern machine-learning frameworks, we impose the exact mathematical relation between the exchange-correlation energy and the potential, leading to a fully consistent method. We demonstrate the feasibility of our approach by looking at one-dimensional systems with two strongly-correlated electrons, where density-functional methods are known to fail, and investigate the behavior and performance of our functional by varying the degree of non-locality.

Tables

Table 1. Table 1: Mean absolute errors (MAE) for the total energy in self-consistent calculations for various kernel sizes of our ML-DFT functional, relative to the error of the one-dimensional LDA of Ref.  Helbig  et al. , 2011 . For reference, the mean absolute error of the LDA of Ref.  Helbig  et al. , 2011 is 1.4 × 10 − 2 absent superscript 10 2 \times 10^{-2}  a.u.

Kernel-size	MAE(ML)/MAE(LDA)[%]
LDA	100
1	38.1
15	21.8
30	8.2
60	8.2
120	7.1
180	6.5

Equations

$$E_{\mathrm{GS}}=\sum _i{ϵ}_i+E_{\mathrm{xc}}[n]-\int {\mathrm{d}}^{3}r{v}_{\mathrm{xc}}(\mathbf{r})n(\mathbf{r})-E_{\mathrm{H}}[n],$$

$$v_{\mathrm{xc}}(\mathbf{r})=\frac{\delta E_{\mathrm{xc}}[n]}{\delta n(\mathbf{r})}.$$

$$\frac{1}{r}\to \frac{1}{\sqrt{1+x^2}},$$

$$H(x_1,x_2)$$

$$v(x)=\sum _{k=1}^3\frac{Z_k}{\sqrt{1+(x-a_k{)}^2}}.$$

$$\begin{gathered} L(\theta ,n_i)=\alpha \text{MSE}(E_{\mathrm{xc}})+\beta \text{MSE}(v_{\mathrm{xc}}) \\ +\gamma \text{MSE}\left(\frac{\text{d}{v}_{\mathrm{xc}}(x)}{\text{d}x}\right)+\delta \text{MSE}\left(E_{\mathrm{xc}}-\int \mathrm{d}x{v}_{\mathrm{xc}}(x)n(x)\right). \end{gathered}$$

Full text

Nowadays density functional theory (DFT) is the cornerstone of computational theoretical physics and quantum chemistry, as it provides the prevalent method for the calculation of the electronic structure of both solids and molecules. Based on the Hohenberg-Kohn theorems Hohenberg and Kohn (1964), DFT reformulates the quantum many-electron problem as a theory of the ground-state electronic density $n(\mathbf{r})$. The success of DFT is to a large extent due to the existence of a system of non-interacting electrons (the Kohn-Sham system) that has the same ground-state density as the interacting electrons. This leads to the Kohn-Sham equations, a set of self-consistent equations for one-particle orbitals Kohn and Sham (1965). In such a formalism the ground-state (GS) energy can be expressed as:

[TABLE]

where $\epsilon_{i}$ are the eigenvalues of the Kohn-Sham Hamiltonian, $v_{\rm xc}(\mathbf{r})$ is the exchange-correlation potential, $E_{\rm H}[n]$ is the Hartree energy, and $E_{\rm xc}[n]$ is the exchange-correlation energy. The exchange-correlation potential is defined as the functional derivative of the universal exchange-correlation energy functional:

[TABLE]

Due to the Hohenberg-Kohn existence theorems, if the exact exchange-correlation energy functional $E_{\rm xc}[n]$ is known, Kohn-Sham DFT then yields the exact ground-state energy and the exact ground-state electronic density.

Traditionally, “educated” formal expressions of the exchange-correlation energy functional have been proposed by a combination of theoretical insight, highly accurate Monte-Carlo Ceperley and Alder (1980) or quantum chemical simulations, or by fitting general expressions to experimental data. In general, functionals can be sorted according to Jacob’s ladder Perdew and Schmidt (2001): the lowest rung of the ladder is occupied by local-density approximations (LDA) that use solely single density points as inputs Vosko et al. (1980); Cole and Perdew (1982); Perdew and Wang (1992). The second rung is occupied by generalized-gradient approximations (GGA) that include the gradient of the density Perdew et al. (1996a, b). This is followed by the meta-GGAs Sun et al. (2016) (that use the kinetic-energy density) and hybrid functionals Becke (1993); Perdew et al. (1996c); Adamo and Barone (1999) (that mix a fraction of non-local Fock exchange) on the subsequent rungs. Note that more than 500 of these functionals have been proposed in the past decades Lehtola et al. (2018), although most of them with rather limited impact.

In spite of the success of DFT in dealing efficiently with electronic systems, it still suffers from stubborn quantitative and qualitative failures. For instance, barriers of chemical reactions, band gaps of materials, or molecular dissociation energies are usually underestimated Becke (2014). Degenerate or near-degenerate states are also poorly described by DFT. While hybrid functionals can alleviate some of the problems of traditional semi-local functionals, they come at a greatly increased computational cost that limits severely the number and size of systems that can be researched. It is believed that many of these problems originate in the delocalization and static correlation errors which plague approximate functionals Cohen et al. (2008a, b); Skone et al. (2014). Roughly speaking, the delocalization error refers to the tendency of DFT functionals to spread out the electron density, while the static correlation arises from the difficulty of describing degenerated states with a single Slater determinant Benavides-Riveros et al. (2017).

More recently, machine learning (ML) has revolutionized many fields of computational sciences, such as image or speech recognition Liu and Tian (2010); Waibel and Lee (1990), and has found countless applications in material science Butler et al. (2018); Schmidt et al. (2019); Ghiringhelli et al. (2015). Within DFT, the application of ML techniques to the formulation of density functionals has already a long history Tozer et al. (1996). In 2012, an ML approximation for the kinetic energy functional $T_{s}[n]$ was constructed for a system of noninteracting spinless fermions Snyder et al. (2012, 2013). Yao et al. Yao and Parkhill (2016) developed a convolutional neural network to reproduce the kinetic energy functional for molecules. They already mentioned the possibility of using the functional derivative of the neural network for minimization purposes. In order to exploit the Hohenberg-Kohn density-potential map, an ML model was later trained to learn the fundamental relation of DFT between external potentials and electronic densities Brockherde et al. (2017). These works focused mainly on developing functionals for the total energy or the non-interacting kinetic energy to facilitate orbital-free DFT calculations. More recently, some works have addressed the problem of training the exchange-correlation potential Liu et al. (2017); Nagal et al. (2018); Lundgaard et al. (2016); Wellendorff et al. (2012); Li et al. (2016). However, this line of research has been limited by the fact that the exchange-correlation potential was not obtained from the exchange-correlation energy through the functional derivative of Eq. (2).

It is true that one can find in the literature a series of approximations to the exchange-correlation functionals that do not fullfil Eq. (2). For example, the Krieger-Lee-Iafrate approximation Krieger et al. (1990) breaks this connection in order to simplify the implementations of orbital functionals using the optimized effective potential method Sharp and Horton (1953); Talman and Shadwick (1976). Sometimes, it is also convenient to approximate directly the potential (e.g., in the van Leeuwen-Baerends GGA from 1994 van Leeuwen and Baerends (1994) or the modified Becke-Johnson potential Becke and Johnson (2006); Tran and Blaha (2009)), leading again to expressions that do not obey Eq. (2). These so-called “stray” functionals Gaiduk and Staroverov (2009) have found some important applications. For example, the modified Becke-Johnson is one of the most successful functionals to calculate electronic band-gaps Borlido et al. (2019). Unfortunately, they are also found to break a series of exact theorems and conditions Gaiduk and Staroverov (2009), severely limiting their universality and range of applicability. By and large, it is highly advantageous to develop consistent functionals that obey the important Eq. (2).

Modern ML frameworks, like pytorch Paszke et al. (2017) and tensorflow Abadi et al. (2015), allow for automatic differentiation with respect to any parameter. Recently, Nagai et al. used this functionality to train exchange-correlation potentials for molecules Nagai et al. (2019). They trained neural networks through a Monte-Carlo updating scheme to reproduce accurate energies and densities of molecules. The functionals by Nagai and coauthors follow the traditional approaches of an LDA, GGA, meta-GGA, and add a related near-region approximation. Although a clear step forward, using traditional forms for the exchange-correlation functional is unlikely to lead to fundamentally better, disruptive approximations to the exchange-correlation functionals. New paradigms have to be sought in order to unleash the power of ML techniques to its full extent.

In this paper we use the auto-differentiation functionality to train neural-network exchange-correlation functionals through back propagation. The networks are trained to reproduce not only the correct exchange-correlation energy $E_{\rm xc}$, but also the exchange-correlation potential $v_{\rm xc}(r)$ consistently as its functional derivative with respect to the density. Consequently, the resulting functional allows for self-consistent calculations and can easily be integrated into existing Kohn-Sham DFT frameworks. Furthermore, these functionals can be made highly non-local by using the information of the density in a finite neighborhood as input to the neural network, allowing for far more non-locality than traditional LDA or GGA functionals, despite having the same computational scaling with system size. Therefore, this approach promises to alleviate the delocalization problems of DFT and to improve its accuracy without the computational expense of hybrid functionals. To demonstrate the feasibility of this approach, we developed an ML functional for the exchange-correlation energy and exchange-correlation potential based on exact results for two electrons in one-dimension (1D).

The letter presents the details of the dataset, training process, and neural networks. The exact dependence of the functional on the degree of locality and its behavior is also discussed, as well as our results for the 1D homogeneous electron gas and the H2 molecule along the dissociation path. Finally, we finish the letter discussing our conclusions and future research directions.

Data. The training data was produced by solving exactly the one-dimensional two-electron problem in the external potential generated by up to three different nuclei. Softening the Coulomb interaction,

[TABLE]

we obtain the 1D Hamiltonian driven by the interaction of the two electrons, namely,

[TABLE]

where the external potential is given by the superposition of three potentials,

[TABLE]

The total charge of the nuclei $Z=\sum_{k}Z_{k}$ is equal to 2 or 3. Qualitatively close to real 3D systems, this 1D model is known as a theoretical laboratory for studying strong correlation and developing exchange-correlation density functionals for DFT Wagner et al. (2012). Since the ground-state problem of the Hamiltonian $H(x_{1},x_{2})$ can be treated as a one-particle problem in two dimensions, the problem can be solved exactly.

We sampled 20 000 systems and calculated their exact ground-state energy and ground-state electronic density. We used a grid spacing of 0.1 a.u., and a box size of 20 a.u., leading to a grid with 201 points. The nuclei positions $a_{i}$ in Eq. (4) were normally distributed with zero mean and variance of 4 a.u. We then solved the corresponding inverse Kohn-Sham problem in octopus Andrade et al. (2015) to find the exact exchange-correlation energy and potential. Since the inversion is known to be numerically unstable Jensen and Wasserman (2018), we removed outliers that result from these instabilities. We used up to 12 800 of these systems for training, 6 400 for validation during the training, and 2 000 systems for the test set. Furthermore, training was considerably improved when removing outliers with $E_{\rm xc}>-0.55$ a.u. from the training set. No outliers were removed from the test set to allow for a completely unbiased evaluation of the functionals.

In general, one would have to double the data by mirroring the systems to learn the correct symmetry. However, in this specific case one can simply build the symmetry directly into the neural network functional, as explained in the next subsection.

Topology of the neural network. Our ML functional scans the density of the total system, as illustrated in Fig. 1. The density in a neighborhood of the test point is used as the input for a 4 or 5 layer fully connected neural network, that then outputs a local exchange-correlation energy. Specifically, the network takes a certain number of density points as input, which we call $\kappa$, the kernel size. This is the degree of locality of the ML functional. At the borders of the system the density is padded with $\kappa-1$ zeros. Starting at one border the network calculates the local exchange-correlation energy. In the next step the input of the network is moved by one grid point, and it is evaluated again. As described in Fig. 1, this process continues until the other border is reached. We arrive at the total exchange-correlation energy of the system by summing over all network outputs. The padding and the scanning with a certain kernel-size are inspired by standard convolutional neural networks and can also be implemented as such by concatenating the data along the channel-dimensions in between standard convolutional layers. Due to the homogeneity of space, the functional has to be symmetric with respect to its input densities. The symmetry is ensured by initializing the weights of the first layer symmetrically along the spacial dimension. To arrive at the final scalar output we used four- or five-layer fully connected neural networks.

The possible selection of activation functions (i.e., the non-linearities that follow each multiplication with a weight matrix of a neural network) was rather limited, because typical functions (e.g., rectified linear units) were not usable due to their lacking differentiability at zero (Using relus actually resulted in piece-wise linear potentials). To avoid this problem, we chose exponential linear functions Clevert et al. (2015). Different numbers and sizes of hidden layers were also tested and we settled on the minimum number of parameters that could be used without underfitting. In this work all networks were built on the basis of the ML framework pytorch Paszke et al. (2017). The library Ignite Ign (2018) was used to simplify the training process and tensorboardX to integrate tensorboard Abadi et al. (2015) into pytorch. The network weights were optimized with Adam Kingma and Ba (2014) using default parameters from pytorch.

For the loss function (i.e., the cost function which is going to be optimized in the learning process), we should keep in mind that the objective is not only to obtain small errors for the exchange-correlation energy. To arrive at the correct density through the solution of the self-consistent Kohn-Sham equations, also the exchange-correlation potential should be as close as possible to the exact one. Furthermore, we want not only to ensure a small error for the potential but also its smoothness. In addition, the error of the exchange-correlation energy, as well as the error of the integral $\int\!\!{\rm d}x\>v_{\rm xc}(x)n(x)$ (that appears in the expression for the total energy (1)), should be minimized.

In order to achieve all these goals concurrently we used the following loss function, where $\theta$ are the parameters of the neural network that have to be optimized:

[TABLE]

This function is a weighted combination of the mean squared errors (MSE) of the exchange-correlation energy, the exchange-correlation potential, its numerical spatial derivative and the difference between the exchange-correlation energy and the integral over the potential. This latter term is part of the formula for the total energy (1) and theoretically allows for some error cancellation. We also attempted to use the integral as a separate term in Eq. (5). Depending on the network, one or the other term produced better results. Finally, the weights $\alpha$, $\beta$, $\gamma$ and $\delta$ in Eq. (5) are also optimized as part of the hyperparameter optimization. Usual values for $\alpha$, $\beta$, $\gamma$, $\delta$ are 1.0, 100.0, 10.0, 1.0.

The training for the exchange-correlation energy converges quite fast after a few hundred epochs (i.e., one complete pass of the training data). The convergence of the potential can take thousands of epochs depending on the training set and batch size. At each training step the model was saved if it improved the validation error for the potential. The model with the lowest validation error was later used for the self-consistent Kohn-Sham calculations. As the amount of memory that is needed per sample is quite limited, very large batchsizes (e.g., 4 096) are possible, allowing for a far more efficient parallelization of the training. Training with larger batchsizes seems to produce better convergence. However, it leads to a strong increase in the error during validation with self-consistent Kohn-Sham calculations. Smaller batchsizes (32, 64, 128) improve the error by up to 50% and provide the best generalization ability of the functionals, in consistency with the literature Masters and Luschi (2018); Goodfellow et al. (2016).

Evaluation. We trained neural networks with various kernel sizes and used them within a self-consistent Kohn-Sham calculations for a test set of 2 000 systems created with the method described above. The self-consistent Kohn-Sham calculations were run using a self-written code. For all kernel sizes, models with different hyper-parameters were evaluated on a validation set of 250 systems. The training was not completely converged at this stage. Only for the best models of each kernel size we continued the training and evaluated the models on the test set. To compare various models with respect to the LDA functional of DFT, we chose energy differences, as these are physically more meaningful.

In Fig. 2 we plot the exchange-correlation potential resulting from self-consistent calculations with a ML functional trained according to the loss function (5) for two test-systems and compared it with the exact and LDA predicted exchange-correlation potentials. In Fig. 3 we plot the same information for one test-system; the ML-functional is this time trained only on the exchange-correlation energy (i.e., not on the potential). Whenever the machine-learned functional is trained both on the energy and the potential, the exchange-correlation potential presents a great improvement in comparison to the traditional LDA functional, while the functional trained only on the energy fails completely. Remarkably, the functionals trained with the loss-function (5) also show a qualitatively closer behavior to the exact exchange-correlation potentials. The results for the predicted total energies of the test set (relative to the energy of H2 at its equilibrium distance) are presented in Table 1. Our ML-LDA (i.e., the functional with kernel size $\kappa=1$) already performs better than the traditional LDA. As the ML-LDA was trained to reproduce the exchange-correlation energy of heterogeneous systems while traditional LDAs are “trained” for constant densities the difference in performance is not surprising. Increasing the kernel size leads to a monotonical decrease of the error, and improves the results by more than a factor of six for the larger sizes. The optimal kernel size is, in our opinion, around 30 (i.e., 3 a.u.), as larger kernels do not provide a significant advantage. Furthermore, some of the functionals with larger kernel sizes also demonstrate unphysical behavior (see below).

Ultimately, the more non-local the functional is, the higher the complexity and the larger the number of parameters. This reason, together with the need to represent more long-range interactions that are based on different physical principles (such as van der Waals interaction), makes the training considerably more difficult. One approach to circumvent this problem is to keep the non-locality limited to ranges on the scale of molecular bonds. This allows for simpler training and still includes most of the non-locality that is required for the exchange-correlation energy. Another possibility would be to enlarge the non-locality by increasing the architectural complexity of the functional.

Efforts to decrease the number of training systems for a kernel-size of 30 lead to a slightly increased error of 11% using 800 samples for training. Although the scaling of the networks to realistic three-dimensional systems is non-trivial, we expect that both the number of trainable parameters and the density points in the training set will grow cubically when transitioning to three-dimensional systems. In this sense, we expect a similar demand for training data as in 1D. Furthermore, realistic systems are usually far larger and therefore provide more “local” training samples per system for the neural network. Recent research by Nagai and coauthors points in the same direction Nagai et al. (2019). Indeed, they only required a few sample molecules and used far more parameters to learn a much more local (and in this sense simpler) functional than the ones used here.

Previously, we tested the ML-functional on sample systems belonging to the distribution of the training data. Now, we go a step further and test how our functionals perform self-consistently on systems outside this distribution as well as a couple of paradigmatic cases. The systems in the training set had external potentials arising from 2 and 3 nuclei. In order to go beyond these systems we tested the functional with kernel-size 30 also on a test set of 150 systems with 4 nuclei. Using the same functional as in Table I we arrive at an error for the total energy more than eight times smaller than with the LDA (MAE(ML)/MAE(LDA)=11.9%). Naturally, the error increased outside the training distribution. However, considering the different nature of the highly charged systems with 4 nuclei, this hints at a good generalization ability of the functional.

In Fig. 4 we present our calculations for the H2 molecule in 1D along the dissociation path with functionals of varying non-locality in comparison to the exact result. The curves in Fig. 4 are shifted to have the same equilibrium energies. As is well-known, the traditional LDA completely fails to produce the correct dissociation limit Helbig et al. (2011). The same behaviour is observed for the ML LDA (with a kernel size of 1). Remarkably, using increasingly more non-local functionals, we can reduce the relative energy error to $3.2\%$ of the LDA error. It is obvious that even the functional with a kernel size of 30 will start failing above a certain distance. This is a conceptual problem of local KS-DFT and can only be alleviated and not eliminated in our approach. It can already be considered a success that our functionals are able to reproduce the dissociation curve reasonably well far beyond their own degree of non-locality. Yet it has to be noted that not all functionals performed that well. Some of the functionals with larger kernels failed to reproduce a physical behavior with respect to the dissociation distance and produced multiple local minima and maxima. Despite these problems, they still return the correct equilibrium distance and on average far better energies than the LDA. This unphysical behavior of some functionals just stresses the fact that a rigorous validation on a multitude of different systems will be essential to arrive at a working functional. It has to be noted that a larger training set and a longer training time was far more beneficial for this validation than for example the average error. As systems similar to the dissociated molecule are most likely outliers of the training data this is not surprising. Note that the need for more training data can however be avoided by active learning and a thoughtful construction of the training set.

Finally, we study the homogeneous-electron gas, a model system that is used in the construction of the majority of exchange-correlation functionals. We can simulate this system with our neural networks by providing them with a constant electronic density as input. The results can then be compared to the numerically exact values for the energy density as obtained, for example, from quantum Monte-Carlo simulations Helbig et al. (2011). Our results are depicted in Fig. 5 as a function of the Wigner-Seitz radius $r_{s}=1/2n$. Some difficulties for our neural-network functionals are evident. First, the functionals were trained only for systems with a specific size while the homogeneous electron gas is, in fact, an infinite periodic system. Second, as the histogram in Fig. 5 illustrates, the training data does not contain almost any samples with high densities ($r_{s}<1$). Naturally, the availability of training data similar to the homogeneous electron gas is even more important for the more non-local functionals as they take into account larger regions of space. The first challenge results in the fact that the non-zero biases in each layer cause the neural networks to output a non-zero value for zero density. When training for different system sizes there are several ways to avoid this failure. First one could solve the problem by adding systems padded with different amounts of zeros at the border to force the neural network to learn the correct relationship. As a second possibility, one could force all biases of all layers to zero, however, this would severely limit the expressibility of the networks. To circumvent this problem, and in order to compare the behavior of the energy with respect to the Wigner-Seitz radius, we shifted the curves in Fig. 5 to yield zero energy for zero density. Despite the small amount of training data at high density, functionals with larger kernel sizes still generalize on average far better to the homogeneous electron gas. While this is the case for most models, there are some rare cases, similar to the problems with the H2 dissociation, where large kernel-sizes produce unphysical behavior.

It is not obvious whether this result will remain true in three dimensions, it is nevertheless promising that the extra non-local information in the larger kernels might help the functionals to be generalized. Constant densities will be an essential feature of a functional for solid state physics. Fortunately, exact training data in the form of quantum Monte-Carlo calculations already exists for this purpose, and can be easily incorporated in our training sets.

Conclusions. In this letter we have demonstrated the viability of learning an exchange-correlation potential, via the differentiation of the exchange-correlation energy in a physically consistent manner. This procedure allows for standard self-consistent Kohn-Sham calculations. From the presented data, it is evident that neural-network functionals trained on the exchange-correlation potential and energy have the potential to be far more precise than previous local DFT functionals. Increasing the non-locality of the functional allows for an extremely precise treatment of the electronic interaction on the scale of at least a few atomic units and, to a certain extent, even solve long-standing problems in DFT like e.g. molecular dissociation.

For simplicity, we trained a neural network to the one-dimensional two-electron problem in the external potential generated by up to to three nuclei. Training three-dimensional systems will have to be accomplished by using data obtained with coupled-cluster, full configuration-interaction, or quantum Monte Carlo. While sufficient data to train a universal functional still has to be created, exchange-correlation energies and potentials for a few small molecules already exists and can provide a good starting point. The density representation on a grid is unfortunately not feasible for more general systems, as grid sizes and forms will vary. However, we think this can easily be circumvented by representing the density locally in some basis sets (e.g., Gaussians).

Finally, there is already a long history within DFT in the development of empirical functionals Wellendorff et al. (2012); Lundgaard et al. (2016); Yu et al. (2016a, b); Mardirossian and Head-Gordon (2016). The machine learning paradigm allows us to drastically increase the amount of data used for the training and the complexity of these functionals. Including known exact conditions of the exchange-correlation functional in the learning process as constraints in the minimization will still be helpful Hollingsworth et al. (2018) and provide further conditions for validation. Furthermore, as the functionals will have to work in practically every density environment, the importance of an extremely in-depth validation cannot be overstated and will be essential to arrive at a widely used functional.

Acknowledgments. We acknowledge partial support from the German DFG through project MA-6786/1.