Learning the exchange-correlation functional fromnature with differentiable density functional theory

Muhammad F. Kasim1 , & Sam M. Vinko1,2
Department of Physics, University of Oxford, Oxford, OX1 3PU, UK
2
Central Laser Facility, STFC Rutherford Appleton Laboratory, Didcot OX11 0QX, UK{firstname.lastname}@physics.ox.ac.uk
1

Abstract
Improving the predictive capability of molecular properties in ab initio simulations
is essential for advanced material discovery. Despite recent progress making
use of machine learning, utilizing deep neural networks to improve quantum
chemistry modelling remains severely limited by the scarcity and heterogeneityof appropriate experimental data. Here we show how training a neural network
to replace the exchange-correlation functional within a fully-differentiable threedimensional Kohn-Sham density functional theory (DFT) simulation can greatly
improve its accuracy and generalizability. Using only eight experimental data
points on diatomic molecules, our trained exchange-correlation networks provided
improved predictions of atomization energies across a collection of 104 molecules
containing new bonds and atoms that are not present in the training.

1

Introduction

Fast and accurate predictive models of atomic and molecular properties are keys in advancing novel
material discovery. With the rapid development of machine learning (ML) techniques, considerable
effort has been dedicated to accelerating simulations, enabling the efficient exploration of configuration space. Much of this work focuses on end-to-end ML approaches [1], where the models are trained
to predict some specific molecular or material properties using tens to hundreds of thousands of data
points, generated via simulations [2]. Once trained, ML models have the advantage of accurately
predicting the simulation outputs much faster than the original simulations [3]. The level of accuracy
achieved by ML models with respect to the simulated data can be very high, and may even reachchemical accuracy on select systems outside the original training dataset [4]. These machine learning
approaches show great promise, but rely heavily on the availability of large datasets for training.
While there has been a concerted effort in accelerating simulations, less attention has been dedicatedto using machine learning to increase the accuracy of the simulation itself with respect to theexperimental data. In particular, as ML models typically rely on some underlying simulation for
the training data, the accuracy of the trained models can never exceed the simulations in terms ofaccuracy. However, training ML models using experimental data is challenging as the available datais highly heterogeneous as well as limited, making it difficult to train accurate ML models tailored toreproduce a specific system or physical property. This limitation also poses a great challenge on thegeneralization capability of ML models to systems outside the training data.
One emerging solution to this problems is to use machine learning to build models which learn theexchange-correlation (xc) functional within the framework of Density Functional Theory (DFT)
simulations [5, 6, 7, 8, 9, 10]. Here, the Kohn-Sham DFT (KS-DFT) [11] can be used to calculate
various ground-state properties of molecules and materials, where its accuracy largely depends on
the quality of the employed xc functional. Works along these lines in the one-dimensional case35th Conference on Neural Information Processing Systems (NeurIPS 2021), Sydney, Australia.

Figure 1: Schematics of the xc neural network and differentiable KS-DFT code. The trainable neuralnetwork is illustrated with dark gray boxes. The solid dark lines show the forward calculation flow
while the dashed blue lines show the gradient backpropagation flow.
show promise [9, 10], but are severely limited by the lack of availability of experimental data fromone-dimensional systems.
In this context, we present here a new machine learning approach to learn the xc functional using
heterogeneous and limited experimental data. We achieve this in a hybrid of physics and machinelearning way by incorporating a deep neural network representing the xc functional in a fully
differentiable three-dimensional KS-DFT calculation. Notably, using a xc neural network (XCNN) ina fully differentiable KS-DFT calculation gives us a framework to train the xc functional using any
available experimental properties, however heterogeneous, which can be simulated via a standard
KS-DFT. Moreover, as the xc neural network is independent of the system (atom type, structure,
number of electrons, etc.), a well-trained XCNN should be generalizable to a wide range of systems,
including those outside the pool of training data. In this paper, we demonstrate this approach using adataset comprising atoms and molecules from the first few rows of the periodic table, but the results
are readily extendable to more complex structures.

2

Differentiable DFT

A single biggest challenge to realizing our method is writing a KS-DFT simulation in a fully
differentiable manner. This is achieved here by writing the KS-DFT code using an automaticdifferentiation library, PyTorch [12], and by providing separately the gradients of all the KS-DFT
components which are not available in PyTorch. The KS-DFT calculation involves self-consistent
iterations of solving the Kohn-Sham equations,
H[n](r)φi (r) = εi φi (r),Xn(r) =
fi |φi (r)|2 ,

(1)
(2)

i

