D4FT: A Deep Learning Approach to Kohn-Sham Density Functional Theory
Authors: Tianbo Li
Min Lin
Zheyuan Hu
Kunhao Zheng
Giovanni Vignale
Kenji Kawaguchi
A.H. Castro Neto
Kostya S. Novoselov
Shuicheng YAN
Venue: ICLR 2023 notable top 25% (2023)
Keywords: AI for Science
Quantum Chemisty
Density Functional Theory
Deep Learning
Kohn-Sham Equation.

Abstract: Kohn-Sham Density Functional Theory (KS-DFT) has been traditionally solved by the Self-Consistent Field (SCF) method. Behind the SCF loop is the physics intuition of solving a system of non-interactive single-electron wave functions under an effective potential. In this work, we propose a deep learning approach to KS-DFT. First, in contrast to the conventional SCF loop, we propose to directly minimize the total energy by reparameterizing the orthogonal constraint as a feed-forward computation. We prove that such an approach has the same expressivity as the SCF method, yet reduces the computational complexity from O(N^4) to O(N^3). Second, the numerical integration which involves a summation over the quadrature grids can be amortized to the optimization steps. At each step, stochastic gradient descent (SGD) is performed with a sampled minibatch of the grids. Extensive experiments are carried out to demonstrate the advantage of our approach in terms of efficiency and stability. In addition, we show that our approach enables us to explore more complex neural-based wave functions. 

Published as a conference paper at ICLR 2023
D4FT: A DEEP LEARNING APPROACH TO KOHN-SHAM DENSITY FUNCTIONAL THEORY
Tianbo Li" s,. Min Lin" a. ZZeyuan Hu Kunhao Zheng", Giovanni Vignale Kenji Kawaguchi", A. H. Castro Neto Kostya S. Novoseloos Shuicheng Yan "SEA AI Lab, "School of Computing. National University of Singapore, IInstitute for Functional Intelligent Materials, National University of Singapore, "Centre for Advanced 2D Materials, National University of Singapore; ("1itb, flinmin) @esaa.esseeccom com
ABSTRACT
Kohn-Sham Density Functional Theory (KS-DFT) has been traditionally solved by the Self-Consistent Field (SCF) method. Behind the SCF loop is the physics intuition of solving a system of non-interactive single-clectron wave functions un- der an effective potential. In this work, we propose a deep learning approach to KS-DFT. First, in contrast to the conventional SCF loop, we propose directly minimizing the total energy by reparameterizing the orthogonal constraint feed-forward computation. We prove that such an approach has the same expres- sivity as the SCF method yet reduces the computational complexity from ((N*) to ((N5). Second, the numerical integration, which involves a summation over the quadrature grids, can be amortized to the optimization steps. At each step. stochastic gradient descent (SGD) is performed with a sampled minibatch of the grids. Extensive experiments are carried out to demonstrate the advantage of our approach in terms of efficiency and stability. In addition, we show that our ap- proach enables us to explore more complex neural-basse based wave functions.
INTRODUCTION
Density functional theory (DFT) the most successful quantum-mechanical method, which is widely used in chemistry and physics for predicting
properties of matters (Szabo) & Ostlund 2012: Levine et al. 22099: Koch & Holthausen] 2015). As scientists are exploring more complex molecules and materials, DFT methods often limited in scale or accuracy due to their computation complexity. On the other hand, Deep Learning (DL) has achieved great success in function approximations (Hornik al., [989), optimization algorithms (Kingma & Ba] 2014), and systems (Bradbury et al. 2018) in the past decade. Many aspects of deep learning can be harnessed improve DFT. Of them, data-driven function fitting is the most straightforward and often the first to be considered. It has been shown that models learned from a sufficient amount of data generalize greatly to unseen data, given that the models have the right inductive bias. The Hohenberg Kohn theorem proves that the ground state energy is a functional of electron density (Hohenberg & Kohn, 1964a), but this functional is not available analytically. This is where data-driven learning can be helpful for DFT. The strong function approximation capability of deep learning gives hope to learn- ing such functionals in a data-driven manner. There have already been initial successes in learning the exchange-correlation functional (Chen et al. 2020a b Dick & Fernandez-Serra 2020). Further- more, deep learning has shifted the mindsets of researchers and engineers towards differentiable programming. Implementing the derivative of function has extra cost if the primal function is implemented with deep learning frameworks. Derivation of functions frequently appears in DFT, e.g., estimating the kinetic energy of wave function; calculating generalized gradient approxi- mation (GGA) exchange-correlatin functional, etc. Using modern automatic differentiation (AD) techniques ease the implementation greatly (Abbott et al. 2021).
Despite the numerous efforts that apply deep learning to DFT, there is still a vast space for explo- ration. For example, the most popular Kohn-Sham DFT (KS-DFT) (Kohn & Sham 1965) utilizes
Equal Contribution. Our code will be available on https://githut wtt.tub..coo/ssllss///////

---

Published as a conference paper at ICLR 2023
the self-consistency field (SCF) method for solving the parameters. At each SCF step, it solves closed-form eigendecomposition problem, which finally leads to energy minimization. However, this method suffers from many drawbacks. Many computational chemists and material scientists criticize that optimizing via SCF is time-consuming for large molecules or solid cells, and that the convergence of SCF is not always guaranteed. Furthermore, DFT methods often utilize the linear combination of basis functions as the ansatz of wave functions, which may not have satisfactory expressiveness approximate realistic quantum systems.
To address the problems of SCF, we propose a deep learning approach for solving KS-DFT. Our ap- proach differs from SCF in the following aspects. First, the eigendecomposition steps in SCF come from the orthogonal constraints on the wave functions; we show in this work that the original objec- tive function for KS-DFT can converted into an unconstrained equivalent by reparameterizing the orthogonal constraints part of the objective function. Second, we further explore amortizing the integral in the objective function over the optimization steps, i.e., using stochastic gradient descent (SGD), which is motivatee both empirically and theoretically for large-scale machine learn- ing (Bottou et al., 2018). We demonstrate the equivalence between our approach and the conven- tional SCF both empirically and theoretically. Our approach reduces the computational complexity from O(N*) to ) which significantly improves the efficiency and scalability of KS-DFT. Third, gradient-based optimization treats all parameters equally. We show that it is possible to op- timize more complex neural-based wave functions instead of optimizing only the coefficients. In this paper, we instantiate this idea with local scaling transformation as an example showing how to construct neural-based wave functions for DFT.
DFT PRELIMINARIES
Density functional theory (DFT) is among the successful quantun-mechanical simulation methods for computing electronic structure and all electron-related properties. DFT defines the ground state energy as a functional of the electron density p R R:
Eg. E(p).
(1)
The Hohenberg- Kohn theorem (Hohenberg & Kohn 1964b) guarantees that such functionals E exists and the ground state energy can be determined uniquely by the electron density. However, the exact definition of such functional has been puzzling obstacle for physicists and chemists. Some approximations, including the famous Thomas- -ermi method and Kohn-Shan method, have been proposed and have later become the most important ab-initio calculation methods.
The Objective Function One of the difficulties in finding the functional of electron density is the lack of an accurate functional of the kinetic energy. The Kohn-Sham method resolves this issue by introducing an orthogonal set of single-particle wave functions {wa} and rewriting the energy a functional of these wave functions. The energy functional connects back the Schrödinger equation. Without compromising the understanding of this paper, we leave the detailed derivation from Schrödinger equation and the motivation of the orthogonality constraint Appendix B.I As far as this paper is conceree,, we focus on the objective function of KS-DFT, defined as,
Egs min H {0f)
(2)
min
Ex Sx + E HE Enlly +1+ Excl(v (+T)
(3)
s.t. (W/l)g) = -ō.
(4)
where of is a wave function mapping 3R3 + C. and vi*h denotes its complex conjugate. For simplicity, we use the bra-ket notation for (v/ w = v (r) (r)dr. dij is the Kronecker delta function. The superscript a {a, 3) denotes the spin Exin: FEST. Fg, Exe the ki- netic, external potential (nuclear attraction), Hartree (Coulomb repulsion between electrons) and
We omit the spin notation in the following sections for simplification reasons.
2

---

Published as a conference paper at ICLR 2023
exchange-correlation energies respectively, defined by
/ (V) dr,
Exin [v
(5)
Σ Σ / Text (xx)
N
(6)
(2i En 108 (r)))) (2) E. 14g (r))
EnlfogH 2
drdr',
(7)
r-r
N
ExclLugH SxclteH-=E>
Exx (r) (r)
(8)
in which text the external potential defined by the molecule's geometry. Exc is the exchange- correlation energy density, which has different instantiations. The entire objective function is given analytically, except for the function v that we replace with parametric functions to be optimized. In the rest of this paper, focus on the algorithms that optimize this objective while minimizing the discussion on its scientific background.
Kohn-Sham Equation and SCF Method The object function above be solved by Eu- ler Lagrange method, which yields the canonical form of the well-known Kohn-Sham equation,
I A2 + vex((r) + / d3y EN |ot''( (x)))))
+ (x)(r) v/ = E
where N the total number of electrons. Vext and Use are the external and exchange-correlation potentials, respectively. This equation is usually solved in iterative manner called Self-Consisten Field (SCF). This method starts with an initial guess the orthogonal set single electron wave functions, which are then used for constructing the Hamiltonian operator. The new wave functions and corresponding electron density can be obtained by solving the eigenvalue equation. This process is repeated until the electron density of the system converges. The derivation of the Kohn-Sham equation and the SCF algorithm is presented in Appendix B.2
The LCAO Method the general case, we can use any parametric approximator for i to trans- form the optimization problem from the function space to the parameter space. In quantum chem- istry, they are usually represented by the linear combination of atomic orbitals (LCAO):
B
(r) = Σ (r))(r),
(10)
where o; are atomic orbitals (or basis functions more generally), which are usually pre-determined analytical functions, e.g., truncated series spherical harmonies or plane waves. We denote the number of the basis functions as B. The single-particle wave functions are linear combinations of these basis functions with Cij as the only optimizable parameters. We introduce vectorized notations
V == (11,4/2, .., "v)] T. == @21.02. ........ ,NN)" T and C := [cy]. so the LCAO wave functions can be written as,
V C
(11)
Classical atomic obitals include Slater-type orbitals, Pople basis (Ditchfield et al.) 1971) correlation-consisten basis sets (Dunning Jr 1989), etc.
3 A DEEP LEARNING APPROACH TO KS-DFT
In this section, we propose a deep learning approach to solving KS-DFT. Our method can be de- scribed by three keywords:
deep learning: our method is deep learning native; is implemented in a widely-used deep learning framework JAX;
differentiable: all the functions are differentiable, and thus the energy can be optimized via purely gradient-based methods;
3