where n(r) is the electron density as a function of a 3D coordinate (r), fi is the occupation numberof the i-th orbital (φi ), εi is the Kohn-Sham orbital energy, and H[n](r) is the Hamiltonian operator,
a functional of electron density profile n and a function of the 3D coordinate. The Hamiltonian
operator consists of kinetic, external potential, Hartree potential, and xc potential operators.
The KS-DFT calculation proceeds by solving Eq. (1) with eigendecompositions and computing
Eq. (2) iteratively until convergence or self-consistency is reached. In this work, the Hamiltonian
operator H[n] is constructed using the contracted Gaussian-type basis [13] as the orbital basis. The
schematics of our differentiable KS-DFT code is shown in Fig. 1. Unless specified otherwise, allcalculations involved in this paper use 6-311++G(3df,3pd) basis sets [14]. More details on solving
the KS equations can be found in the appendix.
Although modern automatic differentiation libraries provide many operations that enable an automatic
propagation of gradient calculations, there are many gradient operations that have to be implemented
specifically for this project. An important example is the self-consistent cycle iteration. Previous
works on differentiable quantum chemistry [15, 10] employed a fixed number of linear mixing
iterations to achieve self-consistency, and let the automatic differentiation library propagate thegradient through the chain of iterations. In contrast, we implement the gradient propagation of the2

self-consistency cycle using an implicit gradient calculation (see appendix) [16, 17], allowing us to
deploy a better self-consistent algorithm, achieving faster and more robust convergence than linear
mixing.
Another example is the gradient calculation of the eigendecomposition with (near-)degenerate
eigenvalues which often produces numerical instabilities in the gradient calculation. To solve thisproblem, we follow the calculation by Kasim [18] to obtain a numerically stable gradient calculationin the (near-)degenerate cases. Besides those two examples, there are more KS-DFT parts that require
specific gradient implementations. The details of those parts can be found in the appendix.
With the fully differentiable KS-DFT code, the xc energy represented by a deep neural network can betrained efficiently. For this work we will test two different types of xc deep neural networks (XCNN)
based on two rungs of Jacob’s ladder – the Local Density Approximation (LDA) and the Generalized
Gradient Approximation (GGA). These xc energies can be written as:
Z
EnnLDA [n] = αELDA(xc) [n] + β n(r)f (n, ξ) dr(3)
Z
EnnPBE [n] = αEPBE(xc) [n] + β n(r)f (n, ξ, s) dr,(4)
where ELDA(xc) is the LDA exchange [11] and correlation energy [19], and EPBE(xc) is the PBE
exchange-correlation energy [20]. Here α and β are two trainable parameters starting from 1 and
0 respectively, and f (·) is the trainable neural network. The LDA and PBE energies are calculatedusing Libxc [21], wrapped with PyTorch to provide the required gradient calculations. The XCNNs
take inputs of the local electron density n, the relative spin polarization ξ = (n↑ − n↓ )/n, and thenormalized density gradient s = |∇n|2 /(24π 2 n4 )1/3 .
For the XCNN training, we compiled a small dataset consisting of experimental atomization energies(AE) of diatomic molecules from the NIST CCCBDB database [22] and calculated density profiles
using a CCSD calculation [23] as a regularizer [24]. The dataset is split into 2 groups: training andvalidation. The training dataset is used directly to update the parameters of the neural network via
the gradient of a loss function (see appendix for details on the loss function). After the training isfinished, we selected a model checkpoint during the training that gives the lowest loss function for
the validation dataset.
The complete list of atoms and molecules used in the training and validation datasets is shown in
Table 1. We note that some atoms are only present in the validation set and not in the training set, toobtain checkpoints that are able to generalize well to new types of atoms outside the training set.

3

Results

To test the performance of the trained XCNNs, we prepared a test dataset consisting of the atomizationenergies of 104 molecules from ref [25] and the ionization potentials of atoms H-Ar from NIST
ASD [27]. The experimental geometric data as well as the atomization energies for the 104 molecules
were obtained from the NIST CCCBDB database. There are actually 146 molecules in ref [25],however, only 104 molecules have experimental data in CCCBDB, so we limit our dataset to those.
The atomization energy test dataset contains molecules with 2-11 atoms ranging from H-Cl, containing
many bonds and atoms that do not present in the training and validation datasets. The molecules inthe test dataset are listed in the appendix.
The results of training the XCNNs using the dataset are presented in Table 2. We see that in thiscase both XCNN-LDA and XCNN-PBE provide significant improvement over their bases, i.e. LDA
and PBE, respectively. In terms of average atomization energy prediction errors across the 104 test
molecules, XCNN-LDA achieves more than 4 times lower error than LDA (PW92) and XCNN-PBE
Table 1: Atoms and molecules presented in the training and validation datasets.
Type
Training
Validation

Atomization energy
H2 , LiH, O2 , CO
N2 , NO, F2 , HF

Density profile
H, Li, Ne, H2 , Li2 , LiH, B2 , O2 , CO
He, Be, N, N2 , F2 , HF

3