---

Published as a conference paper at ICLR 2023
Molecular Configuration
urruy/matrix
Atomic Orbitals
function
Decomposition
overlap matrix
operation functional
LCAO
Integrator
forward computation
gradient wrt. coefficients
Riv Molecular orbitals
gradient wrt. orbitals
FiLE
Figure I: An illustration the computational graph of D4FT. In conventional DFT methods using fixed basis sets, the gradients w.r.t. the basis functions (blue dotted line) are unnecessary.
direct optimization: due to the differentiability, our method does not require self-consistent field iterations, but the convergence result also self-consistent.
We refer to our method as D4FT that represents the above 3 keywords and DFT. Fig. shows an overarching framework of our method. We elaborate on each component in the following parts.
3.1 REPARAMETERIZE THE ORTHONORMAL CONSTRAINTS
Constraints can be handled with various optimization methods, e.g.. the Lagrangian multiplier method or the penalty method. deep learning, into the computation graph. For example, when want the normalization constraint vector v, we can reparameterize to y/ /llly/ such that it converts to solving without constraints. Then this constraint-free optimization can benefit from the existing differentiable frameworks and various gradient descent optimizers. The wave functions in Kohn-Sham DFT have to satisfy the constraint (wl))) ōij given in Eq. 4 Traditionally the constraint is handled by the Lagrangian multiplier method, which leads to the SCF method introduced in Sec.
preferred to reparameterize the constraints
2
To enable direct optimization, we
propose the following reparameterization of the constraint.
Using LCAO wave functions, the constraint translates to E Cu Ei cirhh(r)dr = 8.j- Or in the matrix form
CSC =
(12)
S is called the overlap matrix of the basis functions, where the ij-th entry Sig (d)ld))) The literature on whitening transformation (Kessy et al., 2018) offers many ways to construct C to satisfy Eq. 12 based on different matrix factorization S:
1/2 1/21
PCA whitening, with U AU = S Cholesky whiteming, with LL" S-1
C=
QL T
(13)
Q121
ZCA whitening
Taking PCA whitening functions, A 1/ is fixed. The Q matrix can be any orthonormal matrix; for deep learning, can parameterize Q in differentiable way using QR decomposition of an unconstrained matrix W:
example. Since S be precomputed from the overlap of basis
Q, R QR(QR)())))
(14)
Besides QR decomposition, there are several other differentiable methods construct the thonormal matrix Q. e.g.. Householder transformation (Mathiasen et al.) 2020). and exponential map (Lezcano-Casade Martunez-Rubi Rubio 2019). Finally, the wave functions can be written as:
V ww- QD4.
(15)
In this way, the wave functions Vw always orthogonal given arbitrary W. The searching over the orthogonal function space is transformed into optimizing over the parameter space that W re- sides. Moreover, this parameterization covers all possible sets of orthonormal wave functions in the space spanned by the basis functions. These statements are formalized in the following proposition, and the proof is presented in Appendix C.I
4

---