Table 2: Mean absolute error (MAE) in kcal/mol of the atomization energy and ionization potential
for atoms and molecules in the test dataset. The column “IP” represents the deviation in ionization
potential for atoms H-Ar. Column “AE” is the MAE of atomization energy of a collection of 104
molecules from ref [25]. The following 4 columns present the atomization energies of subsets ofmolecules from the full 104-molecule set: “AE HC” for hydrocarbons, “AE SHC” for substitutedhydrocarbons, “AE NHC-1” for non-hydrocarbon molecules containing only first and second row
atoms, and “AE NHC-2” for non-hydrocarbon molecules containing at least one third row atom. The
bolded values are the best MAE in the respective column for each group of approximations.
Calculation
IP
AE AE HC AE SHC AE NHC-1 AE NHC-2
Local Density Approximations (LDA)
LDA [19]
6.9 70.4
97.0
101.1
58.8
43.7
22.7
19.4
7.8
16.6
XCNN-LDA
50.8 15.3
XCNN-LDA-IP
15.2 18.5
25.4
21.8
8.4
22.9
Generalized gradient approximations (GGA)
15.1
23.2
18.0
9.8
PBE [20]
3.6 16.5
XCNN-PBE
10.7
7.4
5.4
8.4
6.5
8.5
XCNN-PBE-IP
4.1
8.1
6.8
9.6
6.7
8.6
Other approximations
3.8
8.2
4.0
4.6
SCAN [26]3.7
5.2
CCSD (basis: cc-pvqz)
2.0 12.1
11.1
17.7
11.2
9.0
CCSD-T (basis: cc-pvqz)
1.3
3.5
2.2
5.6
2.5
3.5

achieves more than 2 times lower error than PBE. Importantly, although the training and validation
sets contain atomization energies only for 8 diatomic molecules, both XCNNs provide improvements
on atomization energy predictions of the 104-molecule test dataset, which includes molecules withup to 11 atoms.
By looking into the subsets of the atomization energy dataset, both XCNN-LDA and XCNN-PBE
also achieve better atomization energy predictions than their bases on all of the subsets. The
substantial improvements seen on the hydrocarbon and substituted hydrocarbon molecules show thegeneralization capability of the technique to molecules with bonds not present in neither the training
nor the validation dataset (e.g., C–H, C–C, C=C). Moreover, both XCNNs can also predict better
atomization energy on the “AE NHC-2” subset containing third-row atoms (i.e. Na-Ar) that are notpresent in the training and validation datasets, showing great generalization capability of the method.
Despite these encouraging results on the atomization energies, both XCNNs lead to significantly
worse predictions of atomic ionization potentials compared with their bases. Therefore, for the
next experiment we trained the XCNN by adding ionization potential data into the training andvalidation datasets, as shown in Tab. 1. The xc neural networks trained with this additional data arecalled “XCNN-LDA-IP” and “XCNN-PBE-IP”. Adding ionization potential data to the training andvalidation datasets improves the performance of XCNNs on the ionization potential at the cost ofslightly increasing the error on the atomization energy predictions. Although the ionization potential
predictions of XCNN-IPs do not have lower error compared to their bases, they produce considerably
better predictions on the atomization energies of the tested molecules than their bases.
Although the trained XCNNs cannot make better predictions than higher level of approximations, e.g.
SCAN [26] and CCSD(T), the presented method shows that significant improvement can be madewithin the same level of approximation. The level of accuracy can be improved in the future workby using a higher level of approximation, such as moving to the higher rungs of Jacob’s ladder orincorporating more global density information with neural networks in calculating the xc energy.

4Conclusion

We presented a novel approach to train machine-learned xc functionals within the framework offully-differential Kohn-Sham DFT using experimental data. Our xc neural networks, trained on a
few atoms and diatomic molecules, are able to improve the prediction of atomization and ionizationenergies across a set of 104 molecules not present in the training dataset, including larger molecules
containing new types of bonds and atoms.
4

Broader impact
Computational quantum chemistry has been an active research field since decades ago with large
potential in chemistry and material sciences. This paper presents a new way to solve the accuracy
problem in computational quantum chemistry that potentially have a significant impact both positively
and negatively. More accurate quantum chemistry calculation can accelerate materials discovery andwhether it brings positive or negative impact, depends on the use of the materials. For example, if
the computational quantum chemistry is used to find drugs for curing rare diseases, it would bringpositive impact to the society. However, if it is used to find an explosive materials for weapons, itwould be perceived negatively by most people.

Acknowledgments and Disclosure of Funding
We would like to thank Kyle Bystrom for a useful suggestion on compiling experimental data. M.F.K.
and S.M.V. acknowledge support from the UK EPSRC grant EP/P015794/1 and the Royal Society.
S.M.V. is a Royal Society University Research Fellow. The authors declare no conflict of interest.

References[1] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message
passing for quantum chemistry. In International Conference on Machine Learning, pages 1263–1272.
PMLR, 2017.
[2] Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet–adeep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722,
2018.
[3] Justin S Smith, Olexandr Isayev, and Adrian E Roitberg. Ani-1: an extensible neural network potential
with dft accuracy at force field computational cost. Chemical science, 8(4):3192–3203, 2017.
[4] Felix A Faber, Luke Hutchison, Bing Huang, Justin Gilmer, Samuel S Schoenholz, George E Dahl, Oriol
Vinyals, Steven Kearnes, Patrick F Riley, and O Anatole Von Lilienfeld. Prediction errors of molecularmachine learning models lower than hybrid dft error. Journal of chemical theory and computation,
13(11):5255–5264, 2017.
[5] Ryo Nagai, Ryosuke Akashi, and Osamu Sugino. Completing density functional theory by machinelearning hidden messages from molecules. npj Computational Materials, 6(1):1–8, 2020.
[6] John C Snyder, Matthias Rupp, Katja Hansen, Leo Blooston, Klaus-Robert Müller, and Kieron Burke.
Orbital-free bond breaking via machine learning. The Journal of chemical physics, 139(22):224104, 2013.
[7] Jonathan Schmidt, Carlos L Benavides-Riveros, and Miguel AL Marques. Machine learning the physicalnonlocal exchange–correlation functional of density-functional theory. The journal of physical chemistry
letters, 10(20):6425–6431, 2019.
[8] Rodrigo A Vargas-Hernández. Bayesian optimization for calibrating and selecting hybrid-density functionalmodels. The Journal of Physical Chemistry A, 124(20):4053–4061, 2020.
[9] John C Snyder, Matthias Rupp, Katja Hansen, Klaus-Robert Müller, and Kieron Burke. Finding density
functionals with machine learning. Physical review letters, 108(25):253002, 2012.
[10] Li Li, Stephan Hoyer, Ryan Pederson, Ruoxi Sun, Ekin D Cubuk, Patrick Riley, Kieron Burke, et al.
Kohn-sham equations as regularizer: Building prior knowledge into machine-learned physics. PhysicalReview Letters, 126(3):036401, 2021.
[11] Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects.
Physical review, 140(4A):A1133, 1965.
[12] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor
Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang,
Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie
Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In
H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in
Neural Information Processing Systems, volume 32, pages 8026–8037. Curran Associates, Inc., 2019.

5

[13] Benjamin P Pritchard, Doaa Altarawy, Brett Didier, Tara D Gibson, and Theresa L Windus. New basis
set exchange: An open, up-to-date resource for the molecular sciences community. Journal of chemical
information and modeling, 59(11):4814–4820, 2019.
[14] Timothy Clark, Jayaraman Chandrasekhar, Günther W Spitznagel, and Paul Von Ragué Schleyer. Efficientdiffuse function-augmented basis sets for anion calculations. iii. the 3-21+ g basis set for first-row elements,
li–f. Journal of Computational Chemistry, 4(3):294–301, 1983.
[15] Teresa Tamayo-Mendoza, Christoph Kreisbeck, Roland Lindh, and Alán Aspuru-Guzik. Automaticdifferentiation in quantum chemistry with applications to fully variational hartree–fock. ACS central
science, 4(5):559–566, 2018.
[16] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. arXiv preprint arXiv:1909.01377,
2019.
[17] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López,Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint
arXiv:2105.15183, 2021.
[18] Muhammad Firmansyah Kasim. Derivatives of partial eigendecomposition of a real symmetric matrix fordegenerate cases. arXiv preprint arXiv:2011.04366, 2020.
[19] John P Perdew and Yue Wang. Accurate and simple analytic representation of the electron-gas correlationenergy. Physical review B, 45(23):13244, 1992.
[20] John P Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made simple.
Physical review letters, 77(18):3865, 1996.
[21] Miguel AL Marques, Micael JT Oliveira, and Tobias Burnus. Libxc: A library of exchange and correlation
functionals for density functional theory. Computer physics communications, 183(10):2272–2281, 2012.
[22] Russel D Johnson III and NIST CCCBDB Team. NIST Computational Chemistry Comparison and Benchmark Database (Release 21), [Online]. Available: https://dx.doi.org/10.18434/T47C7Z[2020,
August]. National Institute of Standards and Technology, Gaithersburg, MD., 2020.
[23] Jiří Čížek. On the correlation problem in atomic and molecular systems. calculation of wavefunction
components in ursell-type expansion using quantum-field theoretical methods. The Journal of Chemical
Physics, 45(11):4256–4266, 1966.
[24] Michael G Medvedev, Ivan S Bushmarinov, Jianwei Sun, John P Perdew, and Konstantin A Lyssenko.
Density functional theory is straying from the path toward the exact functional. Science, 355(6320):49–52,
2017.
[25] Larry A Curtiss, Krishnan Raghavachari, Paul C Redfern, and John A Pople. Assessment of gaussian-2
and density functional theories for the computation of enthalpies of formation. The Journal of Chemical
Physics, 106(3):1063–1079, 1997.
[26] Jianwei Sun, Adrienn Ruzsinszky, and John P Perdew. Strongly constrained and appropriately normed
semilocal density functional. Physical review letters, 115(3):036402, 2015.
[27] A. Kramida, Yu. Ralchenko, J. Reader, and and NIST ASD Team. NIST Atomic Spectra Database (ver. 5.8),[Online]. Available: https://dx.doi.org/10.18434/T4W30F[2020, October]. National Instituteof Standards and Technology, Gaithersburg, MD., 2020.
[28] Clemens Carel Johannes Roothaan. New developments in molecular orbital theory. Reviews of modern
physics, 23(2):69, 1951.
[29] Qiming Sun. Libcint: An efficient general integral library for g aussian basis functions. Journal ofcomputational chemistry, 36(22):1664–1671, 2015.
[30] Saswata Dasgupta and John M Herbert. Standard grids for high-precision integration of modern density
functionals: Sg-2 and sg-3. Journal of computational chemistry, 38(12):869–882, 2017.
[31] Axel D Becke. A multicenter numerical integration scheme for polyatomic molecules. The Journal ofchemical physics, 88(4):2547–2553, 1988.
[32] Henk A Van der Vorst. Bi-cgstab: A fast and smoothly converging variant of bi-cg for the solution of
nonsymmetric linear systems. SIAM Journal on scientific and Statistical Computing, 13(2):631–644, 1992.