Published as a conference paper at ICLR 2023
Proposition 3.1. Define the original orthogonal function functiooi space F and the transformed search space
Fw by ( (1/1, 1/2, ... WN) V = C , E R NNN, (Bxlvy)) = 6ig) and Fw =
( = (0) 22 V/ )T Vw = (QD.@QQQ.. = QR W)) W E RNNNN. Then, they
W
equivalent to Fw F.
3.2 STOCHASTIC GRADIENT
SGD is the modern workhorse for large-scale machine learning optimizations. It has been harnessed to achieve unprecedented scale of training. which would have been impossible with full batch training. We elaborate in this section on how DFT could also benefit from SGD.
Numerical Quadrature The total energies defined in Equations 5 8 integrals that involve the wave functions. Although the analytical solution to the integrals of commonly-used basis sets does exist, most DFT implementations adopt numerical quadrature integration, which approximate the value of a definite integral using a set of grids g = {(2,, )=) and corresponding weights, respectively:
where æ; and wy are the coordinate
L. [(x)dx R
Σ f(x:)w..
(16)
These grids and weights can be obtained via solving polynomial equations (Golub & Welsch 1969. Abramowitz & Stegun 1964). One key issue that hinders its application in large-scale scale systems that the Hartree energy requires at least On??) calculation between every two grid points. Some large quantum systems will need 100k ~ 10m grid points, which causes out-of-memory errors and hence is not feasible for most devices.
needs to compute the distance
Stochastic Gradient on Quadrature Instead of evaluating the gradient of the total energy at all grid points in g. we randomly sample minibatch g' C g. where g' contains m grid points, m < n, and evaluate the objective and its gradient on this minibatch. For example, for single integral energies such kinetic energy, the gradient can be estimated by,
ek aw
I ( (V (a;))
W1
(17)
OW
The gradients of external and exchange-correlation energies can be defined accordingly. The gradi- ent of Hartree energy, which is double integral, can be defined as,
2m(n 1)
Σ
Σ
V w (x;) 2 E zz,
Oehn (axe) Ow
(18)
OW m(m
1)
can be proved that the expectation of the above stochastic gradient is equivalent to those of the full gradient. Note that the summation over the quadrature grids resembles the summation over all points a dataset in deep learning. Therefore, in our implementation, we can directly rely on AD with minibatch to generate unbiased gradient estimates.
3.3 THEORETICAL PROPERTIES
Asympototic Complexity Here we analyze the computation complexity of D4FT in comparison with SCF. We use N to denote the number of electrons, B to denote the number of basis functions, n for the number of grid points for quadrature integral, and m for the minibatch size when the stochastic gradient used. The major source of complexity comes from the hartree energy Eu,, as it includes a double integral.
In our direct optimization approach, computing the hartree energy involves computing Σ, (1x(())) for all grid points, which takes OnN After that O(n2) computation is needed to compute and aggregate repulsions energies between all the electron pairs. Therefore, the total complexity is OnNN B) + Onn2). P In the SCF approach, it costs OnN O(n a lot more expensive because it computes the Fock matrix instead of scalar energy.
We assume the backward mode gradient computation has the same complexity the forward computation.
5

---

Published as a conference paper at ICLR 2023
Considering both n and B approximately linear to N, the direct optimization approach has more favorable (NV)) compared to (((*)) for SCF. However, since n is often much bigger than N and B practically, the minibatch SGD is indispensable ingredient for computation efficiency. At each iteration, minibatch reduces the factor of n to small constant m, making it B) for D4FT. A full breakdown of the complexity is available in Appendix
Self-Consistency The direct optimization method will converge at self-consistent point. We demonstrate this in the following. We first give the definition of self-consistency as follows. Definition 3.1 (Self-consistency). The wave functions & is said to be self-consistent if the eigen- functions of its corresponding Hamiltonian H ( is V itself. i.e., there exists real number E I, v
such that H [ )
The next proposition states the equivalence between SCF and direct optimization. It states that the convergence point of D4FT is self-consistent. The proof is shown in the Appendix C.2 Proposition 3.2 (Equivalence between SCF and D4FT for KS-DFT). Let Tp be a local optimimal of the ground state energy Egs defined in Eq. such that &.g
O, then V is self-consistent.
4 NEURAL BASIS WITH LOCAL SCALING TRANSFORMATION
In previous sections, the basis functions are considered given and fixed. Now we demonstrate and discuss the possibility of learnable basis set that can be jointly optimized in D4FT. We hope that this extension could serve as stepping stone toward neural-based wave function approximators.
As discussed in Sec. 2 LCAO wave functions are used in the existing DFT calculations. This restricts the wave functions to be in the subspace spanned by the basis functions, and it would then require large number of basis functions in order to make a strong approximator. From the deep learning viewpoint, it is more efficient to increase the depth of the computation. The obstacle to using deep basis functions is that the overlap matrix S will change with the basis, requiring us to do expensive recomputation of whitening matrix D at each iteration. To remedy this problem, we introduce basis functions (p; in the following form
p:(r) == det Jy(r) | (f(())
(19)
where f 3 R3 is parametric bijective function, and det Jet / (r) denotes its Jacobian determi- nant. It verifiable by change of variable that
(p/le)) = f det
( (i())) = (dolor)-
=
Therefore, the overlap matrix S will be fixed and remain unchanged even if f varies.
Within our framework, we can use parameterized function fo and optimize both 0 and W jointly with gradient descent. As proof of concept, we design fo as follows, which we term as neural local scaling:
fo(r) = \o(r)r, (r == an)ge(T)).
(20)
(21)
where 17 is the sigmoid function, and is scaling factor to control the range of Xy- go RR3 R can be arbitrarily complex neural network parametrized by 0. This method has two benefits. First, by introducing additional parameters, the orbitals are more expressive and potentially achieve better ground state energy. Second, the conventional gaussian basis function struggles to tackle the cusps the location of the nuclei. By introducing the local scaling transformation, we have a scaling function that can control the sharpness of the wave functions near the nuclei, and as a consequence, this approximator can model the ground state wave function better. We perform experiments using neural local transformed orbitals on atoms. Details and results will be presented in Section 6.3
5 RELATED WORK
Deep Learning for DFT There are several lines of work using deep learning for DFT. Some of them use neural networks in supervised manner, i.e., data from simulations experiments

---

Published as a conference paper at ICLR 2023
fit neural networks. A representative work is conducted by (Gilmer al. 2017). which proposes a message- passsngg neural network and predicts DFT results on the QM9 dataset within chemical accuracy. (Custódio et al. 2019) and Ellis al. (2021) also follow similar perspectives and predict the DFT ground state energy using feed-forward neural networks. Another active research line falls into learning the functionals, i.e. the kinetic and exchange-correlation functionals. As is difficult to evaluate the kinetic energy given the electron density, some researchers turn to deep learning to fit a neural kinetic functional (Alghadeer al., 2021 Ghasemi Kühhe [021 Ellis et al. (021). which makes energy functional orbital-free (depending only on electron density). Achieve chemical aceuracy is one of the biggest obstacles in DFT, many works (Kirkpatrick al. [2021] Kalita al.
2021 Ryabov et al.] 2020 Kasim & Vinko 2021 Dick & Fernandez-Serra, 202 Nagai et al.,
2020 research into approximating the exchange-correlatio
functional with neural networks in the
past few years, for making the DFT more accurate. It is worth noting that, the above deep learning approaches are mostly data-driven, and their generalization capability is to be tested.
Differentiable DFT As deep learning approaches have been applied to classical DFT methods, many researchers (Tamayo-Mendoza et al.] 2018: Li et al. [02T]] Kasim & Vinko, 202T Dick & Fernandez-Serra [2021a: Laestadius et al. 2019) have discussed how to make the training of these neural-based DFT methods, particularly the SCF loop, differentiable, so that the optimiza- tion of these methods be easily tackled by the auto-differentiation frameworks. Li et al. (2021) uses backward mode differentiation to backpropagate through the SCF loop. In contrast, [Tamayo- Mendoza al. (2018) employs the forward mode differentiation, which has less burden on system memory compared to backward mode. Kasim & Vinko (202I) uses the implicit gradient method, which has the smallest memory footprint, however it requires running the SCF loop convergence.
Direct Optimization in DFT Albeit SCF method dominates the optimization in DFT, there are researches exploring the direct optimization methods (Gillan 1989 Van Voorhis & -ordon & Arias 2000 Vande Vondele & Hutter 2003; Weber et 2008: [vanov et al., orthonormality constraint of the wave functions. A straightforward way to achieve this is via explicit orthonormalizaiton of the orbitals after each update (Gillan 1989). In recent years, some have investigated direct optimization of the total energy with the orthonormality constraint incorporated in the formulation of wave functions. A represen- tative method is to express the coefficients of basis functions as an exponential transformation of a kkew-Hemmitin matrix (Ismail-Bes -Biggi & Arias 2000 Ivanov et al. 2021
2002
2021..
The challenging part is how to preserve
6 EXPERIMENTS
In this section, we demonstrate the accuracy and scalability of D4FT via numerical experiments on molecules. We compare our method with two benchmarks,
PySCF (Sun et al. 2018): one of the most widely-used open-sourced quantum chemistry compu- tation frameworks. uses python/c implementation.
JAX-SCF: our implementation of classical SCF method with Fock matrix momentum mechanism.
We implement our D4FT method with the deep learning framework JAX (Bradbury et al. 2018). For fair comparison with the same software/hardware environment, reimplemented the SCF method in JAX. All the experiments with JAX implementation (D4FT, JAX-SCF) are conducted on an NVIDIA A A100 GPU with 40GB memory. As reference, we also with PySCF on a 64-core Intel Xeon CPU 10GHz with 128GB memory.
6.1 ACCURACY
We first evaluate the accuracy of D4FT two tasks. We first compare the ground state energy obtained from D4FT with PySCF on a series of molecules. We then predict the paramagnetism of molecules to test if our method handles spin correctly. For both tasks, we use the 6-31g basis set with the LDA exchange-correlation functional.
Ground State Energy The predicted ground state energies are presented in Table I It can be that the absolute error between D4FT and PySCF is smaller than 0.02Ha (s=0.54e over all the molecules compared. This validates the equivalence of our SGD optimization and SCF loop.
7

---

Published as a conference paper at ICLR 2023
C36
C60
-2.2
C80
Convergence Time (6)
Fnergy (kHo)
-1.1
-2.4
JAX-SCF converges fäster
-1.2
-2.6
converge
the same time
-2.8
10'
-1.3
-2.2
750
-2.75
C100
4.5
C160
-5.0
C180
Per Spoch Time 6
Elaegg (Mi5)
-3.00
-5.5
-3.25
-6.0
D4FT converges enuch faster
-6.5
-3.75
500
5000
500000
2.50
500
750
Wall clock time (4)
Wall clock time (4)
Wall clock time (9)
# Orbitals
D4FT
JAX-SCF
D4FT convergence point
JAX-SCF convergence point
Figure 2: Convergence speed on different carbon Fullerene molecules. Columns 1-3: the conver- gence curves on different molecules. X-axis: wall clock time (s), Y-axis: total energy (Ik Ha). Right-top: the scale of convergence time. Right-bottom: the scale of per-epoch time. It can be seen that as the number of orbitals increases, D4FT scales much better than JAX-SCF.
Table I: Partial results of ground state energy calculation (LDA, 6-31g, Ha)
Molecule
Hydrogen
Methane
Water
Oxygen
Ethanol
Benzene
PySCF
-1.03864
-39.49700
-75.15381
-148.02180
-152.01379
-227.35910
JAX-SCF
-1.03919
-39.48745
-75.15124
-148.03399
-152.01460
-227.36755
D4FT -1.03982
-39.48155
-75.13754
-148.02304
-152.02893
-227.34860
Magnetism Prediction We predict whether a molecule is paramagnetic or diamagnetic with D4FT. Paramagnetism due the presence of unpaired electrons in the molecule. Therefore, can calculate the ground state energy with different spin magnitude and see if ground state energy is lower with unpaired electrons. We present the experimental results in Table We observe that an oxygen molecule with 2 unpaired electrons achieves lower energy, whereas carbon dioxide is the opposite, in both methods. This coincides with the fact that oxygen molecule paramagnetic, while carbon dioxide is diamagnetic.
Table 2: Prediction on magnetism (LSDA, 6-31g, Ha)
No unpaired
2 unpaired electrons
No unpaired
2 unpaired electrons
Method
Method
electron
electron
PySCF
-148.021 088.02180 180
1448.9143 -148.0836
PySCF
-185.57994
-185.29154
CO2
D4FT 1488.4391
D4FT
118.582288
-185.26463
6.2 SCALABILITY
In this section, we evaluate the scalability of D4FT. The experiments are conducted on carbon Fullerene molecules containing ranging from 20 to 180 carbon atoms. For D4FT, we the Adam optimizer (Kingma & Ba| (014)) with a piecewise constant learning rate decay. The initial learning rate is 0.1, it is reduced 0.01 0.0001 after 60 epochs. In total, we run 200 epochs for each molecule. The results are shown in Fig.
When the system size is small, the acceleration brought by SGD limited. SCF has clear advantage as the closed-form eigendecomposition brings the so- lution very close to optimal in just one step. Both methods become on par when the system reaches a size of 480 orbitals (C80). After that, SGD is significantly more efficient than SCF because exact
II Larger molecule gets larger learning rate decay.
8

---

Published as a conference paper at ICLR 2023
Batch size
Batch size
102400
102400
-700
51200
-710
25600
Ineegy (Ila)
-715
12800
Fowega Fonggy (Ha)
-715
12800
6400
6400
-720
-720
3200
3200
-725
-725
-730
-730
-335
-735
10
15
20
Wall clock time (6)
Epoch
Figure 3: Convergence speed different batch sizes on Fullerene C20. Left: Comparison under wallclock time. Right: Comparison under the number of epochs. Smaller batch sizes are less favorable wallclock time because they can not fully utilize the GPU computation.
numerical integration is prohibitively expensive for large molecules. Column 4 in Fig. shows clear advantage of D4FT JAX-SCF, which verifies the smaller complexity of D4FT.
Influence of Batch Size As repeatedly proved in large-scale machine learning. SGD with minibatch has a more favorable convergence speed as compared to full batch GD. 3 provides evidence that this is also true for D4FT. Smaller batch size generally leads to faster convergence in terms of epochs over the grids. However, when the batch size gets too small to fully utilize the GPU's computation power, it becomes less favorable in wall clock time.
6.3 NEURAL BASIS
In this part, we test the neural basis with local scaling transformation presented in Sec. on atoms. In this experiment, we use simple STO-3g basis set. 3-layer MLP with tanh activation is adopted for Ga. The hidden dimension at each layer is 9. The experimental results (partial) are shown in Table Complete results are in Appendix E.3 The results demonstrate that the local scaling transformation effectively increases its flexibility of STO-3g basis set and achieves lower ground-state energy.
Table 3: Partial results of the neural basis method (LSDA, Ha)
Atom
He
Li
Be
C
N
0
STO-3g
-2.65731
-7.04847
-13.97781
36.47026 -36.52543
-52.82486 -52.86340
-72.78665
STO-3g + local scaling
-2.65978
-7.13649
-13.99859
-72.95256
7 DISCUSSION AND CONCLUSION
We demonstrate in this paper that KS-DFT can be solved in a more deep learning native manner. Our method brings many benefits including unified optimization algorithm, stronger convergence guarantee, and better scalability. It also enables us to design neural-based wave functions that have better approximation capability. However, at this point, there are still several issues that await solu- tions in order to move further along this path. One such issue attributes to the inherent stochasticity of SGD. The convergence of our algorithm is affected by many factors, such the choice of batch sizes and optimizers. Another issue is about the integral which involves fully neural-based wave functions. Existing numerical integration methods rely on truncated series expansions, e.g.. quadra- ture, and Fourier transform. They could face larger truncation error when more potent function approximators such as neural networks are used. Neural networks might not be helpful when we still rely linear series expansions for integration. Monte Carlo methods could be what we should resort to for this problem and potentially it could link to Quantum Monte Carlo methods. On the bright side, these limitations
the same time opportunities for machine learning researchers to contribute to this extremely important problem, whose breakthrough could benefit tremendous applications in material science, drug discovery, etc.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
We would like to extend our sincerest appreciation to all those who made contributions to this research project. We express our gratitude to Pengru Huang of the NUS Institute for Functional Intelligent Materials, Alexandra Carvalho and Keian Noori from the NUS Centre for Advanced 2D Materials, Martin-Isbjorn Trappe of the NUS Center for Quantum Technologies, and Aleksandr Rodin from Yale-NUS College for providing invaluable insights and feedback during the group discussions. Their willingness to share their experiences and perspectives enabled us to gain a deeper understanding of the current computational challenges related to the quantum many-body problem at hand. We thank our research engineer, Zekun Shi, for his assistance in improving the quality of code, paving the way towards high-quality package. We thank Jianhao Li from University of Minnesota Goodpaster Research Group for his valuable comments on the manuscript.
REFERENCES
Adam S Abbott, Boyi Z Abbott, Justin M Turney, and Henry F Schaefer III. Arbitrary deriva- tives of quantum chemical methods via automatic differentiation. The Journal of Physical Chem- istry Letters, 12(12):3232-3239 2021.
Milton Abramowitz and Irene A Stegun. Handbook of mathematical functions with formulas, graphs, and mathematical lables, volume 55. US Government printing office, 1964.
Mohammed Alghadeer, Abdulaziz Al-Aswad, and Fahhad H Alharbi. Highly accurate machine learning model for kinetic energy density functional. Physics Letters A, 414:127621, 2021.
Zhaojun Bai, Ren-Cang Li, and Ding Lu. Optimal convergence rate of self-consistent field iteration for solving cigenvector-dependent nonlinear eigenvalue problems. arXiv preprint arXiv:2009.09022. 2020.
Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311, 2018.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milee and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
7/gi lrgithuubccmm/gooo cotm/ooggggoo//aax
Yunfeng Cai, Lei-Hong Zhang, Zhaojun Bai, and Ren-Cang Li. On an degeevectorrdeppeddent non- linear eigenvalue problem. SIAM Journal on Matrix Analysis and Applications, 1300):1333--3888 2018.
Eric Cancès, Gaspard Kemlin, and Antoine Levitt. Convergence analysis of direct minimization and self-consistent iterations. SIAM Journal on Matrix Analysis and Applications, 42(1):243-274, 2021.
Yixiao Chen, Linfeng Zhang, Han Wang, and Weinan E. Ground state energy functional with fortree-fock efficiency and chemical accuracy. The Journal of Physical Chemistry A. 124(35): 7155-7165, 2020a.
Yixiao Chen, Linfeng Zhang, Han Wang, and Weinan E. Deepks: A comprehensive data-driven approach toward chemically accurate density functional theory. Journal of Chemical Theory and Computation, 17(1):170-181, 2020b.
Caio A Custódio, Érica R Filletti, and Vivian V França. Artificial neural networks for density- functional optimizations in fermionic systems. Scientific reports, 9(1):177, 2019.
Sebastian Dick and Marivi Fernandez-Serra. Machine learning accurate exchange and correlation functionals of the electronic density. Nature communications, 11(1):1-10, 2020.
Sebastian Dick and Marivi Fernandez-Serra Highly accurate and constrained density functional obtained with differentiable programming. Physical Review B, 104(16)LL161 109, 2021a.
10

---

Published as a conference paper at ICLR 2023
Sebastian Dick and Marivi Bermandez-Serra. Using differentiable programming to obtain energy and density-optimized exchange functional. arXiv preprint wrxiv22106.0488/1 202.,
RHWJ Ditchfield, W J_ Hehre, and John Pople. Self-consistent molecular-orbital methods. ix, an extended gaussian-type basis for molecular-orbita. studies of organic molecules. The Journal of Chemical Physics, 54(2)72247728 1971.
Thom H Dunning Jr. Gaussian basis sets for use in correlated molecular calculations. i. the atoms boron through neon and hydrogen. The Journal of chemical physics, 90(2):1007-1023, 1989.
J Austin Ellis, Lenz Fiedler, Gabriel A Popoola, Normand A Modine, J Adam Stephens, Aidan Thompson, Attila Cangi, and Sivasankaran Rajamanickam. Accelerating finite-temperature kohn- sham density functional theory with deep neural networks. Physical Review B, 104(3):035120, 2021.
S Alireza Ghasemi and Thomas D Kühne. Artificial neural networks for the kinetic energy functional of non-interacting fermions. The Journal of Chemical Physics, 154(7)0074107, 2021.
MJ Gillan. Calculation of the vacancy formation energy aluminium. Journal of Physics: Con- densed Malter, 1(4):689, 1989.
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, pp- 1263-1272. PMLR, 2017.
Gene H Golub and John H Welsch. Calculation of gauss quadrature rules. Mathematics of compu- tation, 23(100): 1969.
Hohenberg and WJPR Kohn. Density functional theory (dft). Phys. Rev, 136:B864, 1964a.
Pierre Hohenberg and Walter Kohn. Inhomogeneous electron gas. Physical review, 136(3B):B864, 1964b.
Kurt Hornik, Maxwell Stinchcombe, and Halhert White. Multilayer feedforward networks uni- versal approximators. Neural networks, 2(5):359-366, 1989.
Sohrab Ismail-Beigi and TA Arias. New algebraic formulation of density functional calculation. Computer Physics Communications, 128(1-2):1-45, 2000. 2000.
Aleksei V Ivanov, Elvar Ö Jónsson, Tejs Vegge, and Hannes Jónsson. Direct energy minimiza- tion based on exponential transformation in density functional calculations of finite and extended systems. Computer Physics Communications, 267:108047, 2021.
Bhupalee Kalita, Li Li, Ryan McCarty, and Kieron Burke. Learning approximate density func- tionals. Accounts of Chemical Research, 54(4):818 826, 2021.
M. F. Kasim and S. M. Vinko. Learning the exchange-correlation functional from nature with fully differentiable density functional theory. Phys. Rev. Lett, 127:126403 Sep 2021.
Agnan Kessy, Alex Lewin, and Korbinian Strimmer. Optimal whitening and decorrelation. The
American Statistician, 72(4):309-33 -140, 2018
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6988 2014.
James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander I. Gaunt, James Spencer, Alexander GDG Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, et al. Push- ing the frontiers of density functionals by solving the fractional electron problem. Science, 374 (6573); 1389, 2021.
Wolfram Koch and Max C Holthausen. A chemist's guide density functional theory. John Wiley & Sons, 2015.
W. Kohn and L. J. Sham. Self-consistent equations including exchange and correlation effects. Phys. Rev., A40:A1133-10 A138, Nov 1965.
11

---

Published as a conference paper at ICLR 2023
Andre Laestadius, Erik I Tellgren, Markus Penz, Michael Ruggenthaler, Simen Kvaal, and Trygve Helgaker. Shhm theory with paramagnetic currents: compatibility and functional differen- tiability. Journal of Chemical Theory and Computation, 15(7):403344085 2019.
Ira N Levine, Daryle H Busch, and Harrison Shull. Quantum chemistry, volume 6. Pearson Prentice
Hall Upper Saddle River, NJ, 2009.
Mario Sezcano-Casade and David Martinez-Rubio. Cheap orthogonal constraints in neural net- works: A simple parametrization of the orthogonal and unitary group. In International Confer- ence on Machine Learning, pp. 3794-3803. PMLR, 2019.
Li Stephan Hoyer, Ryan Pederson, Ruoxi Sun, Ekin D Cubuk, Patrick Riley, Kieron Burke, al. Kohn-sham equations as regularizer: Building prior knowledge into machine-learned physics. Physical review letters, 126(3):03640,, 2021.
Xin Liu, Xiao Wang. Zaiwen Wen, and Yaxiang Yuan. On the convergence of the self-consistent density functional theory. SIAM Journal on Matrix Analysis and
field iteration in Applications, 35(2)544-555, 2014.
Xin Liu, Zaiwen Wen, Xiao Wang. Michael Ulbrich, and Yaxiang Yuan. On the analysis of the discretized kohn-sham density functional theory. SIAM Journal on Numerical Analysis, 53(4): 1758 1785, 2015.
Alexander Mathiasen, Frederik Hvilshej, Jakob Rødsgaard Jørgensen, Anshul Nasery, and Davide Mottin. Faster orthogonal parameterization with householder matrices. In ICML, Workshop Pro- ceedings, 2020.
Ryo Nagai, Ryosuke Akashi, and Osamu Sugino. Completing density functional theory by machine learning hidden messages from molecules. npj Computational Materials, 6(1):1-8, 2020.
Robert G Parr. Density functional theory of atoms and molecules. In Horizons of quantum chemistry, pp. 5-1. Springer, 1980.
Alexander Ryabov, Iskander Akhatov, and Petr Zhilyaev. Neural network interpolation of exchange- correlation functional. Scientific reports, 10(1):1-7,200 2020.
Qiming Sun, Timothy C Berkelbach, Nick Blunt, George H Booth, Sheng Guo, Zhendong Li, Junzi Liu, James D McClain, Elvira R Sayfutyarova, Sandeep Sharma, et al. Pysef: the python-based simulations of chemistry framework. Wiley Interdisciplinary Reviews: Computational Molecular Science, 8(40,e1340, 2018
Attila Szabo and Neil S Ostlund. Modern quantum chemistry: introduction to advanced electronic structure theory. Courier Corporation, 2012.
Teresa Manayo-Mendoza, Christoph Kreisheck, Roland Lindh, and Alán Aspuru-Gurii Automatic differentiation in quantum chemistry with applications to fully variational hartree-fock ACS central science, 455:5559-566, 2018
Troy Van Voorhis and Martin Head-Gordon. A geometric approach to direct minimization. Molec- ular Physics, 100(1 1): 11:1))1111111 1713-1172 2002.
Joost Vande Vondele and Jürg Hutter. An efficient orbital transformation method for electronic struc- ture calculations. The Journal of chemical physics, 118(10):4365 43669 2003.
Valery Weber, Joost Vande Vondele, Jürg Hutter, and Anders MN Niklasson. Direct functional minimization under orthogonality constraints. The Journal of chemical physics, 128(8):084113, 2008.
Chao Yang, Juan C Meza, and Lin- Wang Wang. A trust region direct constrained minimization algorithm for the kohn-sham equation. SIAM Journal on Scientific Computing, 29(5): 88401185587 2007.
Chao Yang, Weiguo Gao, and Juan C Meza. On the convergence of the self-consistent field iteration for a class of nonlinear eigenvalue problems. SIAM journal on matrix analysis and applications, 3774):173-11788, 2009.
12

---

Published as a conference paper at ICLR 2023
Appendices
A NOTATIONS
A.1 BRA-KET NOTATION
The bra-ket notation, or the Direc notation, is the conventional notation quantum mechanics for denoting quantum states. It can be viewed way to denote the linear operation in Hilbert space. ket e) represent a quantum state in V, whereas bra (3 represents a linear mapping: C. The inner product of two quantum states, represented by wave functions 3 and dy can be rewritten in the bra-ket notation as
(bl1))) := / vi (T)
(22)
Let Ô be an operator, the expectation of Ô is defined as,
( (v) = / 16 (r))())()))) dr.
(23)
A.2 NOTATION TABLE
We list Trequentlyuseed notations used in this paper the next Table
Table 4: Notation used in this paper
Notation
Meaning
R
real space
complex space
coordinate in 1
eb
single-particle wave function//moeculaa orbitals basis function/atomic orbitals
L
vectorized wave function: IR3 - IR
N
number molecular orbitals
B
number of basis functions
grid point pair (coordinate and weight) number of grid points for numerical quadrature hamiltonian operator
H
electron density overlap matrix spin index energy functionals
B
Egs. Ekin.
13