6

[33] Qiming Sun, Timothy C Berkelbach, Nick S Blunt, George H Booth, Sheng Guo, Zhendong Li, Junzi Liu,James D McClain, Elvira R Sayfutyarova, Sandeep Sharma, et al. Pyscf: the python-based simulations ofchemistry framework. Wiley Interdisciplinary Reviews: Computational Molecular Science, 8(1):e1340,
2018.
[34] Richard M Martin. Electronic structure: basic theory and practical methods. Cambridge university press,
2020.
[35] Karl K Irikura. Experimental vibrational zero-point energies: Diatomic molecules. Journal of physicaland chemical reference data, 36(2):389–397, 2007.
[36] Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Icml,
2010.
[37] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han.
On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv:1908.03265, 2019.

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes] See the last paragraph in section
3.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See“Broader impact” section.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] See Appendix
A.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] For data splits, see Table 1 and Appendix E. For hyperparameters,
see see Appendix D3.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes] See Appendix D3.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A](b) Did you mention the license of the assets? [N/A](c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
See Appendix A.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]7

A

Code availability

The differentiable quantum chemistry (DQC) code used for this paper can be found in
https://github.com/diffqc/dqc/ with Apache-2.0 license. The DQC code uses PyTorch as well as xitorch, an additional differentiable scientific library written specifically for this project. The repository
for xitorch can be found in https://github.com/xitorch/xitorch/ with MIT license.

B

KS-DFT Calculation

The Hamiltonian operator in Eq. (1) is given in atomic units as H[n](r) = −∇2 /2 + v(r) +vH [n](r) + vxc [n](r) with −∇2 /2 representing the kinetic operator, v(r) the external potential
operator including ionic Coulomb potentials, vH [n](r) the Hartree potential operator, and vxc [n](r)the xc potential operator. The xc potential operator is defined as the functional derivative of the xc
energy with respect to the density, i.e. vxc [n](r) = δExc [n]/δn(r), which can be calculated using an
automatic differentiation library.
To discretize the Hamiltonian operator, a contracted Gaussian basis set is chosen as the basis of
the orbitals, bj (r) for j = {1...nb } where nb is the number of basis elements. As the basis set isnon-orthogonal, Eq. (1) becomes the Roothaan’s equation [28],F[n]pi = εi Spi ,(5)

where pRi is the basis coefficient for the i-th orbital, S is the overlap matrix with elements
Src = br (r)b∗c (r)dr, and RF[n] is the Fock matrix as a functional of the electron density profile, n, with elements Frc = br (r)H[n](r)b∗c (r)dr. Equation (5) can be solved by performing aneigendecomposition of the Fock matrix to obtain the orbitals from the electron density profile.
The ionic Coulomb potential, Hartree potential, and the kinetic energy operators in the Fock matrix
are constructed by calculating the Gaussian integrals with libcint [29]. The contribution from thexc potential operator in the Fock matrix is constructed by integrating it with the basis using SG-3
integration grid [30] with Becke scheme [31] to combine multi-atomic grids.

C

Gradient of KS-DFT components

C.1

Self-consistent iteration cycle

The self-consistent iteration cycle in a KS-DFT calculation can be written mathematically as
y = f (y, θ),(6)where y are the parameters to be found self-consistently, and f (·) is the function that needs to besatisfied which may depend additionally on other parameters, here denoted by θ. The self-consistent
iteration takes the parameters θ and an initial guess y0 , and produces the self-consistent parameters ythat satisfy the equation above. In our case, the self-consistent parameters are the elements of theFock matrix.
The gradient of a loss function L, with respect to the parameters θ, given the gradient with respect tothe self-consistent parameters ∂L/∂y, can be expressed as−1  
∂L
∂L
∂f
∂f
=−
I−
,
(7)
∂θ
∂y
∂y
∂θ
where I is the identity matrix. The inverse-matrix vector calculation in the equation above isperformed using the BiCGSTAB algorithm [32].
C.2

Gaussian basis

The differentiable KS-DFT employs a contracted Gaussian basis. The calculation of the Gaussianbasis in 3D space is performed using libcgto, a library within PySCF [33]. Libcgto is a library thatcan generate code to compute a Gaussian basis of any order. As libcgto does not provide an automaticdifferentiation feature, it has to be wrapped by PyTorch to enable the automatic differentiation modethrough the integrals calculated by libcgto.
8

A single Gaussian basis centered at r0 of order l, m, n is expressed as2

glmn (r; α, c, r0 ) = c(x − x0 )l (y − y0 )m (z − z0 )n e−α||r−r0 || ,(8)

where c is the basis coefficient and α is the exponential parameter. The contracted Gaussian basis isjust a linear combination of the single Gaussian bases above.
The gradient back-propagation calculation with respect to the parameters, (c, α, and r0 ) can be
expressed as∂L∂L ∂glmn=∂c∂glmn ∂c∂L
∂L ∂glmn=∂α∂glmn ∂α∂L
∂L ∂glmn=,∂r0∂glmn ∂r0

(9)
(10)
(11)

where
glmn
∂glmn
=
∂c
c

∂glmn
= − g(l+2)mn + gl(m+2)n + glm(n+2)∂α∂glmn= −∇glmn .
∂r0

(12)
(13)
(14)

The parameters required for the calculations above (e.g. the gradient of the basis and higher orderbases) can be provided by libcgto.

C.3

Gaussian integrals

Constructing the Fock matrix and the overlap matrix in Roothaan’s equation requires the evaluation ofGaussian integrals. The integrals typically involves 2-4 Gaussian bases of any order. Those integralsare:Z
Sijk−lmn =gijk (r)glmn (r) drZ
1
Kijk−lmn = −gijk (r)∇2 glmn (r) dr
2
XZ
Zc
Cijk−lmn =glmn (r) dr
gijk (r)|r − rc |cZ
1
Eijk−lmn−pqr−stu = gijk (r)glmn (r)gpqr (r0 )gstu (r0 ) drdr0 ,|r − r0 |

(15)
(16)
(17)
(18)

where S is the overlap term, K is the kinetic term, C is the ionic Coulomb term of an ion of charge
Zc at coordinate rc , and E is the two-electron Coulomb term (i.e. the Hartree potential operator).
These integrals can be calculated efficiently by libcint [29].
9The gradients of a loss function L with respect to the basis parameters θ = {c, α, r0 } are:








∂L
∂L
∂L
∂L
∂L
=
+
+
+
(19)
∂θ
∂θ S
∂θ K
∂θ C
∂θ E


Z

Z
∂L
∂L
∂gijk (r)∂glmn (r)=glmn (r) dr + gijk
dr(20)
∂θ S ∂S
∂θ
∂θ


Z

Z
∂L
1 ∂L
∂gijk (r) 2∂glmn (r)=−∇ glmn (r) dr + gijk ∇2
dr
(21)
∂θ K
2 ∂K
∂θ
∂θ

Z


Z
∂L X
∂gijk (r) Zc
Zc ∂glmn (r)∂L=glmn (r) dr + gijk
dr(22)
∂θ C ∂C c∂θ |r − rc ||r − rc |
∂θ


Z
Z
∂L
∂L
∂gijk (r) glmn (r)gpqr (r0 )gstu (r0 )∂glmn (r) gijk (r)gpqr (r0 )gstu (r0 )
=
dr
+
dr
∂θ E ∂E∂θ|r − r0 |∂θ
|r − r0 |ZZ∂gpqr (r0 ) gijk (r)glmn (r)gstu (r0 )∂gstu (r0 ) gijk (r)glmn (r)gpqr (r0 )
+
dr
+
dr
.
∂θ
|r − r0 |∂θ
|r − r0 |(23)
If the parameter θ does not belong to the basis g· , then ∂g· /∂θ is 0. Otherwise, the values of ∂g· /∂θ
for various parameters are listed in the previous subsection. For the gradient with respect to theatomic position, rc , the expression is given by

Z
Z
∂L ∂L
Zc
Zc
glmn (r) dr + gijk (r)∇glmn (r) dr .
(24)=
∇gijk (r)∂rc ∂C|r − rc ||r − rc |The expression above is obtained by performing partial integration. As the derivative of the Gaussianbasis with respect to its parameters is also a Gaussian basis, the required integrals can be provided bylibcint [29].
C.4

Exchange-correlation energy and potential

The exchange-correlation (xc) energy in our calculation is a hybrid between available xc energyfunctionals (i.e. the LDA [11] and PBE [20] models) and a neural network xc energy. The gradient
back-propagation through the xc neural network is easily provided by PyTorch. However, as we use
Libxc [21] to calculate the established xc energy, it has to be wrapped by PyTorch to provide thegradient back-propagation calculation. As Libxc can also calculate the gradient of the energy withrespect to its parameters, providing the automatic differentiation feature by wrapping it with PyTorchis just a matter of calling the right function in Libxc.
If the PyTorch wrapper for Libxc is provided, calculating the xc potential can be done using theautomatic differentiation feature provided by PyTorch. The component of the Fock matrix from thexc potential for the LDA spin-polarized case is
Z
(LDA)
Vij
= φ∗i (r)vs(LDA) (r)φj (r) dr,
(25)
(LDA)

where the subscript s can be ↑ or ↓ representing the spin of the electron, and vsis the xc potential,
Z∂
vs(LDA) =nε(LDA) (n↑ , n↓ ) dr,(26)∂ns
with ε(LDA) the energy density per unit volume, per electron, provided by Libxc. The 3D spatialintegration is performed using SG-3 integration grid [30] with Becke scheme [31] to combine multiatomic grids. The derivation with respect to ns in the equation above can simply be calculated using
the automatic differentiation feature of PyTorch wrapping the Libxc library.
For spin-polarized GGA, the potential linear operator is [34]
Z
(GGA)
Vij
= φ∗i (r)vs(GGA) (r)φj (r) dr
(27)