---

Published as a conference paper at ICLR 2023
B KOHN-SHAM EQUATION AND SCF METHOD
B.1 KS-DFT OBJECTIVE FUNCTION
The Schrödinger Equation The energy functional of KS-DFT can be derived from the quantum many- many-body Schrödinger Equation. The time-invariant Schrödinger Equation writes
V) >
(24)
where A is the system hamiltonian and (r) denotes the many-body wave function mapping R C. is vector of N 3D particles. The Hamiltonian is
A = V2 ΣΣ
(25)
R
ri Tj
kinetic
external
electron repulsion
where r; is the coordinate of the i-th electron and Rj is the coordinate of the j-th nuclei. The follow- ing objective function is optimized obtain the ground-state energy and corresponding ground-state wave function:
min ( Mr (
(26)
s.t. ((r) (P)V(PP); VP.
(27)
Antisymmetric Constraint Notice there is a constraint in the above objective function, P is any permutation matrix, and (P) denotes the parity of the permutation. This constraint comes from the Pauli exclusion principle, which enforces the antisymmetry of the wave function Psi. When solving the above optimization problem, need variational ansatz, in machine learning terms, function approximator for V. To satisfy the antisymmetry constraint, we can start with any function approximator V, and apply the antisymmetrizer V AV, the antisymmetrizer is defined as
1 VN!
Σ (P)))(P))
(28)
However, there are N! terms summing over P, which is prohibitively expensive. Therefore, the Slater determinant is resorted as much cheaper approximation.
Slater Determinant & Hartree Fock Approximation To efficiently compute the above antisym- metrizer, the mean-field assumption is usually made, i.e. the many-body wave function is made of the product of kiggleparticle wave functions:
(ri)
(12)
s (rN)
(2(rr) (2(122)
(rNN)
slater AW
AI[ di
(29)
pbN (r1)
(x())) dew (rx))
This approximation has much more favorable computation complexity, computing the determinant only takes O(N3). However, it is more restricted in its approximation capability. Notice that the mean-field assumption discards the correlation between different wave functions. Therefore, with the Slater determinant approximation, it omits the correlation between electrons. Plugging this into 26 gives us the Hartree-Foct approximation.
Orthogonal Constraints While the Slater determinant is much cheaper to compute, integration in the Nx3 space is still complex. Introducing an orthogonal between the single-particle wave functions simplifies things greatly.
(Wl1))) 8zg-
(30)
14