Z
Z
∂
∂
vGGAs =nε(GGA) (n↑ , n↓ , ∇n↑ , ∇n↓ ) dr +nε(GGA) dr · ∇,
(28)
∂ns
∂∇ns
where the derivatives with respect to ns and ∇ns can be performed by automatic differentiation, andthe last ∇ term is applied to the basis which can be provided by libcgto from PySCF [33].
10

C.5

(Near-)degenerate case of eigendecomposition

One challenge in propagating the gradient backward for some molecules is the degeneracy of
eigenvalues in the eigendecomposition. The eigendecomposition of a real and symmetric matrix, A,
can be written as
Avi = λi vi(29)with λi and vi as the i-th eigenvalue and eigenvector, respectively.
In the non-degenerate case, the gradient of a loss function L with respect to each element of matrix
A can be written as
XX
∂L
∂L T
=−
(λi − λj )−1 vi viT
vj .
(30)
∂A
∂v
j
j
i6=j

A problem arises when degeneracy appears as a term of 0−1 appears in the equation above.
To solve this problem, we follow the equation from ref [18], which can be written simply as
X X
∂L T
∂L
(λi − λj )−1 vi viT
=−
v ,
∂A
∂vj j
j

(31)

i6=j,λi 6=λjwhich just removes the degenerate terms in the sum. However, the equation above is only valid aslong as the loss function does not depend on which linear combinations of the degenerate eigenvectors
are used and the elements in matrix A depends on some parameters that always make the matrix A
real and symmetric. Otherwise, the equation above becomes invalid and the actual gradient becomesmathematically undefined.

D

Neural network training

D.1

DatasetThe atomization energy dataset is compiled from NIST CCCBDB [22]. CCCBDB contains theexperimental data for enthalpy of atomization energy. To convert it into electronic atomization energy
which can be calculated by KS-DFT calculations, the zero-point vibration energy must be excluded
from the enthalpy of atomization energy. Some diatomic molecules have the experimental valuesof the zero-point energy in CCCBDB from ref [35]. For the rest of the molecules, the zero-point
energy is estimated from the fundamental vibrations of the molecules which are also available from
CCCBDB.
D.2

Neural networks

The neural network is an ordinary feed-forward neural network with 3 hidden layers consist of 32
elements each with softplus [36] activation function to have infinite differentiability of the neuralnetwork. To avoid non-convergence of the self-consistent iterations during the training, the parameters
α and β are initialized to have values of 1 and 0 respectively.
The neural network takes an input of the density, n, spin polarized density, ξ = (n↑ − n↓ )/n, andnormalized density gradient, s = |∇n|2 /(24π 2 n4 )1/3 (for XCNN-PBE). As the density (n) and thenormalized density gradient (s) can have values between 0 up to a very large number, these values
are renormalized by applying log(1 + x), i.e. n ← log(1 + n) and s ← log(1 + s).
D.3

Training

The parameters of XCNN are updated based on the gradient of a loss function with respect to its
parameters. The loss function is given by a combination of two terms
L=

Ndp Z
Nae 
2 w X
wae X
dp
(ae)
(ae)
2
Êi − Ei+[n̂i (r) − ni (r)] dr,
Nae i=1
Ndp i=1(32)

where w(·) are the weights assigned to each type of the data, N(·) are the number of entries ofeach datatype in the dataset, E (ae) and E (ip) are, respectively, the predicted atomization energy and11ionization potential using XCNN and KS-DFT, n is the electron density profile predicted with XCNN
and KS-DFT. The variables with a hat, Ê (·) and n̂, are the target variables that were obtained fromexperimental databases (for energies) and CCSD calculations (for electron density profiles).
The XCNN is trained using the RAdam [37] algorithm with a constant learning rate of 10−4 . The
algorithm and learning rate were chosen arbitrarily without further parameter tuning due to resource
and time limitations. The weights for training of the neural networks are:
• XCNN-LDA: wae = 1340, wdp = 170.
• XCNN-PBE: wae = 1340, wdp = 5360.
The weights were chosen to make the mean error for each component approximately equals to 1.
The XCNN training took 2 weeks using 32 CPU cores. We did not use GPU in this training because
the integral calculations in DQC are not written for GPU.

E

Molecules in the test dataset

The 104 molecules in the test dataset, as well as the subset they belong, to are listed below.
1. LiH (Lithium Hydride): NHC-1
2. BeH (beryllium monohydride): NHC-1
3. CH (Methylidyne): HC
4. CH3 (Methyl radical): HC
5. CH4 (Methane): HC
6. NH (Imidogen): NHC-1
7. NH2 (Amino radical): NHC-1
8. NH3 (Ammonia): NHC-1
9. OH (Hydroxyl radical): NHC-1
10. H2 O (Water): NHC-1
11. HF (Hydrogen fluoride): NHC-1
12. SiH2 (silicon dihydride): NHC-2
13. SiH3 (Silyl radical): NHC-2
14. SiH4 (Silane): NHC-2
15. PH3 (Phosphine): NHC-2
16. H2 S (Hydrogen sulfide): NHC-2
17. HCl (Hydrogen chloride): NHC-2
18. Li2 (Lithium diatomic): NHC-1
19. LiF (lithium fluoride): NHC-1
20. C2 H2 (Acetylene): HC
21. C2 H4 (Ethylene): HC
22. C2 H6 (Ethane): HC
23. CN (Cyano radical): NHC-1
24. HCN (Hydrogen cyanide): NHC-1
25. CO (Carbon monoxide): NHC-1
26. HCO (Formyl radical): NHC-1
27. H2 CO (Formaldehyde): NHC-1
28. CH3 OH (Methyl alcohol): SHC
29. N2 (Nitrogen diatomic): NHC-1
1230. N2 H4 (Hydrazine): NHC-1
31. NO (Nitric oxide): NHC-1
32. O2 (Oxygen diatomic): NHC-1
33. H2 O2 (Hydrogen peroxide): NHC-1
34. F2 (Fluorine diatomic): NHC-1
35. CO2 (Carbon dioxide): NHC-1
36. Na2 (Disodium): NHC-2
37. Si2 (Silicon diatomic): NHC-2
38. P2 (Phosphorus diatomic): NHC-2
39. S2 (Sulfur diatomic): NHC-2
40. Cl2 (Chlorine diatomic): NHC-2
41. NaCl (Sodium Chloride): NHC-2
42. SiO (Silicon monoxide): NHC-2
43. CS (carbon monosulfide): NHC-2
44. SO (Sulfur monoxide): NHC-2
45. ClO (Monochlorine monoxide): NHC-2
46. ClF (Chlorine monofluoride): NHC-2
47. Si2 H6 (disilane): NHC-2
48. CH3 Cl (Methyl chloride): SHC
49. HOCl (hypochlorous acid): NHC-2
50. SO2 (Sulfur dioxide): NHC-2
51. BF3 (Borane trifluoro-): NHC-1
52. BCl3 (Borane trichloro-): NHC-2
53. AlF3 (Aluminum trifluoride): NHC-2
54. AlCl3 (Aluminum trichloride): NHC-2
55. CF4 (Carbon tetrafluoride): NHC-1
56. CCl4 (Carbon tetrachloride): NHC-2
57. OCS (Carbonyl sulfide): NHC-2
58. CS2 (Carbon disulfide): NHC-2
59. CF2 O (Carbonic difluoride): NHC-1
60. SiF4 (Silicon tetrafluoride): NHC-2
61. SiCl4 (Silane tetrachloro-): NHC-2
62. N2 O (Nitrous oxide): NHC-1
63. ClNO (Nitrosyl chloride): NHC-2
64. NF3 (Nitrogen trifluoride): NHC-1
65. O3 (Ozone): NHC-1
66. F2 O (Difluorine monoxide): NHC-1
67. C2 F4 (Tetrafluoroethylene): NHC-1
68. CF3 CN (Acetonitrile trifluoro-): NHC-1
69. CH3 CCH (propyne): HC
70. CH2 CCH2 (allene): HC
71. C3 H4 (cyclopropene): HC
72. C3 H6 (Cyclopropane): HC
73. C3 H8 (Propane): HC
13

74. C4 H6 (Methylenecyclopropane): HC
75. C4 H6 (Cyclobutene): HC
76. CH3 CH(CH3 )CH3 (Isobutane): HC
77. C6 H6 (Benzene): HC
78. CH2 F2 (Methane difluoro-): SHC
79. CHF3 (Methane trifluoro-): SHC
80. CH2 Cl2 (Methylene chloride): SHC
81. CHCl3 (Chloroform): SHC
82. CH3 CN (Acetonitrile): SHC
83. HCOOH (Formic acid): SHC
84. CH3 CONH2 (Acetamide): SHC
85. C2 H5 N (Aziridine): SHC
86. C2 N2 (Cyanogen): SHC
87. CH2 CO (Ketene): SHC
88. C2 H4 O (Ethylene oxide): SHC
89. C2 H2 O2 (Ethanedial): SHC
90. CH3 CH2 OH (Ethanol): SHC
91. CH3 OCH3 (Dimethyl ether): SHC
92. C2 H4 O (Ethylene oxide): SHC
93. CH2 CHF (Ethene fluoro-): SHC
94. CH3 CH2 Cl (Ethyl chloride): SHC
95. CH3 COF (Acetyl fluoride): SHC
96. CH3 COCl (Acetyl Chloride): SHC
97. C4 H4 O (Furan): SHC
98. C4 H5 N (Pyrrole): SHC
99. H2 (Hydrogen diatomic): NHC-1
100. HS (Mercapto radical): NHC-2
101. C2 H (Ethynyl radical): HC
102. CH3 O (Methoxy radical): SHC
103. CH3 S (thiomethoxy): SHC
104. NO2 (Nitrogen dioxide): NHC-1

14