---

Published as a conference paper at ICLR 2023
Plugging this constraint and the Slater determinant into Eq. 26 and with some elementary caleulus, there are the following useful conclusions. From this point on, without ambiguity, the symbol r is also used to denote vector in 3D space when it appears in the single-particn wave function (r)).
First, the total density of the electron is the sum of the density contributed by each wave function. With this conclusion, the external part of the hamiltonian from the many body Schrödinger Equation connects with the external energy DFT in Eq. 3
N
p(r) (Wl))) 88
Σla()))
(31)
Second, the kinetic energy of the wave function in the joint N x 3 space breaks down into the summation of the kinetic energy of each angle-particle wave function in 3D space, which again equals the kinetic term in Eq.
V slater
22
V slattr? (th (w1l)))b
Bis Σ (whel
(32)
Third, the electron repulsion term breaks into two terms, one corresponding to the hartree energy in Eq. and the other is the exact exchange energy.
V stater Eies
V slater)
(33)
rr
ryl
1
dr,dry dr
(34)
2
(r r2
Σ
1 // 1 (m)))) (r r2l XYS
dr dr
(35)
Ex
The above three conclusions link the Schrödinger Equation with the three terms in the KS-DFT objective. To finish the link, notice that when Slater determinant is introduced, we made the mean- field assumption. This causes an approximation error due the electron correlation in the Hartree- Fock formulation we're deriving. In DFT, this correlation error, together with the exchange energy Fx is approximated with functional Exc- The art of DFT thus find the best Exc that minimizes the approximation error.
B.2 DERIVATION THE KOHN-SHAN EQUATION
The energy minimization in the functional space formally reads
min
{ Egs (#1,.... (bn) g/y E H (I=*) (th,, bby) = 8eg, 1 & i,j < N}.
(36)
Vi
where we can denote the electron density as
((r)= = Σ (s(r))
(37)
and have the objective function to be minimized:
N =
1 Σ / (t(r)) (VPv)(r)) dr + Σ / Vext (1))))
/ (E ((((r)) (r')12
1
drdr' + Σ / Uxe (T)
15

---

Published as a conference paper at ICLR 2023
This objective dunction can be solved by Euler-Lagrangian method. The Lagrangian is expressive as,
(1 vi
L , NN) E (1/, ... U2))
Σ
Eij
(38)
s/==
The first-order condition of the above optimization problem can be writen as,
5L
0,
i i= NN
8
(39)
ac
0,
i,j = 1, N
Jeij
The first variation above can be written as,
N
8L
8
Faattree
F 8p
ōp
+
Ety
(40)
8p
8p
8p
Next, derive each components in the above equation. We follow the widely adopted tradition that treats of and 1/s as different functions Parr 1980), and have
6p
(41)
8uh
For EHartes according to definition of functional derivatives:
(o 6p) + 5p) 5p)) - p(r))(r')
Filanree ip + 6p]
El (p] = 2
drdr'
r--rl / (pp)(r)p(r) + (rp)(r')p(r) drdr'
2 /
---
//
(sp)(r))(r') drdr'
(42)
/ (op)(r))(r')
Ir -r- drdr'
rl)
((p)(T) (f
p(r')
|r |-- dr' dr,
where we have used the fact that:
ʃ f((r))2(rr) drdr'= = = fi((') |r 1r-r' rrl
-r-r"
drdr',
(43)
due the symmetry between and r'. We emphasize that the functional derivatives are computed on functionals mapping functions to scalars, and consequently, we need to compute that for the double integral on both r and r'.
Therefore, the functional derivative for L'Hartree can be rewritten into:
E [p]
/
p(r')
r rd dr'.
(44)
6p
Using the chain rule, we have
8 F (p)
p(r')
Op
dr'
8rp²"
r pil
(45)
n((r')
dr'
r rd
For Frs and Fxc:
8 F
5 FEST
Op
Uext 1/:-
83/
5p Tei
(46)
E
δ 0p
Uxct/i.
5p 0y.
16

---

Published as a conference paper at ICLR 2023
For the kinetic energy.
Figi d
2
(47)
8447
2
Taking first-order derivatives, the Lagrangian becomes:
H (VI..... ,(, v A,di.
(48)
where A = Ef== Eij. and the Hamiltonian is explicitly written as
A -
22 Vext (r)
p(r')
dr' + Uxe (r);
(49)
R2 r'
N
(()) =
(x((r)
(50)
which is the desired form of the Kohn-Sham Equation.
B.3 THE SCF ALGORITIM.
The connection between SCF and direct optimization can be summarized as follows. D4FT is min- imizing the energy directly where the constraints are encoded into the corresponding parameter constraints, while SCF deals with the constrained optimization by Lagrangian. They are solving same minimization problem using different approaches, which is reflected in how they deal with the constraint in the problem.
Obviously, D4FT is doing projected gradient descent on the constraint set क. SCF is the correspond- ing Lagrangian, where the Lagrangian enforces the orthogonal constraint. Another constraint on the basis function is reflected in the inner product of SCF.
The SCF algorithm is presented Algo.
Algorithm Self-consistent Field Optimization for Kohn-Shant Sham Equation
Input: set of single particle wave function {v/s). convergence criteria c.
Output: ground state energy, electron density
1: Compute the initial electron density / via Eq. 50
2: while True do
3: Update Hamiltonian via Eq. 49
4: Solve the eigenvalue equations defined in Eq. 48 and get new {th))
5: Compute electron density p via Eq 50
6: if [peew p) then
7: break
8: end if
9: end while
17

---

Published as a conference paper at ICLR 2023
C PROOFS AND DERIVATIONS
C.I PROOF OF PROPOSITION 3.1
Proof. We first prove the following statement as a preparation: (orthogonal constraint) for any (4/1, 22.22... )N)' E Fw. it holds that (wal)),) d.j for all (i,j). Let (1h, 122. ;N)) From the definition,
T E Fw-
(T) = (r),(r,
where Q, is the i-th row vector of Q and (r) = (.(
Thus, since
U I, we have that for any (i.j) [A] x [
Q: 0J - - UEU U2 - QQQ] = -ig-
This proves the statement that for any (4/1, 2/2,. ... dx)T E Fw, it holds that (1b1113)) 8g for all (i.j). This statement on the orthogonal constraint and the definitions of Fw and F implies that
Fw S F. Thus, in the rest of the proof, we need to prove Fw 2 This is equivalent to
(QDT (Q, R) = (R(W), W E RRNN {C : C E CRNN C,BC, = 6:g):
where C: the i-th row vector of C. This is implied if the following holds:
(QD : (Q, R) = QR(W), W E NNNN } {C: C C CRNN C,BC, = 6-z}- (51) Since D is non-singular, for any E (C: CERNNN C.BCJ = dτ) there exists Ö NxN
such that C = CD (and thus C,BC ČDBD Cj diz). Similarly, for any CD E {CD:
Ć E R N CC,DDDTG CJ = d;j), there exists C E RNN such that C = CD (and thus C,BC -.z). Therefore,
[C:C ,C, = = (CD: CCD:CERMX C,DBD'S T CJ = 6zy)
Here,
UEU US
U/C = C.C.. Thus,
{C: C E NNN, C,BCJ = C N C.C, dby.
Substituting this into Eq. 51 the desired statement is implied if the following holds:
(QD: (Q, R) QR(W), W N (CD: C NNN ,C,CJ =
which is implied by
(Q: (Q. QR((Q)))))))) W NNN 2 F C.C,J dgj.
(52)
Finally, we note that for any C E {C E RNN, C.C. = 8ig1, there exists W E NxN such that Q = C where (Q, R) QR(W). Indeed, using the Gram Schmidt process of the QR decomposition, for any C {C RNN C,C, = d, setting W = = RNN suffices to obtain Q C where (Q, R) QR(W). This implies Eq. 52 which implies the desired statement.
C.2 PROOF PROPOSITION
Proof. The total energy minimization problem defined in Eq. 3 is a constraint minimum problem, which can be rewritten as,
Egs
HH (R*)
(1 = =
(53)
with 2 denoted as Sobolev space, which contains all T,2 functions whose first order derivatives are also L integrable. People resort to the associated Euler-Lagrange equations for solution:
L vV) = Eags V
Σ
((ba)))) - 5zz)
$
Letting the first-order derivatives equal to zero, i.e., 5C = 0, the Lagrangian becomes: Á (L) t/ dwī
Xodr, where A Σ5- Eij. and the Hamiltonian H is the same as defined in Eq. 48 To solve
this nonlinear eigenvalue/eigenfunction problem, SCF iteration in Algorithm is adopted. The
18

---

Published as a conference paper at ICLR 2023
output of the (t + 1)th iteration is derived from the Hamiltonian with the previous output, i.e., i (ID w// x The convergence criterion is प = Tpt+1 which is the self-consistency formally defined in Definition 3.1
Since GD and SCF are solving essentially the same problem, we can derive the equivalence be- tween the solutions of GD and SCF. GD is doing projected gradient descent on the constraint set ( ( ) Т]. SCF is the corresponding Lagrangian, where the Lagrangian enforces the or- thogonal constraint.
19

---

Published as a conference paper at ICLR 2023
D COMPLEXITY ANALYSIS
D.1 COMPLEXITY ANALYSIS FOR FULL BATCII
SCF's computation complexity analysis
I p(r1)
1 m 12
vi (ra))
1. Compute (r) for E [1..N] and the n grid points. takes () to compute for each r, and we do it for n points. Therefore, it takes On, producing RR matrix.
2. Compute p(r), it takes an extra O(nN) step to sum the square of the wave functions, on top of step produces I² vector.
3. Compute v (r), it takes an outer product on top step O(nN3), produces " N tensor.
4. Compute
Onn²).
5. Tensor contraction outcomes from step 2, 4. On?
6. In total () + = N33
SGD's s computation complexity analysis
[ plr) ri r2) p p(r2)dridd
1. Using steps 1, 2, and 4 from SCF. Step 3 is not used because only energy is needed, O(nNB)
2. Contract R", R", Rexn into a scalar, O(2²).
3. In total OnN O(n²) O(nN2) On)²).
D.2 COMPLEXITY ANALYSIS FOR MINIBATCH
The complexity of SCF under minibatch size is O(mNB) + Omm2N2), since the number of iteration per epoch is m, the total complexity for one epoch is (nnN + = O(n(1+ ?) since B N, i.e., the number of orbitals equal that of the basis.
Similarly, the complexity of GD is (O(nNB) + O(m²))n/m (ni(N2 + m)).
D.3 LITERATURE ON THE CONVERGENCE OF SCF METHODS.
There is rich literature the convergence of SCF, Cancès al. 12021) provides a literature review and method summarization for the iterative SCF algorithm. Specifically they show that the problem be tackled via either Lagrangian SCF direct gradient descent matrix manifold. Yang et al. (2077) shows that the Lagrangian + SCF algorithm indirectly optimizing the original energy by minimizing sequence of quadratic surrogate functions. Specifically, the fth time step, SCF minimizing the quadrature surrogate energy using the fth Hamiltonian, min A 2 Tr(A* H A), while the correct nonlinear energy should be min 1 Tr (T(AH((AAA) So, our GD directly mini- mizes the nonlinear energy, while SCF deals with a quadrature surrogate energy. Using a concrete wo-dimensionall example, how SCF may fail converge illustrated, since the second order ap- proximation of the original. Yang et al. (2009) shows that SCF will converge to two different limit points, but neither of them is the solution to particular class of nonlinear eigenvector problems. Nevertheless, they also identify the condition under which the SCF iteration becomes contractive which guarantees its convergence, which is the gap between the occupied states and unoccupied states is sufficiently large and the second order derivatives of the exchange-correlation functional are uniformly upper bounded. Bai et al. (2020) provides optimal convergence rate for class of nonlinear eigenvector directly minimizesly; since SCF is iterative algorithm, show that
20

---

Published as a conference paper at ICLR 2023
SCF is contraction mapping:
d(A(t+1), A*)) ,²) d(A40), -A) - 1)s
sup
(54)
t-000
where d measures the distance between two matrices, A is the solution, and 1j is the contraction factor determining the convergence rate of SCF, whose upper bound is given in their result. Liu et al.
(2015 2014) formulates SCF as a fixed point map, whose Jacobian is derived explicitly. from which the convergence of SCF is established, under similar condition to the one in Yang et al. (2009).
Cai et al. (2018) derive nearly optimal local and global convergence rate of SCF
21

---

Published as a conference paper at ICLR 2023
E
MORE EXPERIMENTAL RESULTS
E.1
FULL VERSION OF TABLE
Table 5: Comparison of Ground State Energy (LDA, 6-31g, Ha)
Nuclear Repulsion Energy
Kinetic+ Externel Energy
Hartree Energy
XC Energy
Total
Molecule
Method
Energy
PySCF
0.71375
22.88017
1.28032
-0,55256
-103866
Hydrogen
D4FT
0.71375
-2.48343
1.28357
-0.55371
-1.03982
JaxSCF
0.71375
-2.48448
1.28527
-0.55390
-1.03919
PySCF
13,47203
-79.79071
32.68441
-5.86273
-39,49700
Methane
D4FT
13.47203
-79.77446
32.68037
-5,59499
-39.48155
JaxSCF
13.47203
-79.80889
32.71502
-5.66561
33.448745
PySCF 9.18953
-122.83945
46,58941
-8.93311
-75,15381
Water
D4FT
9.18953
-122.78656
46.54537
-8.08588
-75.13754
JaxSCF
9.18953
-122778241
46,53507
-8.09333
-75.15124
PySCF
28.04748
2661.1997
99.92152
14..79109
148,02180
Oxygen
D4FT
28.04748
2261.13046
99.82705
14.778807
148.03399
JaxSCF
28.04748
2661.16314
99,87551
-14.78290
-488023304
PySCF
82.01074
-371.73603
156.36891
-18.65741
-152.01379
Ethanol
D4FT
82.01074
-371.73431
156.34982
-18665517
-152.02893
JaxSCF
82.01074
-371.64258
156.27433
18.655711
-152,01460
PySCF
203.22654
-713.15807
312.42026
299.84784
-227.35910
Benzene
D4FT 203.22654
-712.71081
311.94665
-29.81097
-227.34860
JaxSCF
203.22654
-712.91053
312.15690
-29,84042
-227.36755
E.2 FULL VERSION OF TABLE
Table 6: Comparison of ground state energy atoms (LSDA, STO-3g, Hartree).
Atom
He
Li
Be
B
C
PySCF
-2.65731
-7.06641
-13.98871
-23.66157
-36.47199
D4FT
-2.65731
-7.04847
-13.97781
-23.64988
-36.47026
NBLST
-2.65978
-7.13649
-13.99859
-23.67838
-3,.52543
Atom
N
O
F
Ne
Na
PySCF
-52.80242
-72.76501
99.93644
-125,38990
-158,40032
DAFT
-52.82486
-72.78665
-96,99325
1623.
1188.16140
NBLST
-52.86340
-72.95256
997.24058
-125.62643
-157.72314
Atom
Mg
AI
Si
P
S
PySCF
-195.58268
-237.25091
-283.59821
-334.79401
-390.90888
D4FT
-195,47165
-237.03756
-283.41116
1035,10473
-391,09967
NBLST
-196.00629
-235.83420
-283.72964
-335.38168
-391.34567
E.3 MORE RESULTS ON THE EFFICIENCY COMPARISON
We test the efficiency and scalability of our methods against couple of mature Quantum Chemistry softwares. We compute the gtaund-state energy using each implementation and record the wall- clock running time. The softwares compare with are:
22

---

Published as a conference paper at ICLR 2023
PySCE. PySCF is implemented mainly with C (Qint and libxc), and use python the interface.
GPAW. A Plane-wave method for solid and molecules.
Psi4. Psi4 is mainly written C++.
We use sto-3g basis and Ida exchange-correlatie functional. For GPAW, apply the Icao mode, and use dzp basis set. The number of bands calculated the double of the number of electrons. The number of grids for evaluating the XC and Coulomb potentials 400*400*400 by default. The maximum number of iteration is 100 for all the methods. Our D4FT ran on an NVIDIA A100 40G GPU, whereas other methods did on Intel Xeon 64-core CPU with 60G memory. The wall-clock running time for convergence (by default setting) are shown in the following table.
Table 7: Comparison of wall-clock running time (s).
Molecule
C20
C60
C80
C100
C180
D4FT
11.46
23.14
62.98
195.13
1103.19
PySCF
20.75
186.67
>3228.20
672.98
1925.25
GPAW 27783.08
Psi4
>4..12
510.35
22144.49
2555.16
14321.64
Notation represent that the experiment cannot finish converge given the default convergence condition. can be seen from the results that our D4FT is the fastest among all the packages.
23

---

Published as a conference paper at ICLR 2023
E.4 CASES THAT PYSCF FAILS.
We show two cases that do not converge with PySCE. We test two unstable systems, C120 and C140, which do exist in the real world. Both of them truncated from C180 Fullerene molecule. Results are shown in Fig.
PySCF on C120
D4FT on C120
-3200
-3400
-1000
-3600
energy
-2000
energy
-3800
-4000
-3000
-4200
44000
-4400
0
20
40
60
80
100
50
100
150
200
iteration
iteration
PySCF on C140
D4FT on C140
-1500
-3800
-2000
-4000
-2500
-4200
-3000
energy
energy
-4400
-3500
-4600
-4000
-4800
-4500
-5000
-5000
-200
20
40
60
80
100
50
100
150
200
iteration
iteration
Figure 4: Cases that do not converge with PySCF, but do with D4FT. Due to the good convergence property SGD, D4FT is more likely to converge.
24


Reviews:
[{"review_id": "xVIRTvesFkM", "reviewer": "Reviewer_fw91", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this work, a deep learning approach for solving KS-DFT was proposed, which converts the objective function for KS-DFT into an unconstrained equivalent by reparameterizing the orthogonal constraints as a feed-forward computation. By using stochastic gradient descent, the integral was amortized over the optimization steps. Its equivalence to the conventional SCF method was proven both empirically and theoretically. The comparison between this work and some other deep-learning-assisted DFT approaches were discussed. The accuracy and scalability were tested via experiments on molecules by comparing with two benchmarks (PySCF and JAX-SCF). The authors further raised several issues that await solutions based on the limitaions from this paper.", "summary_of_the_review": "The paper developed a deep learning approach for solving KS-DFT by converting the objective function for KS-DFT into an unconstrained equivalent by reparameterizing the orthogonal constraints as a feed-forward computation. Its equivalence to the conventional SCF method is proven both empirically and theoretically. The presentation in the manuscript is clear. I recommend it for publication if authors address the above mentioned comments and offer some more substantial experimental data.", "strength_and_weaknesses": "Strength:\n\nThe approach proposed in this paper was proven to has the same expressivity as the SCF method, while reducing the computational complexity from O(N4) to O(N3). The methodology is expected to be also suitable for more complex neural-based wave functions.\n\nWeaknesses:\n\n1. For a fair comparison with the same software/hardware environment, the authors reimplemented the SCF method in JAX. Nevertheless, its equivalence with the conventional SCF solvers in terms of efficiency and scalability was not clarified.\n\n2. To demonstrate the efficiency and stability, the authors carried out some calculating experiments. However, the coverage of experiments was not sufficient. When evaluating the scalability of D4FT, carbon Fullerene molecules containing ranging from 20 to 180 carbon atoms were simulated. The size of the system seems to be not large enough. As the main limitation of the conventional SCF is its inefficiency for large molecules or solid cells, a machine learning based method is expected to behave better on large scale systems. This part of data might be implemented to show its advantage in reducing the computational complexity from O(N4) to O(N3).\n\n3. The approach raised was claimed to bring a stronger convergence guarantee compared with conventional SCF, while this was not supported by a related instance. A successfully converged calculation done by D4FT which failed in conventional SCF would be valuable.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "s6VA9IGh6L", "reviewer": "Reviewer_GibD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a deep learning model for Kohn-Sham density functional theory. Unfortunately, I know nothing about this theory and related works, so I cannot accurately summaries this work. ", "summary_of_the_review": "Unfortunately, I am not in this area. ", "strength_and_weaknesses": "Unfortunately, I am not in this area so cannot identify the possible strength and weaknesses. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "BTeLbEFBxW", "reviewer": "Reviewer_BHRj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the computation of wave functions in Kohn-Sham Density Functional Theory (KS-DFT). Differently from conventional works, the authors use neural network (NN) models with a stochastic optimization method in machine learning to tackle the problem. The problem is formulated as the minimization problem of energy function depending on the set of orthogonal wave functions. The integral in the energy function is approximated by the finite sum based on the numerical quadrature. The stochastic gradient descent (SGD) is then employed to approximate the optimization for the objective function consisting of a finite sum. The parameter transformation with the Jacobian is introduced to deal with the orthogonal constraint among wave functions. The neural network model is used as the model for the transformation. Numerical experiments show that the proposed method provides comparable accuracy and high scalability to exciting methods.", "summary_of_the_review": "The paper is clearly written. The neural basis function is an interesting model to deal with the problem considered in the paper. However, the theoretical insight is limited. Showing other applications of neural basis function would be beneficial for readers. ", "strength_and_weaknesses": "Strength\n- Overall the paper is clearly written. The significance of the problem is well introduced though the problem setting may not be well-known in the machine learning community. The modeling of the local scaling transformation using NNs is interesting. \n\n Weaknesses\n- There are some unclear points in the algorithm. The Jacobian appears in (19), which depends on a NN model. It is unclear how to compute the Jacobian matrix efficiently. I guess the back-propagation for NNs is applied to calculate the derivative w.r.t. the input vector instead of the weight parameters. A detailed description of the algorithm would be helpful to readers. \n- The authors provide theoretical contributions as some propositions, but the theoretical results are relatively straightforward. \n- Numerical experiments are not sufficient. Hyper-parameter tuning is important to achieve high accuracy in the computation. In the discussion, the author pointed out, \"One potential remedy for this problem is the second-order optimizers.\" However, the current paper does not provide a systematic study of that problem. Ablation studies should be provided to assess which factor is significant to attain satisfactory results. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%