# M-OFDFT: Overcoming the Barrier of Orbital-Free Density Functional Theory
for Molecular Systems Using Deep Learning

He Zhang1,2†♯, Siyuan Liu2†♯, Jiacheng You2♯, Chang Liu2∗, Shuxin Zheng2∗,
Ziheng Lu2, Tong Wang2, Nanning Zheng1, Bin Shao2∗.
1 Xi’an Jiaotong University.
2 Microsoft Research AI4Science.
∗ Corresponding authors. E-mails: {chang.liu, shuxin.zheng,
<EMAIL_ADDRESS>
† These authors contributed equally to this work.
♯ These authors did this work during an internship at Microsoft Research
AI4Science

###### Abstract

Orbital-free density functional theory (OFDFT) is a quantum chemistry
formulation that has a lower cost scaling than the prevailing Kohn-Sham DFT,
which is increasingly desired for contemporary molecular research. However,
its accuracy is limited by the kinetic energy density functional, which is
notoriously hard to approximate for non-periodic molecular systems. In this
work, we propose M-OFDFT, an OFDFT approach capable of solving molecular
systems using a deep-learning functional model. We build the essential
nonlocality into the model, which is made affordable by the concise density
representation as expansion coefficients under an atomic basis. With
techniques to address unconventional learning challenges therein, M-OFDFT
achieves a comparable accuracy with Kohn-Sham DFT on a wide range of molecules
untouched by OFDFT before. More attractively, M-OFDFT extrapolates well to
molecules much larger than those in training, which unleashes the appealing
scaling for studying large molecules including proteins, representing an
advancement of the accuracy-efficiency trade-off frontier in quantum
chemistry.

##  1 Introduction

Density functional theory (DFT) is a powerful quantum chemistry method for
solving electronic states hence energy and properties of molecular systems. It
is among the most popular choices due to its appropriate accuracy-efficiency
trade-off, and has fostered many scientific discoveries [1, 2]. The prevailing
Kohn-Sham formulation (KSDFT) [3] solves a system by minimizing the electronic
energy as a functional of $N$ orbital functions
$\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$, where $N$ is the number of electrons.
Although the orbitals allow explicitly calculating the non-interacting part of
kinetic energy, optimizing $N$ functions deviates from the original idea of
DFT [4, 5, 6, 7] to optimize one function, the (one-body reduced) electron
density $\rho(\mathbf{r})$, hence immediately increases the cost scaling by an
order of $N$ (Fig. 1(a)). This higher complexity is increasingly undesired for
the current research stage where large-scale system simulations for practical
applications are in high demand. For this reason, there is a growing interest
in methods following the original DFT spirit, now called _orbital-free DFT_
(OFDFT) [8, 9, 10].

The central task in OFDFT is to approximate the non-interacting part of
kinetic energy as a density functional (KEDF) $T_{\textnormal{S}}[\rho]$.
Classical approximations are developed based on the uniform electron gas
theory [11, 12, 13, 14, 15], and have achieved many successes of OFDFT for
periodic material systems [16, 17, 10]. But the accuracy is still limited for
molecules [18, 19, 20, 21], mainly due to that the electron density in
molecules is far from uniform.

For approximating a complicated functional, recent triumphant progress in deep
machine learning creates new opportunities. Yet, existing explorations for
OFDFT are still in an early stage. These methods use a regular grid to
represent density as the model input, which is not efficient enough to
represent the uneven density in molecular systems. Even an irregular grid
requires unaffordably many points for a nonlocal calculation, while the
nonlocality has been found indispensable to approximate KEDF [22, 14, 8, 23].
As a result, these works only studied molecules of up to a dozen atoms, either
due to the unaffordable cost of a nonlocal calculation [24, 25, 26, 27, 28] or
limited accuracy of a (semi-)local approximation [29, 30, 31, 32]. Moreover,
few work showed the accuracy on molecules much larger than those in training,
but such an extrapolation study is imperative as it is on molecules larger
than other methods could afford to generate abundant data that an OFDFT method
could demonstrate the dominating value of its scaling advantage.

In this work, we develop an OFDFT method called M-OFDFT that can handle
Molecules using a deep-learning KEDF model. To account for the nonlocal nature
of KEDF with affordable cost, we take the expansion coefficients of the
density on an atomic basis set as the model input (Fig. 1(b)), which
constitute a much more concise representation than a grid-based
representation. Each coefficient represents a density component around an
atom, and can be treated as a feature associated to that atom. To process such
input, we build a deep-learning model based on the Graphormer architecture
[33, 34], a variant of Transformer [35] for processing molecular data. The
model iteratively processes features on each atom, with the interaction with
features on other atoms through the attention mechanism, which covers the
nonlocal effect. We note that learning a functional model faces unconventional
challenges, for which we propose method to generate multiple density
datapoints with gradient labels per molecular structure, and techniques to
handle geometric invariance and vast gradient range. After the KEDF model is
learned, M-OFDFT solves a given molecular system by optimizing the density
coefficients, where the KEDF model is used to construct the energy objective
(Fig. 1(c)).

We demonstrate the practical utility and advantage in the following aspects.
(1) M-OFDFT achieves _chemical accuracy_ compared to KSDFT on a range of
molecular systems in similar scales as those in training. This is hundreds
times more accurate than classical OFDFT. The optimized density shows a clear
shell structure, which is regarded challenging for an orbital-free approach.
(2) M-OFDFT achieves an attractive _extrapolation capability_ that its per-
atom error stays constant or even decreases on increasingly larger molecules
all the way to 10 times (224 atoms) beyond those in training. The absolute
error is still much smaller than classical OFDFT. In contrast, the per-atom
error keeps increasing by end-to-end energy prediction counterparts. M-OFDFT
also shows a more efficient utilization of a few large-scale data after
trained on abundant affordable-scale data. (3) With the accuracy and
extrapolation capability, M-OFDFT unleashes the scaling advantage of OFDFT to
large-scale molecular systems. We find its empirical time complexity is
$O(N^{1.46})$, indeed lower by order-$N$ over $O(N^{2.49})$ of KSDFT. The
absolute time is always shorter, achieving a 27.4-fold speedup on the protein
B system (2,750 electrons). In all, M-OFDFT pushes the accuracy-efficiency
trade-off frontier in quantum chemistry, and provides a powerful tool for
solving large-scale molecular science problems.

Figure 1: Overview of M-OFDFT. (a) Kohn-Sham DFT solves the properties of a
molecular structure $\mathcal{M}$ by optimizing $N$ orbital functions
$\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$, while orbital-free DFT only needs to
optimize one density function $\rho(\mathbf{r})$ if the kinetic energy density
functional (KEDF) $T_{\textnormal{S}}[\rho]$ is available, which reduces the
complexity by an order of $N$. (b) The proposed M-OFDFT uses a deep-learning
model to approximate KEDF, which is learned from data. The model incorporates
nonlocal interaction of density over the space, which is made affordable by
inputting the concise density representation of expansion coefficients
$\mathbf{p}$ on an atomic basis $\\{\omega_{a,\tau}(\mathbf{r})\\}_{a,\tau}$.
Each basis function concentrates around an atom, and they altogether span a
similar pattern as the density, making the representation concise. (c) M-OFDFT
solves a molecular structure $\mathcal{M}$ by optimizing the density
coefficients $\mathbf{p}$, for which the learned KEDF model
$T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$ constitutes the energy
objective.

##  2 Results

###  2.1 Overview of M-OFDFT

OFDFT solves the electronic state of a molecular structure $\mathcal{M}$ by
minimizing the electronic energy as a functional of the electron density
$\rho$, which is typically decomposed in the same way as KSDFT:
$E[\rho]=T_{\textnormal{S}}[\rho]+E_{\textnormal{H}}[\rho]+E_{\textnormal{XC}}[\rho]+E_{\textnormal{ext}}[\rho]$,
where $T_{\textnormal{S}}[\rho]$ is the kinetic energy density functional
(KEDF) covering the non-interacting kinetic energy, $E_{\textnormal{H}}[\rho]$
covers the classical internal potential energy, the exchange-correlation (XC)
functional $E_{\textnormal{XC}}[\rho]$ accounts for the rest of kinetic and
internal potential energy, and $E_{\textnormal{ext}}[\rho]$ is the external
potential energy (Supplementary Sec. A.1). Terms $E_{\textnormal{ext}}[\rho]$
and $E_{\textnormal{H}}[\rho]$ have exact expressions, and
$E_{\textnormal{XC}}[\rho]$ already has accurate approximations. As for the
non-interacting kinetic energy, in KSDFT it can be calculated from orbital
solutions, but as a density functional for conducting OFDFT, the expression
$T_{\textnormal{S}}[\rho]$ is unknown and requires an accurate approximation.

The proposed M-OFDFT uses a deep machine learning model to approximate the
KEDF (Fig. 1(b)). For an efficient density representation to allow a nonlocal
architecture, we adopt an atomic orbital basis set
$\\{\omega_{\mu}(\mathbf{r})\\}_{\mu=1}^{M}$ to expand the density
$\rho(\mathbf{r})=\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}(\mathbf{r})$, and
take the coefficients $\mathbf{p}$ as the model input. Each basis function
$\omega_{\mu=(a,\tau)}(\mathbf{r})$ depicts a density pattern $\tau$ around
atom $a$, which aligns with the pattern of electron density in a molecule
where electrons distribute around atoms. Moreover, they are designed to mimic
the nuclear cusp condition [36] for sculpting the sharp density change near a
nucleus. They also naturally form a shell structure, _i.e._ , concentrate on
different distances from the center atom. These features further fit the
details of density in molecules, facilitating the representation with high
efficiency. As a result, commonly required basis functions is thousands times
fewer than grid points.

Under this representation, the KEDF model follows the form
$T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$, where $\theta$ is the
learnable parameters, and the molecular structure
$\mathcal{M}:=\\{\mathbf{X},\mathbf{Z}\\}$ is required for specifying the
locations and types of basis functions, where
$\mathbf{X}:=\\{\mathbf{x}^{(a)}\\}_{a=1}^{A}$ and
$\mathbf{Z}:=\\{Z^{(a)}\\}_{a=1}^{A}$ are the coordinates and atomic numbers
(types) of the atoms (conformation and constitution). As each coefficient
$\mathbf{p}_{\mu=(a,\tau)}$ can be associated to one atom $a$, the input
$(\mathcal{M},\mathbf{p})$ is a set of pinpointed atoms each with a type and
density coefficient features (Fig. 1(b)). To process such input, we build a
graph neural network based on the Graphormer architecture [33, 34], which
improves graph-theoretical expressiveness over Transformer [35] by
incorporating pairwise features (_e.g._ , distance features) into the
attention mechanism, the module responsible for nonlocal interactions between
density features on two different atoms (Supplementary Sec. B.1). In contrast
to commonly used multi-layer perceptrons (MLPs), the Transformer-based model
captures “relative relation” instead of “absolute values” of the input
feature, which generalizes better for varying-length inputs.This nonlocal
formulation has a low cost prefactor by the virtue of the concise density
representation, while is indeed crucial for KEDF according to our results in
Supplementary Sec. D.4.2.

Perhaps unexpectedly, learning a functional model is more challenging beyond
conventional machine learning. Primarily, the model is used as an optimization
objective. This requires a higher quality than for end-to-end prediction, as
the error would accumulate during optimization hence deviate or even diverge
the process. The model needs to capture the energy landscape on the
coefficient space, for which only one datapoint per molecular structure is far
from sufficient. We hence design methods to produce _multiple_ coefficient
data points, each also with a _gradient_ (w.r.t the coefficients) label, for
each molecular structure (Methods 4.1). Moreover, the input coefficients are
tensors equivariant to the rotation of coordinate system, but the output
energy is invariant. To guarantee this geometric invariance in the model, we
employ atom-wise _local frames_. They also stabilize the coefficients for the
same type of bonds (Methods 4.2). Finally, the model is aimed at a physical
mechanism by which the output energy would increase sharply when the input
density deviates from the ground state. To allow the model expressing such
large gradients, we introduce a series of _enhancement modules_ that balances
the sensitivity over coefficient dimensions, rescales the gradient in each
dimension, and offsets the gradient with a reference (Methods 4.3).

After the $T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$ model is
learned, M-OFDFT solves the electronic energy and density of a given molecular
structure $\mathcal{M}$ through the density optimization procedure (Fig.
1(c)):

$\displaystyle\min_{\mathbf{p}:\mathbf{p}^{\top}\mathbf{w}=N}E_{\theta}(\mathbf{p},\mathcal{M}):=T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})+E_{\textnormal{H}}(\mathbf{p},\mathcal{M})+E_{\textnormal{XC}}(\mathbf{p},\mathcal{M})+E_{\textnormal{ext}}(\mathbf{p},\mathcal{M}),$
(2)

where $E_{\textnormal{H}}$, $E_{\textnormal{ext}}$ and $E_{\textnormal{XC}}$
can be computed from $(\mathbf{p},\mathcal{M})$ by the conventional way
(Supplementary Sec. A.3.2). The constraint on $\mathbf{p}$ fulfills a
normalized density, where
$\mathbf{w}_{\mu}:=\int\omega_{\mu}(\mathbf{r})\,\mathrm{d}\mathbf{r}$ is the
basis normalization vector. The optimization is solved by gradient descent:

$\displaystyle\mathbf{p}^{(k+1)}:=\mathbf{p}^{(k)}-\varepsilon\left\lparen\mathbf{I}-\frac{\mathbf{w}\mathbf{w}^{\top}}{\mathbf{w}^{\top}\mathbf{w}}\right\rparen\nabla_{\mathbf{p}}E_{\theta}(\mathbf{p}^{(k)},\mathcal{M}),$
(3)

where $\varepsilon$ is a step size, and the gradient is projected onto the
admissible plane in respect to the linear constraint. Notably, due to directly
operating on density, the complexity of M-OFDFT in each iteration is
$O(N^{2})$ (Supplementary Sec. A.3.2), which is order-$N$ less than that
$O(N^{3})$ (with density fitting; Supplementary Sec. A.3.1) of KSDFT.

###  2.2 M-OFDFT Achieves Chemical Accuracy on Molecular Systems

We first evaluate the performance of M-OFDFT on molecules in similar scales
but unseen in training. We generate datasets based on two settings: ethanol
structures from the MD17 dataset [37, 38] for studying conformational space
generalization, and molecular structures from the QM9 dataset [39, 40] for
studying chemical space generalization. Each dataset is split into three parts
for the training and validation of the KEDF model, and the test of M-OFDFT.
For ease of training, we use the APBE functional [41] as a base KEDF and let
the deep-learning model learn the residual (Supplementary Sec. B.4.1).

We evaluate M-OFDFT in terms of the mean absolute error (MAE) from KSDFT
results in energy, as well as in the Hellmann-Feynman (HF) force
(Supplementary Sec. C.5). The results are 0.18$\,\mathrm{kcal/mol}$ and
1.18$\,\mathrm{kcal/mol/\text{\r{A}}}$ on ethanols, and
0.93$\,\mathrm{kcal/mol}$ and 2.91$\,\mathrm{kcal/mol/\text{\r{A}}}$ on QM9
(Supplementary Sec. D.1.1 shows more results). We see M-OFDFT achieves
chemical accuracy (1$\,\mathrm{kcal/mol}$ energy MAE) in both cases.

To show the significance of this result, we compare M-OFDFT with classical
OFDFT using well-established KEDFs, including the Thomas-Fermi (TF) KEDF [4,
5] which is exact in the uniform electron gas limit, its corrections
TF+$\frac{1}{9}$vW [12] and TF+vW [42] with the von Weizsäcker (vW) KEDF [43],
and the base KEDF APBE (Supplementary Sec. C.3). We note that different KEDFs
may have different absolute energy biases, so for the energy error we compare
the MAE in relative energy. On ethanol structures, the relative energy is
taken w.r.t the energy on the equilibrium conformation. On QM9, as each
molecule only has one conformation, we evaluate the relative energy between
every pair from the 6,095 isomers of $\mathrm{C_{7}H_{10}O_{2}}$ in the QM9
dataset. These isomers can be seen as different conformations of the same set
of atoms [44]. As shown in Fig. 2(a), M-OFDFT still achieves chemical accuracy
on relative energy, and is two orders more accurate than classical OFDFT.

As a qualitative investigation of M-OFDFT, we visualize the density on a test
ethanol structure optimized by these methods in Fig. 2(b) (Supplementary Sec.
D.1.2 shows more). Radial density by spherical integral around the oxygen atom
is plotted. We find that the M-OFDFT curves coincide with the KSDFT curve
precisely. Particularly, the two major peaks around 0 Å and 1.4 Å correspond
to the density of core electrons of the oxygen atom and the bonded carbon
atom, while the minor peak in between reflects the density of electrons in the
covalent bonds with the hydrogen atom and the carbon atom. M-OFDFT
successfully recovers this shell structure, which is deemed difficult for
OFDFT. In comparison, the classical OFDFT using the APBE KEDF does not align
well with the true density around the covalent bonds. These results suggest
that M-OFDFT is a working OFDFT for molecular systems.

Figure 2: Results of M-OFDFT compared with classical OFDFT on molecular
systems. (a) Relative energy and Hellmann-Feynman (HF) force results in mean
absolute error (MAE) from KSDFT, with 95% confidence interval bars. (b)
Visualization of optimized density. Each curve plots the integrated density on
spheres with varying radii centered at the oxygen atom in an ethanol
structure.

###  2.3 M-OFDFT Extrapolates Well to Larger-Scale Molecules

To wield the advantage of the lower cost scaling of M-OFDFT for a more
meaningful impact, we evaluate its accuracy on molecular systems with a scale
beyond affordable for generating abundant training data. For running on large
molecules, we train the deep-learning model targeting the sum of the kinetic
and XC energy to get rid of the demanding calculation on grid (Supplementary
Sec. B.4.2). This modification does not lead to obvious accuracy lost
(Supplementary Sec. D.1.1).

To evaluate the significance of the extrapolation performance, we compare
M-OFDFT with a natural variant of deep machine learning method that directly
predicts the ground-state energy from the molecular structure $\mathcal{M}$ in
an end-to-end manner, which we call M-PES (following “potential energy
surface”). We also consider a variant named M-PES-Den that additionally takes
the MINAO initialized density into input for investigating the effect of
density feature on extrapolation. Both variants use the same nonlocal model
architecture and training settings as M-OFDFT for fair comparison
(Supplementary Sec. C.4).

##### QMugs

Figure 3: Extrapolation performance of M-OFDFT compared with other deep-
learning methods. Considered are M-PES and M-PES-Den that use deep-learning
models to predict the ground-state energy end-to-end. The shades and bars show
95% confidence intervals. (a) Mean absolute error (MAE) in per-atom energy on
increasingly larger molecules from the QMugs dataset, using models trained on
molecules with no more than 15 heavy atoms from QM9 and QMugs datasets. (b)
Energy error on 50 QMugs molecules with 56-60 heavy atoms, using models
trained on a series of datasets containing increasingly larger QMugs molecules
up to 30 heavy atoms. The horizontal dotted black line marks the performance
of M-OFDFT trained on the first dataset. (c) Relative energy error on 50
Chignolin structures, using models trained on all peptides (lengths 2-5). Also
shown is the result of the classical OFDFT using APBE. (d) Energy error on
1,000 Chignolin structures, using models trained on a series of datasets
including increasingly longer peptides. (e) Energy error on 50 Chignolin
structures, using models trained on all peptides without (‘Pretrain’) and with
‘Finetune’ on 500 Chignolin structures. Also marked are error reduction ratios
by the finetuned models over models trained ‘FromScratch’ on the 500 Chignolin
structures only.

We first study the extrapolation on the QMugs dataset [45], containing much
larger molecules than those in QM9 which have no more than 9 heavy atoms. We
train the models on QM9 together with QMugs molecules with no more than 15
heavy atoms, and test the methods on larger QMugs molecules up to 101 heavy
atoms, which are grouped according to the number of heavy atoms into bins of
width 5, and are randomly subsampled to ensure the same number (50) of
molecular structures in each bin to eliminate statistical effects.

The result is shown in Fig. 3(a). We see that the per-atom MAE of M-OFDFT is
always orders smaller than M-PES and M-PES-Den in absolute value, even though
M-PES and M-PES-Den achieve a lower validation error (Supplementary Table 6).
More attractively, the error of M-OFDFT keeps constant and even decreases
(note the negative exponent) when the molecule scale increases, while the
errors of M-PES and M-PES-Den keep increasing, even though they use the same
nonlocal architecture capable of capturing long-range effects, and M-PES-Den
also has a density input. We attribute the qualitatively better extrapolation
to appropriately formulating the machine-learning task. The ground-state
energy of a molecular structure is the _result_ of an intricate, many-body
interaction among electrons and nuclei, leading to a highly challenging
function to extrapolate from one region to another. M-OFDFT converts the task
into learning the objective function for the target output. The objective only
needs to capture the _mechanism_ that the particles interact, which has a
reduced level of complexity, while transferring a large portion of complexity
to the optimization process, for which optimization tools can handle
effectively without an extrapolation issue. Similar phenomena have also been
observed recently in machine learning that learning an objective shows better
extrapolation than learning an end-to-end map [46, 47].

To further substantiate the significance of the extrapolation capability of
M-OFDFT, we investigate the magnitude by which the training molecule scale
must be increased for M-PES and M-PES-Den to achieve the same level of
performance as M-OFDFT on a given workload of large-scale molecules. We take
50 QMugs molecules with 50-60 heavy atoms as the extrapolation benchmark, and
train the models on a series of equal-sized datasets that include increasingly
larger molecules up to 30 heavy atoms. As shown in Fig. 3(b), M-PES and M-PES-
Den require at least twice as large molecules in the training dataset (30
_vs._ 15 heavy atoms) to achieve a commensurate accuracy
(0.068${\,\mathrm{kcal/mol}}$) as M-OFDFT provides. These extrapolation
results suggest that M-OFDFT can be applied to systems much larger than
training to exploit the scaling advantage, and is more affordable to develop
for solving large-scale molecular systems.

##### Chignolin

An increasingly important portion of the demand for large-scale quantum
chemistry calculation comes from biomolecular systems, particularly proteins,
which are not touched by OFDFT previously. We assess the capability of M-OFDFT
for protein systems on the Chignolin protein (10 residues, 168 atoms after
neutralization). We consider the practical setup that it is unaffordable to
generate abundant data for the large target system hence requiring
extrapolation. We generate training data on smaller-scale systems of short
peptide structures containing 2 to 5 residues, cropped from 1,000 Chignolin
structures selected from [48]. To account for non-covalent effects, training
data also include systems containing two peptides of lengths 2 and 2, and 2
and 3, where each peptide pair is cropped from the same Chignolin structure.
See more details in Supplementary Sec. C.1.4. For this task, we let the model
target the total energy for a learning stability consideration (Supplementary
Sec. B.4.3).

We first train the model on all available peptides, and compare the relative
energy error on Chignolin with other methods in Fig. 3(c). Notably, M-OFDFT
achieves a significantly lower per-atom error than the classical OFDFT using
the APBE KEDF (0.098${\,\mathrm{kcal/mol}}$ _vs._
0.684${\,\mathrm{kcal/mol}}$), providing an effective OFDFT method for
biomolecular systems. M-OFDFT also outperforms deep-learning variants M-PES
and M-PES-Den, indicating a better extrapolation capability. To investigate
extrapolation in more detail, we train the deep-learning models on peptides
with increasingly larger scale and plot the error on Chignolin in Fig. 3(d)
(similar to the setting of Fig. 3(b)). Remarkably, M-OFDFT consistently
outperforms end-to-end energy prediction methods M-PES and M-PES-Den across
all lengths of training peptides, and halves the required length for the same
level of accuracy. We note the spikes of M-PES and M-PES-Den at peptide length
3 despite extensive hyperparameter tuning, possibly due to that their harder
extrapolation difficulty magnifies the gap between in-scale validation and
larger-scale performance in this case.

After being trained on data in accessible scale, which is called “pretraining”
in the following context, a deep-learning model for a larger-scale workload
can be further improved if a few larger-scale data are available for
finetuning. In this situation, a method capable of good extrapolation could be
roughly aligned with the larger-scale task in advance using accessible data,
more efficiently leveraging the limited larger-scale data, and outperforming
the model trained from scratch on these limited data only. To investigate the
benefit of M-OFDFT in this scenario, we build a finetuning dataset on 500
Chignolin structures. Results in Fig. 3(e) show that M-OFDFT achieves the most
gain from pretraining, reducing the energy error by 35.4% over training from
scratch, showing the appeal of extracting a more generalizable rule from
accessible-scale data. With finetuning, M-OFDFT still gives the best absolute
accuracy. These results suggest that M-OFDFT could effectively handle as large
a molecular system as a protein, even without abundant training data on the
same large scale.

###  2.4 M-OFDFT Has a Lower Empirical Time Complexity than KSDFT

Figure 4: Empirical time cost of M-OFDFT compared with KSDFT on molecules in
various scales. Each plotted value is the average of running times on
molecules whose number of electrons falls in the corresponding bin of width
20.

After validating the accuracy and extrapolation capability, we now demonstrate
the scaling advantage of M-OFDFT empirically. The time cost for running both
methods on increasingly large molecules from the QMugs dataset [45] is plotted
in Fig. 4. We see the absolute running time of M-OFDFT is always shorter than
that of KSDFT, achieving up to 6.7-fold speedup. The empirical complexity of
M-OFDFT is $O(N^{1.46})$, which is indeed at least order-$N$ less than the
empirical complexity $O(N^{2.49})$ of KSDFT. Supplementary Sec. C.2 details
the running setup.

To further wield the advantage, we run M-OFDFT on two molecular systems as
large as proteins: (1) the peripheral subunit-binding domain BBL-H142W (PDB
ID: 2WXC) [49] containing 2,676 electrons (709 atoms), and (2) the K5I/K39V
double mutant of the Albumin binding domain of protein B (PDB ID: 1PRB) [50]
containing 2,750 electrons (738 atoms). Such a scale exceeds the typical
workload of KSDFT [51]. M-OFDFT costs 0.41 and 0.45 hours on the two systems,
while using KSDFT costs 10.5 and 12.3 hours, hence a 25.6-fold and 27.4-fold
speedup is achieved. Supplementary Sec. D.3 provides more details.

##  3 Conclusion and Discussion

This work has developed M-OFDFT, an orbital-free density functional theory
that works successfully on molecules. The central task to approximate the
kinetic energy density functional (KEDF) is regarded challenging, especially
for molecular systems. We have shown that such approximation can be achieved
much more accurately by modern deep machine learning models with proper
architecture design and training techniques. M-OFDFT achieves working accuracy
on molecules and shows desirable extrapolation capability, unleashing the
attractive scalability of OFDFT for large molecular systems.

This work introduces a few technical improvements for learning a functional
model. Instead of a grid-based representation, we used the density
coefficients on an atomic basis as model input, whose much lower
dimensionality allows our construction of a nonlocal architecture to enhance
accuracy and extrapolation. Some works [52, 53] on learning the XC functional
also adopt the coefficient input, but without the molecular structure input,
hence cannot properly capture inter-atomic density feature interaction.
Regarding the additional challenge for learning an objective, we generated
_multiple_ data points each also with a _gradient_ label for each molecular
structure. Although the possibility has been noted by previous works [24, 25],
none has fully leveraged such abundant data for training (some only
incorporated gradient [29, 30, 32, 28]). There are other ways to regularize
the optimization behavior of a functional model [54, 55, 56, 53], but our
trials in Supplementary Sec. D.4.4 show that they are not as effective. To
express intrinsically large gradient, we introduce enhancement modules in
addition to a conventional neural network. For stable density optimization
using a learned model, prior works [24, 25, 27, 29] used projection onto the
training-data manifold in each step, while M-OFDFT only needs the
initialization be on the manifold (Methods 4.4).

While statistical guarantees are established for in-distribution
generalization [57], reliable extrapolation remains an open challenge and has
long bothered the application of machine learning in the science domain [58,
59]. This work has demonstrated the improved extrapolation by choosing an
appropriate formulation of quantum chemistry: learning a density functional
extrapolates qualitatively better than direct energy prediction. Incorporating
exact properties of KEDF into the model also benefits extrapolation. We have
geometric invariance built into the model using local frame. Nevertheless, it
is not always straightforward to gain benefits from these properties, since
some would introduce more training challenges or unintended model capacity
restrictions. For example, we tried the von Weizsäcker KEDF as the base KEDF
which introduces positivity to the residual model [60, Thm. 1.1], but the
resulting gradient labels are too large to learn effectively. The KEDF also
has a scaling property, but it cannot be translated into an exact equation
under atomic basis (Supplementary Sec. A.5).

For better extrapolation, another possibility from the machine learning
perspective is using more data and larger model size with a proper
architecture. Recent progress in large language model [61, 62] has shown the
emergence of the capacity to solve seemingly all language tasks given large
enough data and model. A similar trend in the Graphormer architecture is
hinted by a recent study [63] for equilibrium distribution, indicating
opportunity to further improve the universality of the functional model.

In summary, M-OFDFT represents an advancement of the frontier of accuracy-
efficiency trade-off in quantum chemistry, creating a new powerful tool for
exploring large and complex molecules with a higher level of detail and scale.

##  4 Methods

In response to the more challenges beyond conventional (deep) machine
learning, we describe methodological details for KEDF model training (Methods
4.1), additional design for geometric invariance (Methods 4.2) and large
gradient capacity (Methods 4.3), and density optimization strategies (Methods
4.4) of M-OFDFT.

###  4.1 Training the KEDF Model

Although learning the KEDF model
$T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$ can be converted to a
supervised machine learning task, it is more challenging than the conventional
form. The essential difference is rooted in the way that the model is used:
instead of as an end-to-end mapping to predict the kinetic energy of
$(\mathbf{p},\mathcal{M})$ queries, the model is used as the objective to
optimize the density coefficients $\mathbf{p}$ for a given molecular structure
$\mathcal{M}$ (Fig. 1(c)). To eliminate instability and achieve accurate
optimization result, the model is required to capture how to vary with
$\mathbf{p}$ for a fixed $\mathcal{M}$, _i.e._ , the optimization landscape on
the coefficient space. Conventional data format
$\\{\mathcal{M}^{(d)},\mathbf{p}^{(d)},T_{\textnormal{S}}^{(d)}\\}_{d}$ does
not effectively convey such information, since only one labeled $\mathbf{p}$
data point is seen for each $\mathcal{M}$. Hence the first requirement on
training data is _multiple_ coefficient data points per structure, following
the format
$\\{\mathcal{M}^{(d)},\\{\mathbf{p}^{(d,k)},T_{\textnormal{S}}^{(d,k)}\\}_{k}\\}_{d}$.
On such data, the model is trained by minimizing:

$\displaystyle\sum_{d}\sum_{k}\left\lvert
T_{\textnormal{S},\theta}(\mathbf{p}^{(d,k)},\mathcal{M}^{(d)})-T_{\textnormal{S}}^{(d,k)}\right\rvert.$
(4)

After some trials, we found this is still not sufficient. The trained model,
although accurately predicts the kinetic energy value, still decreases the
electronic energy in density optimization (Eq. (3)) even starting from the
ground-state density. This indicates the gradient
$\nabla_{\mathbf{p}}T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$ w.r.t
the coefficients is still not accurately recovered. We hence also desire a
_gradient_ label for each data point, which constitutes data in the format
$\\{\mathcal{M}^{(d)},\\{\mathbf{p}^{(d,k)},T_{\textnormal{S}}^{(d,k)},\nabla_{\mathbf{p}}T_{\textnormal{S}}^{(d,k)}\\}_{k}\\}_{d}$.
As only the projected gradient matters for density optimization following Eq.
(3), the gradient data is used for training the model by minimizing:

$\displaystyle\sum_{d}\sum_{k}\bigg{\lVert}\bigg{(}\mathbf{I}-\frac{\mathbf{w}^{(d)}{\mathbf{w}^{(d)}}^{\top}}{{\mathbf{w}^{(d)}}^{\top}\mathbf{w}^{(d)}}\bigg{)}\Big{(}\nabla_{\mathbf{p}}T_{\textnormal{S},\theta}(\mathbf{p}^{(d,k)},\mathcal{M}^{(d)})-\nabla_{\mathbf{p}}T_{\textnormal{S}}^{(d,k)}\Big{)}\bigg{\rVert}.$
(5)

The gradient label provides additional information on the local landscape near
each coefficient data point. As the model is used in density optimization only
through its gradient, the gradient data directly stabilizes and regularizes
density optimization, and enforces stationary-point condition for correct
convergence. Supplementary Sec. D.4.1 verifies the improvement empirically
through an ablation study.

To generate such multiple-coefficient and gradient-labeled data, we note that
it is tractable from running the conventional KSDFT on each molecular
structure $\mathcal{M}^{(d)}$, which conducts a self-consistent field (SCF)
iteration. The rationale is that, the task in each SCF step $k$ is to solve a
non-interacting fermion system in an effective one-body potential constructed
from previous steps. The ground-state wavefunction solution is a Slater
determinant specified by the $N$ orbital solutions in that step, by which the
non-interacting kinetic energy $T_{\textnormal{S}}^{(d,k)}$ can be directly
calculated. The corresponding density coefficients $\mathbf{p}^{(d,k)}$ can be
calculated from these orbitals by density fitting [64]. For the gradient
label, since $\mathbf{p}^{(d,k)}$ represents the ground-state density of the
non-interacting system, it minimizes the energy of the non-interacting system
as a function of density coefficient,
$T_{\textnormal{S}}(\mathbf{p},\mathcal{M}^{(d)})+\mathbf{p}^{\top}\mathbf{v}_{\textnormal{eff}}^{(d,k)}$,
where $\mathbf{v}_{\textnormal{eff}}^{(d,k)}$ is the effective potential in
SCF step $k$ in vector form under the atomic basis. This indicates
$\nabla_{\mathbf{p}}T_{\textnormal{S}}(\mathbf{p}^{(d,k)},\mathcal{M}^{(d)})=-\mathbf{v}_{\textnormal{eff}}^{(d,k)}$
up to the normalization projection. Supplementary Sec. A.2 elaborates more on
the reasoning, and Supplementary Sec. A.4 provides calculation details,
including an efficient implementation to generate the gradient label.

In our implementation of M-OFDFT, the atomic basis for representing density is
taken as the even-tempered basis set [65] with tempering ratio $\beta=2.5$.
For generating data, restricted-spin KSDFT is conducted at the
PBE/6-31G(2df,p) level, which is sufficient for the considered systems which
are uncharged, in near-equilibrium conformation, and only involve light atoms
(up to fluorine).

###  4.2 Geometric Invariance

Another challenge beyond conventional machine learning is that the target
physical functional exhibits symmetry w.r.t transformations on the input
$(\mathbf{p},\mathcal{M}=\\{\mathbf{X},\mathbf{Z}\\})$ arising from the
translation and rotation of the molecule. This is formally referred to as
$\mathrm{SE}(3)$-invariance, following “3-dimensional special Euclidean group”
that comprises these transformations. This is because the non-interacting
kinetic energy of electrons does not change with the translation and rotation
of the molecule, but the input atomic coordinates $\mathbf{X}$ do, and the
input density coefficients $\mathbf{p}$ also change with the rotation. The
change of $\mathbf{p}$ is due to that the electron density rotates with the
molecule, but the atomic basis functions do not, since their orientations are
aligned with the (global) coordinate system, a.k.a frame. Formally, such input
features are geometric vectors and tensors that change equivariantly with the
translation and/or rotation of the molecule. Subsequently, the model is
expected to have this $\mathrm{SE}(3)$-invariance built-in. This allows the
model to learn the essential dependency of the energy on the density
irrespective of geometric variability, reducing the problem space, and
facilitating data efficiency and effective training. The invariance also
enhances generalization and extrapolation performance, as an important
physical property is always guaranteed [66, 67].

For the invariance w.r.t atomic coordinates $\mathbf{X}$, the neural network
model of Graphormer is naturally $\mathrm{SE}(3)$-invariant, since the model
only uses relative distances of atom pairs for later processing, which are
inherently invariant w.r.t the translation and rotation of the molecule. To
ensure the invariance of the model w.r.t the density coefficients
$\mathbf{p}$, we introduce a transformation on $\mathbf{p}$ under _local
frames_ to make invariant coefficient features. Each local frame is associated
to an atom, and specifies the orientation of atomic basis functions on that
atom. It is determined by the relative positions among nearby atoms, hence the
basis function orientations rotate with the molecule and the density, making
the density coefficients under the local frame invariant. Specifically, the
local frame on the atom located at $\mathbf{x}_{a}^{(0)}$ is determined
following previous works (_e.g._ , [68, 69]): the x-axis unit vector
$\hat{\mathbf{x}}:=\mathrm{Normalize}(\mathbf{x}_{a}^{(1)}-\mathbf{x}_{a}^{(0)})$
is pointed to its nearest heavy atom located at $\mathbf{x}_{a}^{(1)}$, then
the z-axis is pointed to
$\hat{\mathbf{z}}:=\mathrm{Normalize}\big{(}\hat{\mathbf{x}}\times(\mathbf{x}_{a}^{(2)}-\mathbf{x}_{a}^{(0)})\big{)}$,
where $\mathbf{x}_{a}^{(2)}$ is the coordinates of the second-nearest heavy
atom not collinear with the nearest one, and finally the y-axis is pointed to
$\hat{\mathbf{y}}:=\hat{\mathbf{z}}\times\hat{\mathbf{x}}$ following a right-
handed system. See Supplementary Sec. B.2 for more details.

Moreover, the local frame approach offers an additional benefit that the
coefficient features are stabilized for local molecular substructures, _e.g._
, bond or functional group, of the same type. Such substructures on one
molecule may have different orientations relative to the whole molecule, but
the electron density on them are naturally close, up to a rotation. Other
invariant implementations, _e.g._ , using an equivariant global frame [70, 71]
or processing tensorial input invariantly [72, 73, 74, 75], bind the basis
orientations on different atoms together, so the resulting coefficients on the
substructures appear vastly different. In contrast, using local frames, basis
orientations on different atoms are decoupled, and since they are determined
only by nearby atoms, the basis functions rotate from one substructure to
another accordingly. Hence, the resulting density coefficients are aligned
together, whose difference only indicates the minor density fluctuation on the
same type of substructure but not the different orientations of the copies.
This makes the model much easier to identify that such local density
components follow the same pattern and contribute similarly to the energy.
Supplementary Sec. B.2 provides an illustrative explanation. We numerically
demonstrate the benefit in Supplementary Figs. 9-10 that using local frame
instead of equivariant global frame significantly reduces the variance of both
density coefficients and gradients on atoms of each type. Especially, on most
basis functions of hydrogen, the coefficient and gradient scales are reduced
by over 60%. This significantly stabilizes the training process and
immediately reduces training error, resulting in a considerable improvement of
overall performance as empirically verified in Supplementary Sec. D.4.3.

###  4.3 Enhancement Modules for Vast Gradient Range

After reducing the geometric variability of data using local frame, the raw
gradient values still show a vast range, which conventional neural networks
are not designed for (_e.g._ , [76]) and indeed causes training difficulties
in our trials. This is an intrinsic challenge for learning a physical
functional since we require non-ground-state density in the data, which would
increase the energy steeply. The large gradient range cannot be trivially
reduced by conventional data normalization techniques, since its scale is
associated with the scale of energy and coefficient, hence downscaling the
gradient would either proportionally downscale the energy values which
requires a higher prediction resolution, or inverse-proportionally upscale the
coefficients which is also numerically unfriendly to process. To handle this
challenge, we introduce a series of enhancement modules to allow expressing a
vast gradient range, including dimension-wise rescaling, a reparameterization
of the density coefficients, and an atomic reference module to offset the
large mean of gradient.

##### Dimension-wise Rescaling

We first upgrade data normalization more flexibly to trade-off coefficient-
gradient scales dimension-wise. Considering the number of coefficient
dimensions vary from different molecules, we propose to center and rescale the
coefficients using biases $\bar{\mathbf{p}}_{Z,\tau}$ and factors
$\lambda_{Z,\tau}$ each specific to one coefficient/gradient dimension $\tau$
associated with one _atom type_ (_i.e._ , chemical element) $Z$ but not one
atom. The bias
$\bar{\mathbf{p}}_{Z,\tau}:=\mathrm{mean}\\{\mathbf{p}_{a,\tau}^{(d,k)}\\}_{a:Z^{(a)}=Z,\;k,\,d}$
for $(Z,\tau)$ is the average over coefficient values in dimension $\tau$ on
all atoms of type $Z$ in all molecular structures in the training dataset.
After centering the coefficients using the bias (which does not affect
gradients), the scaling factor $\lambda_{Z,\tau}$ is determined by upscaling
the centered coefficient and simultaneously inverse-proportionally downscaling
the gradient, until the gradient achieves a chosen target scale
$s_{\textnormal{grad}}$ or the coefficient exceeds a chosen maximal scale
$s_{\textnormal{coeff}}$. In equation:

$\displaystyle\lambda_{Z,\tau}=\begin{dcases}\min\bigg{\\{}\frac{\mathrm{max\\_grad}_{Z,\tau}}{s_{\textnormal{grad}}},\,\frac{s_{\textnormal{coeff}}}{\mathrm{std\\_coeff}_{Z,\tau}}\bigg{\\}},&\text{if
}\mathrm{max\\_grad}_{Z,\tau}>s_{\textnormal{grad}},\\\
1,&\text{otherwise},\end{dcases}$ (6)

where the scales of gradient and coefficient for $(Z,\tau)$ are measured by
the maximum of gradient
$\mathrm{max\\_grad}_{Z,\tau}:=\max\\{\nabla_{\mathbf{p}_{a,\tau}}T_{\textnormal{S}}^{(d,k)}\\}_{a:Z^{(a)}=Z,\;k,\,d}$
and standard derivation of coefficient
$\mathrm{std\\_coeff}_{Z,\tau}:=\mathrm{std}\\{\mathbf{p}_{a,\tau}^{(d,k)}\\}_{a:Z^{(a)}=Z,\;k,\,d}$
on the dataset. Using the rescaling factors, each centered coefficient is
rescaled by
$\mathbf{p}^{\prime}_{a,\tau}:=\lambda_{Z^{(a)},\tau}\,\mathbf{p}_{a,\tau}$,
and gradient by
$\nabla_{\mathbf{p}_{a,\tau}}T_{\textnormal{S}}^{\prime}:=\nabla_{\mathbf{p}_{a,\tau}}T_{\textnormal{S}}\,/\,\lambda_{Z^{(a)},\tau}$
($\lambda_{Z,\tau}>1$ in most cases).

##### Natural Reparameterization

On quite a few dimensions, both the coefficient and gradient scales are large,
making dimension-wise rescaling ineffective. We hence introduce _natural
reparameterization_ applied before rescaling to balance the rescaling
difficulties across dimensions hence reduce the worst-case difficulty. The
unbalanced scales come from the different sensitivities of the density
function on different coefficient dimensions: the change of density function
from a coefficients change $\Delta\mathbf{p}$ is measured by the L2-metric in
the function space,
$\int\left\lvert\Delta\rho(\mathbf{r})\right\rvert^{2}\,\mathrm{d}\mathbf{r}$,
which turns out to be $\Delta\mathbf{p}^{\top}\mathbf{W}\Delta\mathbf{p}$, in
which different dimensions indeed contribute with different weights since the
overlap matrix
$\mathbf{W}_{\mu\nu}:=\int\omega_{\mu}(\mathbf{r})\omega_{\nu}(\mathbf{r})\,\mathrm{d}\mathbf{r}$
therein is generally anisotropic. The reparameterized coefficients
$\tilde{\mathbf{p}}$ are expected to contribute equally across the dimensions:
$\int\left\lvert\Delta\rho(\mathbf{r})\right\rvert^{2}\,\mathrm{d}\mathbf{r}=\Delta\tilde{\mathbf{p}}^{\top}\Delta\tilde{\mathbf{p}}$.
We hence take:

$\displaystyle\tilde{\mathbf{p}}:=\mathbf{M}^{\top}\mathbf{p},$ (7)

where $\mathbf{M}$ is a square matrix satisfying
$\mathbf{M}\mathbf{M}^{\top}=\mathbf{W}$. See Supplementary Sec. B.3.2 for
more details. This reparameterization also leads to natural gradient descent
[77] in density optimization, which is known to converge faster than vanilla
gradient descent.

##### Atomic Reference Module

Recall that in dimension-wise rescaling, the large bias of coefficients can be
offset by the mean on a dataset, but this does not reduce the bias scale of
gradient labels. To further improve the coefficient-gradient scale trade-off,
we introduce an _atomic reference module_ :

$\displaystyle
T_{\mathrm{AtomRef}}(\mathbf{p},\mathcal{M}):=\bar{\mathbf{g}}_{\mathcal{M}}^{\top}\mathbf{p}+\bar{T}_{\mathcal{M}},$
(8)

which is linear in the coefficients $\mathbf{p}$ and whose output is added to
the neural network output as the kinetic energy value. By this design, the
gradient of the atomic reference model
$\nabla_{\mathbf{p}}T_{\mathrm{AtomRef}}(\mathbf{p},\mathcal{M})=\bar{\mathbf{g}}_{\mathcal{M}}$
is a constant, which offsets the target gradient for the neural network to
capture, effectively reducing the scale of gradient labels and facilitating
neural network training. The weights
$\bar{\mathbf{g}}_{\mathcal{M}}:=\mathrm{concat}\\{\bar{\mathbf{g}}_{Z^{(a)},\tau}\\}_{\tau,\,a\in\mathcal{M}}$
and bias
$\bar{T}_{\mathcal{M}}:=\sum_{a\in\mathcal{M}}\bar{T}_{Z^{(a)}}+\bar{T}_{\mathrm{global}}$
of the linear model are constructed by tiling and summing the per-type
statistics, which are derived over all atoms of each type in a dataset. The
per-type gradient statistics is defined by
$\bar{\mathbf{g}}_{Z,\tau}:=\mathrm{mean}\\{\nabla_{\mathbf{p}_{a,\tau}}T_{\textnormal{S}}^{(d,k)}\\}_{a:Z^{(a)}=Z,\;k,\,d}$,
which represents the average response to $T_{\textnormal{S}}$ from the change
of coefficients on an atom of type $Z$. Per-type bias statistics
$\\{\bar{T}_{Z}\\}_{Z}$ and $\bar{T}_{\mathrm{global}}$ are fit by least
squares. See Supplementary Sec. B.3.3 for more details.

The final KEDF model is constructed from these enhancement modules and the
neural network model in the following way (Supplementary Fig. 6(a)): the
density coefficients are first transformed under local frame and processed by
natural reparameterization; the processed coefficients, through one branch,
are fed to the atomic reference module to calculate the reference part of
output energy, and through another branch, are processed by dimension-wise
rescaling and then input to the neural network model which produces the rest
part of output energy. Comparative results in Supplementary Sec. B.3 and D.4.3
highlight the empirical benefits of each module.

###  4.4 Density Optimization

In the deployment stage, M-OFDFT solves the ground state of a given molecular
structure $\mathcal{M}$ by minimizing the electronic energy as a function of
density coefficients $\mathbf{p}$, where the learned KEDF model
$T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$ is used to construct the
energy function (Fig. 1(a)). As described in Results 2.1, we use gradient
descent to optimize $\mathbf{p}$ (Eq. (3)), since it is unnatural to formulate
the optimization problem into a self-consistent iteration. Gradient descent
has also been used in KSDFT, which bears the merit of being more stable [78].

Figure 5: Typical density optimization curves of M-OFDFT for a QM9 molecule
with different initialization methods. MINAO, the common KSDFT initialization,
leads the optimization to a large gap from the target energy, since it is not
from the eigensolution to an effective Hamiltonian matrix, hence lies off the
training-data manifold (out of distribution). Hückel initialization solves an
eigenvalue problem, which indeed converges the curve with a much smaller
energy error of 5.34 $\mathrm{kcal/mol}$. ProjMINAO initialization uses a
deep-learning model to project the MINAO density onto the training-data
manifold, which also converges the curve and achieves the best result of 0.60
$\mathrm{kcal/mol}$ energy error. The inset figure highlights the role of
density optimization even though the ProjMINAO density is close to the ground-
state density.

A subtlety in density optimization using a learned functional model is that
the model may be confronted with densities far from the training-data manifold
(or “out of distribution” in machine-learning term), which may lead to
unstable optimization. Such an issue has been observed in previous machine-
learning OFDFT [24, 25, 27], which mitigates the problem by projecting the
density onto the training-data manifold in each optimization step. A similar
phenomenon is also observed in M-OFDFT. As shown in Fig. 5, when starting from
the MINAO initialization [79] which is common for KSDFT, the density
optimization process leads to an obvious gap from the target KSDFT energy. We
note that the initial density by MINAO already lies off the manifold
inherently: each density entry in the training data comes from the
eigensolution to an effective one-electron Hamiltonian matrix, which exactly
solves an effective non-interacting fermion system (Supplementary Sec.
A.2),while the MINAO density comes from the superposition of orbitals of each
atom in isolation, which is a different mechanism.

We hence propose using two other initialization methods to resolve the
mismatch. The first approach is to use an established initialization that
solves an eigenvalue problem, for which we choose the Hückel initialization
[80, 81]. As shown in Fig. 5, although the Hückel density shows a much larger
energy error than MINAO density at initialization, it ultimately indeed leads
the optimization process to converge closely to the target energy.

The second choice is to project the MINAO density onto the training-data
manifold, which we call ProjMINAO. In contrast to previous methods, M-OFDFT
conducts optimization on the coefficient space which varies with molecular
structure, so the training-data manifold of coefficient is unknown for an
unseen molecular structure. We hence use another deep-learning model
$\Delta\mathbf{p}_{\theta}(\mathbf{p},\mathcal{M})$ to predict the required
correction to project the input coefficient $\mathbf{p}$ towards the ground-
state coefficient $\mathbf{p}^{\star}$ of the input molecular structure
$\mathcal{M}$, which is always on the manifold. See Supplementary Sec. B.5.2
for details. From Fig. 5, we see ProjMINAO initialization indeed converges the
optimization curve close to the target energy, even better than Hückel
initialization. Note that even though ProjMINAO already closely approximates
the ground state density, density optimization still continues to improve the
accuracy. This suggests a potential advantage over end-to-end ground-state
density prediction followed by energy prediction from ground-state density,
which may also encounter extrapolation challenges similar to M-PES and M-PES-
Den. Remarkably, Fig. 5 indicates that M-OFDFT only requires an on-manifold
initialization but does not need projection in each optimization step,
suggesting better robustness than previous methods. M-OFDFT in Results 2 is
conducted using ProjMINAO, although using Hückel still achieves a reasonable
accuracy; see Supplementary Sec. D.1.1. Supplementary Sec. B.5.1 provides
curves in density error and comparison with classical KEDFs.

## Acknowledgements

We thank Paola Gori Giorgi, William Chuck Witt, Sebastian Ehlert, Zun Wang,
Lixue Cheng, Jan Hermann and Ziteng Liu for insightful discussions and
constructive feedback; Xingheng He and Yaosen Min for suggestions on protein
preprocessing; Yu Shi for suggestions and feedback on model design and
optimization; and Jingyun Bai for helping with figure design.

## Author information

### Author contributions

C.Liu led the research under the support from B.Shao and N.Zheng. C.Liu,
S.Zheng and B.Shao conceived the project. H.Zhang, S.Zheng and J.You designed
and implemented the deep-learning model. S.Liu, C.Liu, H.Zhang and J.You
derived and implemented methods for data generation, the enhancement modules,
training pipeline, and density optimization. H.Zhang and S.Liu conducted the
experiments. Z.Lu and T.Wang contributed to the experiment design and
evaluation protocol. C.Liu, H.Zhang, S.Liu and S.Zheng wrote the paper with
inputs from all the authors.

### Corresponding authors

Correspondence to Chang Liu, Shuxin Zheng, and Bin Shao.

## References

  * Seminario [1996] Jorge M Seminario.  Recent developments and applications of modern density functional theory.  1996\.
  * Jain et al. [2013] Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, and Kristin A. Persson.  Commentary: The Materials Project: A materials genome approach to accelerating materials innovation.  _APL Materials_ , 1(1):011002, 07 2013.  ISSN 2166-532X.  doi: 10.1063/1.4812323.  URL https://doi.org/10.1063/1.4812323.
  * Kohn and Sham [1965] Walter Kohn and Lu Jeu Sham.  Self-consistent equations including exchange and correlation effects.  _Physical review_ , 140(4A):A1133, 1965.
  * Thomas [1927] Llewellyn H Thomas.  The calculation of atomic fields.  In _Mathematical proceedings of the Cambridge philosophical society_ , volume 23, pages 542–548. Cambridge University Press, 1927.
  * Fermi [1928] Enrico Fermi.  Eine statistische methode zur bestimmung einiger eigenschaften des atoms und ihre anwendung auf die theorie des periodischen systems der elemente.  _Zeitschrift für Physik_ , 48(1):73–79, 1928\.
  * Slater [1951] John C Slater.  A simplification of the Hartree-Fock method.  _Physical review_ , 81(3):385, 1951.
  * Hohenberg and Kohn [1964] Pierre Hohenberg and Walter Kohn.  Inhomogeneous electron gas.  _Physical review_ , 136(3B):B864, 1964.
  * Wang and Carter [2000] Yan Alexander Wang and Emily A Carter.  Orbital-free kinetic-energy density functional theory.  _Theoretical methods in condensed phase chemistry_ , 5:117–84, 2000.
  * Karasiev et al. [2014] Valentin V Karasiev, Debajit Chakraborty, and SB Trickey.  Progress on new approaches to old ideas: Orbital-free density functionals.  In _Many-electron approaches in physics, chemistry and mathematics: a multidisciplinary view_ , pages 113–134. Springer, 2014.
  * Witt et al. [2018] William C Witt, G Beatriz, Johannes M Dieterich, and Emily A Carter.  Orbital-free density functional theory for materials research.  _Journal of Materials Research_ , 33(7):777–795, 2018.
  * Hodges [1973] C. H. Hodges.  Quantum corrections to the Thomas–Fermi approximation — the Kirzhnits method.  _Canadian Journal of Physics_ , 51(13):1428–1437, 1973.  doi: 10.1139/p73-189.  URL https://doi.org/10.1139/p73-189.
  * Brack et al. [1976] M. Brack, B.K. Jennings, and Y.H. Chu.  On the extended Thomas-Fermi approximation to the kinetic energy density.  _Physics Letters B_ , 65(1):1–4, 1976.  ISSN 0370-2693.  doi: https://doi.org/10.1016/0370-2693(76)90519-0.  URL https://www.sciencedirect.com/science/article/pii/0370269376905190.
  * Wang and Teter [1992] Lin-Wang Wang and Michael P Teter.  Kinetic-energy functional of the electron density.  _Physical Review B_ , 45(23):13196, 1992.
  * Wang et al. [1999] Yan Alexander Wang, Niranjan Govind, and Emily A Carter.  Orbital-free kinetic-energy density functionals with a density-dependent kernel.  _Physical Review B_ , 60(24):16350, 1999.
  * Huang and Carter [2010] Chen Huang and Emily A Carter.  Nonlocal orbital-free kinetic energy density functional for semiconductors.  _Physical Review B_ , 81(4):045206, 2010.
  * Zhou et al. [2005] Baojing Zhou, Vincent L Ligneres, and Emily A Carter.  Improving the orbital-free density functional theory description of covalent materials.  _The Journal of chemical physics_ , 122(4), 2005.
  * Qiu et al. [2017] Ruizhi Qiu, Haiyan Lu, Bingyun Ao, Li Huang, Tao Tang, and Piheng Chen.  Energetics of intrinsic point defects in aluminium via orbital-free density functional theory.  _Philosophical Magazine_ , 97(25):2164–2181, 2017\.
  * García-Aldea and Alvarellos [2007] David García-Aldea and JE Alvarellos.  Kinetic energy density study of some representative semilocal kinetic energy functionals.  _The Journal of chemical physics_ , 127(14), 2007.
  * Xia et al. [2012] Junchao Xia, Chen Huang, Ilgyou Shin, and Emily A Carter.  Can orbital-free density functional theory simulate molecules?  _The Journal of chemical physics_ , 136(8):084102, 2012.
  * Sillaste and Thompson [2022] Spencer Sillaste and Russell B Thompson.  Molecular bonding in an orbital-free-related density functional theory.  _The Journal of Physical Chemistry A_ , 126(2):325–332, 2022.
  * Teale et al. [2022] Andrew M Teale, Trygve Helgaker, Andreas Savin, Carlo Adamo, Bálint Aradi, Alexei V Arbuznikov, Paul W Ayers, Evert Jan Baerends, Vincenzo Barone, Patrizia Calaminici, et al.  DFT exchange: sharing perspectives on the workhorse of quantum chemistry and materials science.  _Physical chemistry chemical physics_ , 24(47):28700–28781, 2022.
  * García-González et al. [1996] P. García-González, J. E. Alvarellos, and E. Chacón.  Nonlocal kinetic-energy-density functionals.  _Phys. Rev. B_ , 53:9509–9512, Apr 1996.  doi: 10.1103/PhysRevB.53.9509.  URL https://link.aps.org/doi/10.1103/PhysRevB.53.9509.
  * Mi et al. [2018] Wenhui Mi, Alessandro Genova, and Michele Pavanello.  Nonlocal kinetic energy functionals by functional integration.  _The Journal of Chemical Physics_ , 148(18):184107, 05 2018.  ISSN 0021-9606.  doi: 10.1063/1.5023926.  URL https://doi.org/10.1063/1.5023926.
  * Snyder et al. [2012] 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\.
  * Snyder et al. [2013] 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.
  * Li et al. [2016] Li Li, Thomas E Baker, Steven R White, Kieron Burke, et al.  Pure density functional for strong correlation and the thermodynamic limit from machine learning.  _Physical Review B_ , 94(24):245129, 2016.
  * Brockherde et al. [2017] Felix Brockherde, Leslie Vogt, Li Li, Mark E Tuckerman, Kieron Burke, and Klaus-Robert Müller.  Bypassing the Kohn-Sham equations with machine learning.  _Nature communications_ , 8(1):1–10, 2017.
  * del Mazo-Sevillano and Hermann [2023] Pablo del Mazo-Sevillano and Jan Hermann.  Variational principle to regularize machine-learned density functionals: the non-interacting kinetic-energy functional.  _arXiv preprint arXiv:2306.17587_ , 2023.
  * Meyer et al. [2020] Ralf Meyer, Manuel Weichselbaum, and Andreas W Hauser.  Machine learning approaches toward orbital-free density functional theory: Simultaneous training on the kinetic energy density functional and its functional derivative.  _Journal of chemical theory and computation_ , 16(9):5685–5694, 2020.
  * Fujinami et al. [2020] Mikito Fujinami, Ryo Kageyama, Junji Seino, Yasuhiro Ikabata, and Hiromi Nakai.  Orbital-free density functional theory calculation applying semi-local machine-learned kinetic energy density functional and kinetic potential.  _Chemical Physics Letters_ , 748:137358, 2020.
  * Ryczko et al. [2021] Kevin Ryczko, Sebastian J Wetzel, Roger G Melko, and Isaac Tamblyn.  Orbital-free density functional theory with small datasets and deep learning.  _arXiv preprint arXiv:2104.05408_ , 2021.
  * Imoto et al. [2021] Fumihiro Imoto, Masatoshi Imada, and Atsushi Oshiyama.  Order-N orbital-free density-functional calculations with machine learning of functional derivatives for semiconductors and metals.  _Physical Review Research_ , 3(3):033198, 2021\.
  * Ying et al. [2021] Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu.  Do Transformers really perform badly for graph representation?  _Advances in Neural Information Processing Systems_ , 34:28877–28888, 2021.
  * Shi et al. [2022] Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie-Yan Liu.  Benchmarking Graphormer on large-scale molecular modeling datasets.  _arXiv preprint arXiv:2203.04810_ , 2022.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Kato [1957] Tosio Kato.  On the eigenfunctions of many-particle systems in quantum mechanics.  _Communications on Pure and Applied Mathematics_ , 10(2):151–177, 1957.
  * Chmiela et al. [2017] Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller.  Machine learning of accurate energy-conserving molecular force fields.  _Science advances_ , 3(5):e1603015, 2017.
  * Chmiela et al. [2019] Stefan Chmiela, Huziel E Sauceda, Igor Poltavsky, Klaus-Robert Müller, and Alexandre Tkatchenko.  sGDML: Constructing accurate and data efficient molecular force fields using machine learning.  _Computer Physics Communications_ , 240:38–45, 2019.
  * Ruddigkeit et al. [2012] Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond.  Enumeration of 166 billion organic small molecules in the chemical universe database GDB-17.  _Journal of chemical information and modeling_ , 52(11):2864–2875, 2012.
  * Ramakrishnan et al. [2014] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific Data_ , 1, 2014.
  * Constantin et al. [2011] Lucian A Constantin, E Fabiano, S Laricchia, and F Della Sala.  Semiclassical neutral atom as a reference system in density functional theory.  _Physical review letters_ , 106(18):186406, 2011\.
  * Karasiev and Trickey [2012] Valentin V Karasiev and Samuel B Trickey.  Issues and challenges in orbital-free density functional calculations.  _Computer Physics Communications_ , 183(12):2519–2527, 2012.
  * Weizsäcker [1935] CF von Weizsäcker.  Zur theorie der kernmassen.  _Zeitschrift für Physik_ , 96(7):431–458, 1935.
  * Smith et al. [2017] 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.
  * Isert et al. [2022] Clemens Isert, Kenneth Atz, José Jiménez-Luna, and Gisbert Schneider.  QMugs, quantum mechanical properties of drug-like molecules.  _Scientific Data_ , 9(1):1–11, 2022.
  * Du and Mordatch [2019] Yilun Du and Igor Mordatch.  Implicit generation and generalization in energy-based models.  _arXiv preprint arXiv:1903.08689_ , 2019.
  * Du et al. [2022] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch.  Learning iterative reasoning through energy minimization.  In _International Conference on Machine Learning_ , pages 5570–5582. PMLR, 2022.
  * Lindorff-Larsen et al. [2011] Kresten Lindorff-Larsen, Stefano Piana, Ron O Dror, and David E Shaw.  How fast-folding proteins fold.  _Science_ , 334(6055):517–520, 2011.
  * Neuweiler et al. [2009] Hannes Neuweiler, Timothy D Sharpe, Trevor J Rutherford, Christopher M Johnson, Mark D Allen, Neil Ferguson, and Alan R Fersht.  The folding mechanism of BBL: Plasticity of transition-state structure observed within an ultrafast folding protein family.  _Journal of molecular biology_ , 390(5):1060–1073, 2009.
  * Wang et al. [2004] Ting Wang, Yongjin Zhu, and Feng Gai.  Folding of a three-helix bundle at the folding speed limit.  _The Journal of Physical Chemistry B_ , 108(12):3694–3697, 2004.
  * Rayson and Briddon [2008] Mark J Rayson and Patrick R Briddon.  Rapid iterative method for electronic-structure eigenproblems using localised basis functions.  _Computer Physics Communications_ , 178(2):128–134, 2008.
  * Dick and Fernandez-Serra [2020] Sebastian Dick and Marivi Fernandez-Serra.  Machine learning accurate exchange and correlation functionals of the electronic density.  _Nature communications_ , 11(1):3509, 2020.
  * Chen et al. [2021] 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, 2021.
  * Nagai et al. [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\.
  * Kirkpatrick et al. [2021] James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander L Gaunt, James S Spencer, Alexander GDG Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, et al.  Pushing the frontiers of density functionals by solving the fractional electron problem.  _Science_ , 374(6573):1385–1389, 2021.
  * Li et al. [2021a] 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.  _Physical review letters_ , 126(3):036401, 2021a.
  * Kawaguchi et al. [2022] K. Kawaguchi, Y. Bengio, and L. Kaelbling.  _Generalization in Deep Learning_ , page 112–148.  Cambridge University Press, 2022.  doi: 10.1017/9781009025096.003.
  * Huang et al. [2023] Bing Huang, Guido Falk von Rudorff, and O Anatole von Lilienfeld.  The central role of density functional theory in the AI age.  _Science_ , 381(6654):170–175, 2023.
  * Zhang et al. [2023] Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, et al.  Artificial intelligence for science in quantum, atomistic, and continuum systems.  _arXiv preprint arXiv:2307.08423_ , 2023.
  * Lieb [1983] Elliott H Lieb.  Density functionals for Coulomb systems.  _International Journal of Quantum Chemistry_ , 24(3):243–277, 1983.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Zhao et al. [2023] Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al.  A survey of large language models.  _arXiv preprint arXiv:2303.18223_ , 2023.
  * Zheng et al. [2023] Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, Zhang He, Shidi Tang, Hongxia Hao, Peiran Jin, Chi Chen, Frank Noé, Haiguang Liu, and Tie-Yan Liu.  Towards predicting equilibrium distributions for molecular systems with deep learning.  _arXiv preprint arXiv:2306.05445_ , 2023.
  * Dunlap [2000] Brett I. Dunlap.  Robust and variational fitting.  _Phys. Chem. Chem. Phys._ , 2:2113–2116, 2000.  doi: 10.1039/B000027M.  URL http://dx.doi.org/10.1039/B000027M.
  * Bardo and Ruedenberg [1974] Richard D. Bardo and Klaus Ruedenberg.  Even-tempered atomic orbitals. VI. optimal orbital exponents and optimal contractions of Gaussian primitives for hydrogen, carbon, and oxygen in molecules.  _The Journal of Chemical Physics_ , 60(3):918–931, 1974.  doi: 10.1063/1.1681168.  URL https://doi.org/10.1063/1.1681168.
  * Satorras et al. [2021] Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling.  E (n) equivariant graph neural networks.  In _International conference on machine learning_ , pages 9323–9332. PMLR, 2021.
  * Thomas et al. [2018] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley.  Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds.  _arXiv preprint arXiv:1802.08219_ , 2018.
  * Han et al. [2018] Jiequn Han, Linfeng Zhang, Roberto Car, et al.  Deep Potential: A general representation of a many-body potential energy surface.  _Communications in Computational Physics_ , 23(3), 2018\.
  * Li et al. [2022] He Li, Zun Wang, Nianlong Zou, Meng Ye, Runzhang Xu, Xiaoxun Gong, Wenhui Duan, and Yong Xu.  Deep-learning density functional theory Hamiltonian for efficient ab initio electronic-structure calculation.  _Nature Computational Science_ , 2(6):367–377, 2022.
  * Li et al. [2021b] Feiran Li, Kent Fujiwara, Fumio Okura, and Yasuyuki Matsushita.  A closer look at rotation-invariant deep point cloud analysis.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 16218–16227, 2021b.
  * Puny et al. [2021] Omri Puny, Matan Atzmon, Heli Ben-Hamu, Ishan Misra, Aditya Grover, Edward J Smith, and Yaron Lipman.  Frame averaging for invariant and equivariant network design.  _arXiv preprint arXiv:2110.03336_ , 2021.
  * Fuchs et al. [2020] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling.  Se (3)-transformers: 3d roto-translation equivariant attention networks.  _Advances in neural information processing systems_ , 33:1970–1981, 2020.
  * Schütt et al. [2021] Kristof Schütt, Oliver Unke, and Michael Gastegger.  Equivariant message passing for the prediction of tensorial properties and molecular spectra.  In _International Conference on Machine Learning_ , pages 9377–9388. PMLR, 2021.
  * Batzner et al. [2022] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky.  E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.  _Nature communications_ , 13(1):2453, 2022.
  * Musaelian et al. [2023] Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky.  Learning local equivariant representations for large-scale atomistic dynamics.  _Nature Communications_ , 14(1):579, 2023.
  * Fazlyab et al. [2019] Mahyar Fazlyab, Alexander Robey, Hamed Hassani, Manfred Morari, and George Pappas.  Efficient and accurate estimation of lipschitz constants for deep neural networks.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Amari [1998] Shun-Ichi Amari.  Natural gradient works efficiently in learning.  _Neural computation_ , 10(2):251–276, 1998.
  * Yoshikawa and Sumita [2022] Naruki Yoshikawa and Masato Sumita.  Automatic differentiation for the direct minimization approach to the Hartree–Fock method.  _The Journal of Physical Chemistry A_ , 126(45):8487–8493, 2022.
  * Sun et al. [2018] 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 of chemistry framework.  _Wiley Interdisciplinary Reviews: Computational Molecular Science_ , 8(1):e1340, 2018.
  * Hoffmann [1963] Roald Hoffmann.  An Extended Hückel Theory. I. Hydrocarbons.  _The Journal of Chemical Physics_ , 39(6):1397–1412, 06 1963.  ISSN 0021-9606.  doi: 10.1063/1.1734456.  URL https://doi.org/10.1063/1.1734456.
  * Lehtola [2019] Susi Lehtola.  Assessment of initial guesses for self-consistent field calculations. Superposition of atomic potentials: Simple yet efficient.  _Journal of chemical theory and computation_ , 15(3):1593–1604, 2019.
  * Levy [1979] Mel Levy.  Universal variational functionals of electron densities, first-order density matrices, and natural spin-orbitals and solution of the v-representability problem.  _Proceedings of the National Academy of Sciences_ , 76(12):6062–6065, 1979.
  * Bobrowicz and Goddard [1977] Frank W. Bobrowicz and William A. Goddard.  _The Self-Consistent Field Equations for Generalized Valence Bond and Open-Shell Hartree-Fock Wave Functions_ , pages 79–127.  Springer US, Boston, MA, 1977.  ISBN 978-1-4757-0887-5.  doi: 10.1007/978-1-4757-0887-5˙4.  URL https://doi.org/10.1007/978-1-4757-0887-5_4.
  * Levine et al. [2009] Ira N Levine, Daryle H Busch, and Harrison Shull.  _Quantum chemistry_ , volume 6.  Pearson Prentice Hall Upper Saddle River, NJ, 2009.
  * Blinder [1965] SM Blinder.  Basic concepts of self-consistent-field theory.  _American journal of physics_ , 33(6):431–443, 1965.
  * Pulay [1982] P. Pulay.  Improved scf convergence acceleration.  _Journal of Computational Chemistry_ , 3(4):556–560, 1982.  doi: https://doi.org/10.1002/jcc.540030413.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/jcc.540030413.
  * Kudin et al. [2002] Konstantin N. Kudin, Gustavo E. Scuseria, and Eric Cancès.  A black-box self-consistent field convergence algorithm: One step closer.  _The Journal of Chemical Physics_ , 116(19):8255–8261, 04 2002.  ISSN 0021-9606.  doi: 10.1063/1.1470195.  URL https://doi.org/10.1063/1.1470195.
  * Dupuis et al. [1976] Michel Dupuis, John Rys, and Harry F. King.  Evaluation of molecular integrals over Gaussian basis functions.  _The Journal of Chemical Physics_ , 65(1):111–116, 07 1976.  ISSN 0021-9606.  doi: 10.1063/1.432807.  URL https://doi.org/10.1063/1.432807.
  * Rys et al. [1983] J. Rys, M. Dupuis, and H. F. King.  Computation of electron repulsion integrals using the rys quadrature method.  _Journal of Computational Chemistry_ , 4(2):154–157, 1983.  doi: https://doi.org/10.1002/jcc.540040206.  URL https://onlinelibrary.wiley.com/doi/abs/10.1002/jcc.540040206.
  * Sun [2015] Qiming Sun.  Libcint: An efficient general integral library for Gaussian basis functions.  _Journal of computational chemistry_ , 36(22):1664–1671, 2015.
  * Baydin et al. [2018] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind.  Automatic differentiation in machine learning: a survey.  _Journal of Marchine Learning Research_ , 18:1–43, 2018\.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.
  * Perdew et al. [1996] John P Perdew, Kieron Burke, and Matthias Ernzerhof.  Generalized gradient approximation made simple.  _Physical review letters_ , 77(18):3865, 1996\.
  * Parr and Yang [1989] Robert G Parr and Weitao Yang.  Density-functional theory of atoms and molecules.  1989\.
  * Hollingsworth et al. [2018] Jacob Hollingsworth, Li Li, Thomas E Baker, and Kieron Burke.  Can exact conditions improve machine-learned density functionals?  _The Journal of chemical physics_ , 148(24):241743, 2018.
  * Kalita et al. [2021] Bhupalee Kalita, Li Li, Ryan J McCarty, and Kieron Burke.  Learning to approximate density functionals.  _Accounts of Chemical Research_ , 54(4):818–826, 2021.
  * Changyong et al. [2014] Feng Changyong, Wang Hongyue, Lu Naiji, Chen Tian, He Hua, Lu Ying, et al.  Log-transformation and its implications for data analysis.  _Shanghai archives of psychiatry_ , 26(2):105, 2014.
  * Wang et al. [2019] Jiang Wang, Simon Olsson, Christoph Wehmeyer, Adrià Pérez, Nicholas E Charron, Gianni De Fabritiis, Frank Noé, and Cecilia Clementi.  Machine learning of coarse-grained molecular dynamics force fields.  _ACS central science_ , 5(5):755–767, 2019.
  * Pérez-Hernández et al. [2013] Guillermo Pérez-Hernández, Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé.  Identification of slow molecular order parameters for markov model construction.  _The Journal of chemical physics_ , 139(1):07B604_1, 2013.
  * Schwantes and Pande [2013] Christian R Schwantes and Vijay S Pande.  Improvements in markov state model construction reveal many non-native interactions in the folding of NTL9.  _Journal of chemical theory and computation_ , 9(4):2000–2009, 2013.
  * McGibbon et al. [2015] Robert T McGibbon, Kyle A Beauchamp, Matthew P Harrigan, Christoph Klein, Jason M Swails, Carlos X Hernández, Christian R Schwantes, Lee-Ping Wang, Thomas J Lane, and Vijay S Pande.  Mdtraj: a modern open library for the analysis of molecular dynamics trajectories.  _Biophysical journal_ , 109(8):1528–1532, 2015\.
  * Eastman et al. [2017] Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al.  OpenMM 7: Rapid development of high performance algorithms for molecular dynamics.  _PLoS computational biology_ , 13(7):e1005659, 2017.
  * Case et al. [2008] David A Case, Tom A Darden, Thomas E Cheatham, Carlos L Simmerling, Junmei Wang, Robert E Duke, Ray Luo, MRCW Crowley, Ross C Walker, Wei Zhang, et al.  Amber 10.  Technical report, University of California, 2008.
  * Becke [1993] Axel D. Becke.  Density-functional thermochemistry. III. The role of exact exchange.  _The Journal of Chemical Physics_ , 98, 1993.  ISSN 00219606.  doi: 10.1063/1.464913.
  * Stephens et al. [1994] Philip J Stephens, Frank J Devlin, Cary F Chabalowski, and Michael J Frisch.  Ab initio calculation of vibrational absorption and circular dichroism spectra using density functional force fields.  _The Journal of physical chemistry_ , 98(45):11623–11627, 1994.
  * Weigend [2006] Florian Weigend.  Accurate coulomb-fitting basis sets for h to rn.  _Phys. Chem. Chem. Phys._ , 8:1057, 2006.  doi: 10.1039/b515623h.
  * Ho et al. [2008] Gregory S Ho, Vincent L Lignères, and Emily A Carter.  Introducing PROFESS: A new program for orbital-free density functional theory calculations.  _Computer physics communications_ , 179(11):839–854, 2008.
  * Enkovaara et al. [2010] Jussi Enkovaara, Carsten Rostgaard, J Jørgen Mortensen, Jingzhe Chen, M Dułak, Lara Ferrighi, Jeppe Gavnholt, Christian Glinsvad, V Haikola, HA Hansen, et al.  Electronic structure calculations with GPAW: a real-space implementation of the projector augmented-wave method.  _Journal of physics: Condensed matter_ , 22(25):253202, 2010.
  * Mi et al. [2016] Wenhui Mi, Xuecheng Shao, Chuanxun Su, Yuanyuan Zhou, Shoutao Zhang, Quan Li, Hui Wang, Lijun Zhang, Maosheng Miao, Yanchao Wang, et al.  ATLAS: A real-space finite-difference implementation of orbital-free density functional theory.  _Computer Physics Communications_ , 200:87–95, 2016.
  * Shao et al. [2021] Xuecheng Shao, Kaili Jiang, Wenhui Mi, Alessandro Genova, and Michele Pavanello.  DFTpy: An efficient and object-oriented platform for orbital-free dft simulations.  _Wiley Interdisciplinary Reviews: Computational Molecular Science_ , 11(1):e1482, 2021.
  * Hellman [1937] Hans Hellman.  Einführung in die Quantenchemie.  _Franz Deuticke, Leipzig_ , 285, 1937.
  * Feynman [1939] R. P. Feynman.  Forces in molecules.  _Phys. Rev._ , 56:340–343, Aug 1939.  doi: 10.1103/PhysRev.56.340.  URL https://link.aps.org/doi/10.1103/PhysRev.56.340.
  * Pulay [1969] Peter Pulay.  Ab initio calculation of force constants and equilibrium geometries in polyatomic molecules: I. Theory.  _Molecular Physics_ , 17(2):197–204, 1969.
  * Schütt et al. [2018] Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller.  Schnet–a deep learning architecture for molecules and materials.  _The Journal of Chemical Physics_ , 148(24), 2018.
  * Branch et al. [1999] Mary Ann Branch, Thomas F Coleman, and Yuying Li.  A subspace, interior, and conjugate gradient method for large-scale bound-constrained minimization problems.  _SIAM Journal on Scientific Computing_ , 21(1):1–23, 1999.
  * Virtanen et al. [2020] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, İlhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.  doi: 10.1038/s41592-019-0686-2.

## Appendix

Table 1: Notations Basic concepts
---
$\mathbf{r}\in\mathbb{R}^{3}$ | Electron coordinate
$N$ | Number of electrons in a molecular system
$\psi(\mathbf{r}^{(1)},\cdots,\mathbf{r}^{(N)})$ |  $N$-electron wavefunction
$\left\langle f\middle|g\right\rangle:=\int f(\mathbf{r})g(\mathbf{r})\,\mathrm{d}\mathbf{r}$ | The standard inner product in function space
$\langle f|\hat{O}|g\rangle:=\int f(\mathbf{r})(\hat{O}g)(\mathbf{r})\,\mathrm{d}\mathbf{r}$ | Function-space inner product with operator
$(f|g):=\int\frac{f(\mathbf{r})g(\mathbf{r}^{\prime})}{\left\lVert\mathbf{r}-\mathbf{r}^{\prime}\right\rVert}\,\mathrm{d}\mathbf{r}\mathrm{d}\mathbf{r}^{\prime}$ | Coulomb integral
Molecular system
$\mathbf{x}\in\mathbb{R}^{3}$ | Atom coordinates
$a,b\in\\{1,\,2,\,\cdots,A\\}$ | Indices for atoms in a molecule
$\mathbf{X}:=\\{\mathbf{x}^{(a)}\\}_{a=1}^{A}$ | Molecular conformation/geometry
$\mathbf{Z}:=\\{Z^{(a)}\\}_{a=1}^{A}$ | Molecular composition
$\mathcal{M}:=\\{\mathbf{X},\mathbf{Z}\\}$ | Molecular structure
$\\{\mathcal{M}^{(d)}\\}_{d=1}^{D}$ | Molecular structures in a dataset
Density functional theory
$\Phi:=\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$ | Orbitals
$i,j\in\\{1,\,2,\,\cdots,N\\}$ | Indices for orbitals or electrons
$\psi_{[\Phi]}(\mathbf{r}^{(1)},\cdots,\mathbf{r}^{(N)}):=\det[\phi_{i}(\mathbf{r}^{(j)})]_{ij}$ | Slater determinant from orbitals $\Phi=\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$
$\\{\eta_{\alpha}(\mathbf{r})\\}_{\alpha=1}^{B}$ | Orbital basis
Subscripts $\alpha,\,\beta,\,\gamma,\,\delta\in\\{1,\,2,\,\cdots,B\\}$ | Indices for orbital basis
$\mathbf{C}_{\alpha i}$ | Orbital coefficients
$\mathbf{\Gamma}_{\alpha\beta}:=\sum_{i=1}^{N}\mathbf{C}_{\alpha i}\mathbf{C}_{\beta i}$ | Density matrix
$\mathbf{S}_{\alpha\beta}:=\left\langle\eta_{\alpha}\middle|\eta_{\beta}\right\rangle$ | Overlap matrix of orbital basis
$\mathbf{D}_{\alpha\beta,\gamma\delta}:=\left\langle\eta_{\alpha}\eta_{\beta}\middle|\eta_{\gamma}\eta_{\delta}\right\rangle$ | Overlap matrix of paired orbital basis
$\tilde{\mathbf{D}}_{\alpha\beta,\gamma\delta}:=(\eta_{\alpha}\eta_{\beta}|\eta_{\gamma}\eta_{\delta})$ | 4-center-2-electron Coulomb integral of orbital basis
$\rho(\mathbf{r})$ | (One-electron reduced) density (function)
$\\{\omega_{\mu}(\mathbf{r})\\}_{\mu=1}^{M}$ | Density basis
Subscripts $\mu,\nu\in\\{1,\,2,\,\cdots,M\\}$ | Indices for density basis
Subscript $\mu=(a,\tau)$ | Atom assignment decomposition of basis index
$\mathbf{p}_{\mu}$ | Density coefficient
$\mathbf{w}_{\mu}:=\int\omega_{\mu}(\mathbf{r})\,\mathrm{d}\mathbf{r}$ | Density basis normalization vector
$\mathbf{W}_{\mu\nu}:=\left\langle\omega_{\mu}\middle|\omega_{\nu}\right\rangle$ | Overlap matrix of density basis
$\tilde{\mathbf{W}}_{\mu\nu}:=(\omega_{\mu}|\omega_{\nu})$ | 2-center-2-electron Coulomb integral of density basis
$\mathbf{L}_{\mu,\alpha\beta}:=\left\langle\omega_{\mu}\middle|\eta_{\alpha}\eta_{\beta}\right\rangle$ | Overlap matrix between density basis and orbital basis
$\mathbf{F}$ | Fock matrix in KSDFT
$k$ | Step index for SCF iteration or density optimization process
| ($\star$ for the converged step)
$(\mathbf{V}_{\textnormal{eff}})_{\alpha\beta}:=\langle\eta_{\alpha}|V_{\textnormal{eff}}|\eta_{\beta}\rangle$ | Effective potential matrix under orbital basis
$(\mathbf{v}_{\textnormal{eff}})_{\mu}:=\langle\omega_{\mu}|V_{\textnormal{eff}}\rangle$ | Effective potential vector under density basis
$\mu\in\mathbb{R}$ | Chemical potential
$\bm{\upvarepsilon}:=\operatorname{Diag}[\varepsilon_{1},\cdots,\varepsilon_{N}]$ | The diagonal $N\times N$ matrix of orbital energies
$U[\rho]$ | Universal functional
$E_{\textnormal{XC}}[\rho]$ | Exchange-correlation (XC) functional
$T_{\textnormal{S}}[\rho]$ | Kinetic energy density functional (KEDF)
$T_{\textnormal{S}}(\mathbf{p},\mathcal{M})$ | KEDF under atomic basis of density
$T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$ | KEDF model/approximation
$T_{\textnormal{APBE}}$ | The APBE kinetic functional as the base functional $T_{\textnormal{S,base}}$
$T_{\textnormal{S,res}}$ | Residual KEDF on top of the base functional $T_{\textnormal{S,base}}$
$E_{\textnormal{TXC}}$ | Kinetic and XC functional
$\mathbf{f}$ | Hellmann-Feynman force

##  Appendix A Mechanism of Density Functional Theory

In this section, we introduce details in relevant theory on DFT, including the
formulation of OFDFT under atomic basis, and the mechanism and details to use
KSDFT to generate value and gradient data for learning KEDF. Atomic units are
used through out the paper.

###  A.1 Basic Formulation of Density Functional Theory

For brevity, the following formulation is for spinless fermions therefore only
consider spacial states. For the restricted Kohn-Sham calculation we adopt for
data generation, a pair of electrons of opposite spins share a common spacial
orbital, which amounts to duplicate the orbitals in the following formulation.

The mechanism of DFT may be more intuitively introduced under Levy’s
constrained search formulation [82]. The $N$-electron Schrödinger equation for
ground state is equivalent to the following optimization problem on
$N$-electron wavefunctions $\psi(\mathbf{r}^{(1)},\cdots,\mathbf{r}^{(N)})$
under the variational principle:111 In this paper we only consider real-valued
wavefunctions (and subsequently orbitals), since we only need to solve for the
ground state of stationary Schrödinger equation without spin-orbit
interaction, for which the Hamiltonian operator
$\hat{H}:=\hat{T}+\hat{V}_{\textnormal{ee}}+\hat{V}_{\textnormal{ext}}$ is
Hermitian and real.

$\displaystyle
E^{\star}=\min_{\psi:\,\text{antisym},\left\langle\psi\middle|\psi\right\rangle=1}\langle\psi|\hat{T}+\hat{V}_{\textnormal{ee}}+\hat{V}_{\textnormal{ext}}|\psi\rangle,$
(9)

where $\hat{T}:=-\frac{1}{2}\sum_{i=1}^{N}{\nabla^{(i)}}^{2}$ is the kinetic
operator, $\hat{V}_{\textnormal{ee}}:=\sum_{1\leqslant i<j\leqslant
N}\frac{1}{\lVert\mathbf{r}^{(i)}-\mathbf{r}^{(j)}\rVert}$ is the electron-
electron Coulomb interaction (internal potential), and
$\hat{V}_{\textnormal{ext}}:=\sum_{i=1}^{N}V_{\textnormal{ext}}(\mathbf{r}^{(i)})$
comes from a one-body external potential $V_{\textnormal{ext}}(\mathbf{r})$
that commonly arises from the electrostatic field of the nuclei specified by
the given molecular structure $\mathcal{M}=\\{\mathbf{X},\mathbf{Z}\\}$ where
$\mathbf{X}:=(\mathbf{x}^{(1)},\cdots,\mathbf{x}^{(A)})$ and
$\mathbf{Z}:=(Z^{(1)},\cdots,Z^{(A)})$:

$\displaystyle
V_{\textnormal{ext}}(\mathbf{r})=-\sum_{a=1}^{A}\frac{Z^{(a)}}{\left\lVert\mathbf{r}-\mathbf{x}^{(a)}\right\rVert}.$
(10)

Although the optimization problem is exactly defined, directly optimizing the
$N$-electron wavefunction is very challenging computationally. Specifying the
wavefunction $\psi$ and evaluating the energy already require an exponential
cost in principle, as $\psi$ is a function on $\mathbb{R}^{3N}$ whose
dimension increases with $N$. To make an easier optimization problem, it is
then desired to optimize a functional of the one-electron reduced density,

$\displaystyle\rho_{[\psi]}(\mathbf{r}):=N\int\lvert\psi(\mathbf{r},\mathbf{r}^{(2)},\cdots,\mathbf{r}^{(N)})\rvert^{2}\,\mathrm{d}\mathbf{r}^{(2)}\cdots\mathrm{d}\mathbf{r}^{(N)},$
(11)

which has an intuitive physical interpretation of charge density even under a
classical view, and more importantly, the cost to specify a density is
constant (w.r.t $N$) in principle, as the density is a function on
$\mathbb{R}^{3}$ whose dimension is constant. This is the starting point of
density functional theory (DFT) [4, 5, 6], and is first formally verified by
Hohenberg and Kohn [7].

In terms of the density, the external potential energy, _i.e._ , the last term
in Eq. (9), is already an explicit density functional, since the external
potential is one-body:

$\displaystyle\langle\psi|\hat{V}_{\textnormal{ext}}|\psi\rangle=\int\rho_{[\psi]}(\mathbf{r})V_{\textnormal{ext}}(\mathbf{r})\,\mathrm{d}\mathbf{r}=E_{\textnormal{ext}}[\rho_{[\psi]}],\quad\text{where
}E_{\textnormal{ext}}[\rho]:=\int\rho(\mathbf{r})V_{\textnormal{ext}}(\mathbf{r})\,\mathrm{d}\mathbf{r}.$
(12)

For the other energy terms, using the density as an intermediate, the
optimization problem in Eq. (9) can be equivalently222 The correspondence
between the optimization space of $\rho$ and of $\psi$ to allow this
equivalence is analyzed by Lieb [60].  carried out in two levels:

$\displaystyle E^{\star}$ $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\Big{\lparen}\min_{\psi:\,\text{antisym},\rho_{[\psi]}=\rho}\langle\psi|\hat{T}+\hat{V}_{\textnormal{ee}}|\psi\rangle\Big{\rparen}+E_{\textnormal{ext}}[\rho]$
(13) $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\bigg{\\{}E[\rho]:=U[\rho]+E_{\textnormal{ext}}[\rho]\bigg{\\}}.$
(14)

Here, the result of the first-level optimization problem carrying out a
constrained search in Eq. (13) defines a density functional,

$\displaystyle
U[\rho]:=\min_{\psi:\,\text{antisym},\rho_{[\psi]}=\rho}\langle\psi|\hat{T}+\hat{V}_{\textnormal{ee}}|\psi\rangle,$
(15)

called the universal functional, as it is independent of system specification
(_i.e._ , $\mathcal{M}$ or $V_{\textnormal{ext}}$). It is composed of the
kinetic and internal potential energy of the electrons. The optimization
objective is then formally converted to a density functional $E[\rho]$ as
shown in Eq. (14).

To carry out practical computation, variants of the kinetic and internal
potential energy that allow explicit calculation or have known properties are
introduced, to cover the major part of the corresponding energies in
$U[\rho]$. The internal potential energy is covered by its classical version,
_i.e._ assuming no correlation, called the Hartree energy:

$\displaystyle
E_{\textnormal{H}}[\rho]:=\frac{1}{2}\int\frac{\rho(\mathbf{r})\rho(\mathbf{r}^{\prime})}{\left\lVert\mathbf{r}-\mathbf{r}^{\prime}\right\rVert}\,\mathrm{d}\mathbf{r}\mathrm{d}\mathbf{r}^{\prime}.$
(16)

The kinetic energy is covered by the kinetic energy density functional (KEDF),
which is defined in a similar way as the universal functional:

$\displaystyle
T_{\textnormal{S}}[\rho]:=\min_{\psi:\,\text{antisym},\rho_{[\psi]}=\rho}\langle\psi|\hat{T}|\psi\rangle.$
(17)

The remainder in the universal functional is called the _exchange-correlation
(XC) functional_ :

$\displaystyle
E_{\textnormal{XC}}[\rho]:=U[\rho]-T_{\textnormal{S}}[\rho]-E_{\textnormal{H}}[\rho],$
(18)

which is by definition also a density functional. Under this decomposition,
the density optimization problem in Eq. (14) becomes:

$\displaystyle E^{\star}=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\bigg{\\{}E[\rho]=T_{\textnormal{S}}[\rho]+\underbrace{E_{\textnormal{H}}[\rho]+E_{\textnormal{XC}}[\rho]+E_{\textnormal{ext}}[\rho]}_{=:E_{\textnormal{eff}}[\rho]}\bigg{\\}}.$
(19)

Here $E_{\textnormal{eff}}[\rho]$ is defined for future convenience and
denoted after the effective-potential interpretation of its variation detailed
later. Using carefully designed explicit expressions or machine-learning
models to approximate the $T_{\textnormal{S}}[\rho]$ and
$E_{\textnormal{XC}}[\rho]$ functionals, practical computation can be
conducted. This is the formulation of orbital-free density functional theory
(OFDFT). Indeed, the object to be optimized is the electron density, which is
one function on the constant-dimensional space of $\mathbb{R}^{3}$, hence
greatly reduces computation complexity over the original variational problem
Eq. (9). Under properly designed $T_{\textnormal{S}}[\rho]$ and
$E_{\textnormal{XC}}[\rho]$ approximations, the complexity is favorably
$O(N^{2})$ under atomic basis.

Considering the KEDF $T_{\textnormal{S}}[\rho]$ is more challenging to
approximate than the XC functional $E_{\textnormal{XC}}[\rho]$, Kohn and Sham
[3] leverage an equivalent formulation of KEDF to allow its accurate
calculation, at the cost of increasing the complexity. The alternative
formulation optimizes _determinantal_ wavefunctions. A determinantal
wavefunction for $N$ electrons is specified by $N$ one-electron wavefunctions
$\Phi:=\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$, a.k.a orbitals, following the
form:

$\displaystyle\psi_{[\Phi]}(\mathbf{r}^{(1)},\cdots,\mathbf{r}^{(N)}):=\frac{1}{\sqrt{N!}}\det[\phi_{i}(\mathbf{r}^{(j)})]_{ij},\quad\text{given
orbitals }\Phi:=\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}.$ (20)

The equivalent optimization problem in parallel with Eq. (17) is:333 The
equivalence on defining $T_{\textnormal{S}}[\rho]$ by Eq. (17) and Eq. (23) is
currently only known to be guaranteed if the density $\rho$ comes from the
ground state of a non-interacting system which is non-degenerate [60, Thm.
4.6]. Even there exists a density $\rho$ that makes the determinantally-
defined $T_{\textnormal{S}}[\rho]$ by Eq. (23) different [60, Thm. 4.8], the
determinantally-defined $T_{\textnormal{S}}[\rho]$ still gives the right
ground-state energy in density optimization (up to a closure) [60, Thm. 4.9].

$\displaystyle
T_{\textnormal{S}}[\rho]=\min_{\\{\phi_{i}\\}_{i=1}^{N}:\rho_{[\Phi]}=\rho}\langle\psi_{[\Phi]}|\hat{T}|\psi_{[\Phi]}\rangle=\min_{\vbox{\Let@\restore@math@cr\default@tag\halign{\hfil$\m@th\scriptstyle#$&$\m@th\scriptstyle{}#$\hfil\cr\\{\phi_{i}\\}_{i=1}^{N}:&\text{orthonormal},\\\
&\rho_{[\Phi]}=\rho\crcr}}}\sum_{i=1}^{N}\langle\phi_{i}|\hat{T}|\phi_{i}\rangle,$
(23) where
$\displaystyle\rho_{[\Phi]}(\mathbf{r}):=\sum_{i=1}^{N}\left\lvert\phi_{i}(\mathbf{r})\right\rvert^{2},\quad\text{given
orthonormal orbitals }\Phi:=\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}.$ (24)

In Eq. (23), the second equality holds since the density normalizes to $N$,
and a set of (non-collinear) functions can always be orthogonalized, _e.g._ ,
using the Gram-Schmidt process ([83, Sec. 3.1.4]; [84, Sec. 14.3]). This
equivalence can be understood from the interpretation of
$T_{\textnormal{S}}[\rho]$ as the _non-interacting_ portion of kinetic energy.
Indeed, for a non-interacting system, there are only kinetic energy and
external potential energy (_i.e._ , taking $\hat{V}_{\textnormal{ee}}=0$), so
the two-level optimization in parallel with Eq. (13) becomes:

$\displaystyle E^{\star}$
$\displaystyle=\min_{\psi:\,\text{antisym},\left\langle\psi\middle|\psi\right\rangle=1}\langle\psi|\hat{T}+\hat{V}_{\textnormal{ext}}|\psi\rangle$
(25) $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\Big{\lparen}\min_{\psi:\,\text{antisym},\rho_{[\psi]}=\rho}\langle\psi|\hat{T}|\psi\rangle\Big{\rparen}+E_{\textnormal{ext}}[\rho]$
(26) $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}T_{\textnormal{S}}[\rho]+E_{\textnormal{ext}}[\rho],$
(27)

from which we see that $T_{\textnormal{S}}[\rho]$ is the ground-state kinetic
energy of the non-interacting system whose ground-state density is $\rho$. On
the other hand, it is known that the ground-state wavefunction of a non-
interacting system is commonly determinantal (at least when the ground state
is non-degenerate [60, Thm. 4.6]).Hence the optimization can be rewritten as:

$\displaystyle E^{\star}$
$\displaystyle=\min_{\\{\phi_{i}\\}_{i=1}^{N}}\langle\psi_{[\Phi]}|\hat{T}+\hat{V}_{\textnormal{ext}}|\psi_{[\Phi]}\rangle$
(28) $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\Big{\lparen}\min_{\\{\phi_{i}\\}_{i=1}^{N}:\rho_{[\Phi]}=\rho}\langle\psi_{[\Phi]}|\hat{T}|\psi_{[\Phi]}\rangle\Big{\rparen}+E_{\textnormal{ext}}[\rho]$
(29) $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\bigg{\lparen}\min_{\vbox{\Let@\restore@math@cr\default@tag\halign{\hfil$\m@th\scriptstyle#$&$\m@th\scriptstyle{}#$\hfil\cr\\{\phi_{i}\\}_{i=1}^{N}:&\text{orthonormal},\\\
&\rho_{[\Phi]}=\rho\crcr}}}\sum_{i=1}^{N}\langle\phi_{i}|\hat{T}|\phi_{i}\rangle\bigg{\rparen}+E_{\textnormal{ext}}[\rho],$
(32)

which indicates Eq. (23).

Back to the main problem, leveraging this knowledge about
$T_{\textnormal{S}}[\rho]$ in the variational problem Eq. (19) gives:

$\displaystyle E^{\star}$ $\displaystyle=\min_{\rho:\geqslant
0,\int\rho(\mathbf{r})\,\mathrm{d}\mathbf{r}=N}\bigg{\lparen}\min_{\vbox{\Let@\restore@math@cr\default@tag\halign{\hfil$\m@th\scriptstyle#$&$\m@th\scriptstyle{}#$\hfil\cr\\{\phi_{i}\\}_{i=1}^{N}:&\text{orthonormal},\\\
&\rho_{[\Phi]}=\rho\crcr}}}\sum_{i=1}^{N}\langle\phi_{i}|\hat{T}|\phi_{i}\rangle\bigg{\rparen}+E_{\textnormal{H}}[\rho]+E_{\textnormal{XC}}[\rho]+E_{\textnormal{ext}}[\rho],$
(35)

which can be converted into directly optimizing the orbitals in a single-level
optimization:

$\displaystyle
E^{\star}=\min_{\\{\phi_{i}\\}_{i=1}^{N}:\text{orthonormal}}\bigg{\\{}E[\Phi]:=\sum_{i=1}^{N}\langle\phi_{i}|\hat{T}|\phi_{i}\rangle+\underbrace{E_{\textnormal{H}}[\rho_{[\Phi]}]+E_{\textnormal{XC}}[\rho_{[\Phi]}]+E_{\textnormal{ext}}[\rho_{[\Phi]}]}_{E_{\textnormal{eff}}[\rho_{[\Phi]}]}\bigg{\\}}.$
(36)

This is the formulation of Kohn-Sham density functional theory (KSDFT). With
decades of development of XC functional approximations, KSDFT has achieved
remarkable success and becomes among the most popular quantum chemistry
method. In its formulation, the object to be optimized is a set of orbitals
$\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$, which are $N$ functions on
$\mathbb{R}^{3}$. This is still substantially cheaper than optimizing an
$N$-electron wavefunction on $\mathbb{R}^{3N}$, but has a complexity at least
$O(N)$ times more than OFDFT, due to $N$ times more $\mathbb{R}^{3}$ functions
to optimize. Under atomic basis, KSDFT has a complexity at least $O(N^{3})$
(using density fitting) without further approximations. In this triumphant era
of deep machine learning, approximating a complicated functional is not as
challenging as before. Powerful deep-learning models create the opportunity to
approximate KEDF accurately enough to match successful XC functional
approximations. This would enable accurate and practical OFDFT calculation,
unleashing its power of lower complexity to push the accuracy-efficiency
trade-off in quantum chemistry.

###  A.2 KSDFT Calculation Produces Labels of KEDF

We now explain why a KSDFT calculation procedure could provide value and
gradient labels of KEDF. Computation details under atomic basis are postponed
in Supplementary Sec. A.4. We start by describing the typical algorithm to
solve the optimization problem in KSDFT. To determine the optimal solution of
orbitals $\Phi:=\\{\phi_{i}(\mathbf{r})\\}_{i=1}^{N}$, the variation of the
energy functional $E[\Phi]$ in Eq. (36) w.r.t each orbital $\phi_{i}$ is
required:

$\displaystyle\frac{\updelta E[\Phi]}{\updelta\phi_{i}}(\mathbf{r})={}$
$\displaystyle\frac{\updelta\sum_{j=1}^{N}(-1/2)\langle\phi_{j}|\nabla^{2}|\phi_{j}\rangle}{\updelta\phi_{i}}(\mathbf{r})+\int\frac{\updelta
E_{\textnormal{eff}}[\rho_{[\Phi]}]}{\updelta\rho}(\mathbf{r}^{\prime})\frac{\updelta\rho_{[\Phi]}(\mathbf{r}^{\prime})}{\updelta\phi_{i}(\mathbf{r})}\,\mathrm{d}\mathbf{r}^{\prime}$
(37) $\displaystyle={}$ $\displaystyle
2\hat{T}\phi_{i}(\mathbf{r})+2V_{\textnormal{eff}[\rho_{[\Phi]}]}(\mathbf{r})\phi_{i}(\mathbf{r}),$
(38) $\displaystyle\text{where}\quad
V_{\textnormal{eff}[\rho]}(\mathbf{r}):={}$ $\displaystyle\frac{\updelta
E_{\textnormal{eff}}[\rho]}{\updelta\rho}(\mathbf{r})=\underbrace{\int\frac{\rho(\mathbf{r}^{\prime})}{\left\lVert\mathbf{r}^{\prime}-\mathbf{r}\right\rVert}\,\mathrm{d}\mathbf{r}^{\prime}}_{=:V_{\textnormal{H}[\rho]}(\mathbf{r})}+\underbrace{\frac{\updelta
E_{\textnormal{XC}}[\rho]}{\updelta\rho}(\mathbf{r})}_{=:V_{\textnormal{XC}[\rho]}(\mathbf{r})}+V_{\textnormal{ext}}(\mathbf{r}).$
(39)

The term $V_{\textnormal{eff}[\rho_{[\Phi]}]}$ arises as a variation w.r.t the
density $\rho$ since the orbitals affect the energy component
$E_{\textnormal{eff}}$ apart from $T_{\textnormal{S}}$ (defined in Eq. (19))
only through the density $\rho_{[\Phi]}$ they define. By Eq. (24) we have
$\frac{\updelta\rho_{[\Phi]}(\mathbf{r}^{\prime})}{\updelta\phi_{i}(\mathbf{r})}=2\phi_{i}(\mathbf{r})\delta(\mathbf{r}-\mathbf{r}^{\prime})$,
which then gives Eq. (38). Combining Eq. (38) with the variation of the
orthonormal constraint yields the optimality equation for the problem Eq.
(36):

$\displaystyle\hat{F}_{[\rho_{[\Phi]}]}\phi_{i}:={}$
$\displaystyle\hat{T}\phi_{i}+V_{\textnormal{eff}[\rho_{[\Phi]}]}\phi_{i}=\varepsilon_{i}\phi_{i},\quad\forall
i=1,\cdots N.$ (40)

In the derivation, only the Lagrange multipliers $\varepsilon_{i}$ for the
normalization constraints
$\left\langle\phi_{i}\middle|\phi_{i}\right\rangle=1$ are imposed, since from
the resulting equations (40), $\\{\phi_{i}\\}_{i=1}^{N}$ are eigenstates of an
Hermitian operator $\hat{F}_{[\rho_{[\Phi]}]}$ called the Fock operator, hence
are naturally orthogonal in the general case of non-degeneracy. These
equations resemble the Schrödinger equation for $N$ non-interacting fermions,
where $V_{\textnormal{eff}[\rho_{[\Phi]}]}(\mathbf{r})$, as a function on
$\mathbb{R}^{3}$, acts as an _effective_ one-body external potential, hence
the name.

Note that $V_{\textnormal{eff}[\rho_{[\Phi]}]}$ is unknown beforehand, as
itself depends on the solution of orbitals. Hence a fixed-point iteration is
employed: starting from a set of initial orbitals
$\Phi^{(0)}:=\\{\phi_{i}^{(0)}\\}_{i=1}^{N}$, construct the Fock operator
using results in previous iterations,

$\displaystyle\hat{F}^{(k)}:=\hat{T}+V_{\textnormal{eff}}^{(k)},$ (41)

where $V_{\textnormal{eff}}^{(k)}$ is taken as
$V_{\textnormal{eff}[\rho_{[\Phi^{(k-1)}]}]}$ following this derivation,
_i.e._ , $\hat{F}^{(k)}=\hat{F}_{[\rho_{[\Phi^{(k-1)}]}]}$, and solve the
corresponding eigenvalue problem for the orbitals in the current iteration:

$\displaystyle\hat{F}^{(k)}\phi_{i}^{(k)}=\varepsilon_{i}^{(k)}\phi_{i}^{(k)},\quad\forall
i=1,\cdots,N.$ (42)

The iteration stops until “self-consistency” is achieved, _i.e._ , the
eigenstate solution $\Phi^{(k)}:=\\{\phi_{i}^{(k)}\\}_{i}$ in the current step
coincides (up to an acceptable error) with the orbitals
$\Phi^{(k-1)}:=\\{\phi_{i}^{(k-1)}\\}_{i=1}^{N}$ in the previous step that
define $\hat{F}^{(k)}$. This is the self-consistent field (SCF) method [85].

An important fact of SCF is that, in each iteration $k$, the solution
$\\{\phi_{i}^{(k)}\\}_{i=1}^{N}$ exactly defines the ground-state of a _non-
interacting_ system of $N$ fermions moving in the effective one-body potential
$V_{\textnormal{eff}}^{(k)}$ as the external potential $V_{\textnormal{ext}}$.
Indeed, the variational problem Eq. (32) that describes the non-interacting
system can be reformulated into a single-level optimization as:

$\displaystyle
E^{\star}=\min_{\\{\phi_{i}\\}_{i=1}^{N}:\text{orthonormal}}\sum_{i=1}^{N}\langle\phi_{i}|\hat{T}|\phi_{i}\rangle+\int\rho_{[\Phi]}(\mathbf{r})V_{\textnormal{eff}}^{(k)}(\mathbf{r})\,\mathrm{d}\mathbf{r},$
(43)

whose variation coincides with Eq. (42) thus solved by
$\\{\phi_{i}^{(k)}\\}_{i=1}^{N}$. This reveals the relation of SCF solution to
the KEDF: this solution of orbitals $\\{\phi_{i}^{(k)}\\}_{i=1}^{N}$ achieves
the minimum non-interacting kinetic energy
$\sum_{i=1}^{N}\langle\psi_{[\Phi]}|\hat{T}|\psi_{[\Phi]}\rangle$ among all
orthonormal orbitals that lead to the same density $\rho_{[\Phi^{(k)}]}$
(otherwise Eq. (43) can be further minimized; can also be seen from the
equivalence to Eq. (32)); by the alternative form of KEDF Eq. (23) as non-
interacting ground-state kinetic energy, we thus have:

$\displaystyle
T_{\textnormal{S}}[\rho_{[\Phi^{(k)}]}]=\langle\psi_{[\Phi^{(k)}]}|\hat{T}|\psi_{[\Phi^{(k)}]}\rangle=\sum_{i=1}^{N}\langle\phi_{i}^{(k)}|\hat{T}|\phi_{i}^{(k)}\rangle.$
(44)

This indicates that _every SCF iteration produces a label for
$T_{\textnormal{S}}$_. Moreover, as the non-interacting variation problem Eq.
(43) is equivalent to its two-level optimization form Eq. (32), which is in
turn equivalent to the density optimization form using KEDF Eq. (27) (which
explains the alternative KEDF form Eq. (23)), the density
$\rho_{[\Phi^{(k)}]}$ from the solution
$\Phi^{(k)}:=\\{\phi_{i}^{(k)}\\}_{i=1}^{N}$ of each SCF iteration minimizes
Eq. (27). Therefore, it satisfies the variation equation (Euler equation) of
Eq. (27) (taking $V_{\textnormal{ext}}$ as $V_{\textnormal{eff}}^{(k)}$)
subject to the normalization constraint with Lagrange multiplier (chemical
potential) $\mu^{(k)}$:

$\displaystyle\frac{\updelta
T_{\textnormal{S}}[\rho_{[\Phi^{(k)}]}]}{\updelta\rho}+V_{\textnormal{eff}}^{(k)}=\mu^{(k)}.$
(45)

The variation of KEDF $\frac{\updelta T_{\textnormal{S}}}{\updelta\rho}$ is
related to the gradient w.r.t density coefficients when the density $\rho$ is
expanded on a basis (see Supplementary Sec. A.4.3). Hence, _every SCF
iteration also produces a label for the gradient of $T_{\textnormal{S}}$, up
to a projection_.

It is worth noting that these arguments still hold when the effective
potential $V_{\textnormal{eff}}^{(k)}$ in SCF iteration $k$ is _not_
$V_{\textnormal{eff}[\rho_{[\Phi^{(k-1)}]}]}$, since the deductions from Eq.
(43) to Eq. (45) only require $V_{\textnormal{eff}}^{(k)}$ to be a one-body
potential. This allows more flexible data generation process since in common
DFT calculation settings, $V_{\textnormal{eff}}^{(k)}$ indeed deviates from
$V_{\textnormal{eff}[\rho_{[\Phi^{(k-1)}]}]}$ for more stable and faster
convergence, _e.g._ when using the “direct inversion in the iterative
subspace” (DIIS) method [86, 87]. This also indicates that even when the XC
functional used in data generation is not accurate, the generated value and
gradient labels for $T_{\textnormal{S}}[\rho]$ are still exact, since the XC
functional still gives an effective one-body potential to define the non-
interacting system Eq. (42) or Eq. (43), as long as it is pure (_i.e._ , only
depends on density features). In this sense, data generation for KEDF is
easier than that for the XC functional.

###  A.3 Formulation under Atomic Basis

For practical calculation, KSDFT typically uses an atomic basis
$\\{\eta_{\alpha}(\mathbf{r})\\}_{\alpha=1}^{B}$ to expand the orbitals for
conducting the SCF iteration in Eq. (42) for molecular systems. The expansion
gives:

$\displaystyle\phi_{i}(\mathbf{r})=\sum_{\alpha}\mathbf{C}_{\alpha
i}\eta_{\alpha}(\mathbf{r}),$ (46)

which converts solving for eigenfunctions into the common problem of solving
for eigenvectors of a matrix. On the other hand, as emphasized in Introduction
1 and Results 2.1, we also hope to represent the density on an atomic basis
$\\{\omega_{\mu}(\mathbf{r})\\}_{\mu=1}^{M}$ for efficient OFDFT
implementation,

$\displaystyle\rho(\mathbf{r})=\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}(\mathbf{r}).$
(47)

The left-hand-sides of Eq. (46) and Eq. (47) may also be denoted as
$\phi_{i,\mathbf{C}}$ or $\Phi_{\mathbf{C}}$ and $\rho_{\mathbf{p}}$ to
highlight the dependency on the coefficients. Note that both the orbital basis
$\\{\eta_{\alpha}(\mathbf{r})\\}_{\alpha=1}^{B}$ and density basis
$\\{\omega_{\mu}(\mathbf{r})\\}_{\mu=1}^{M}$ depend on the molecular structure
$\mathcal{M}$, as the location and type of each basis function is determined
by the coordinates $\mathbf{x}^{(a)}$ and atomic number $Z^{(a)}$ of the
corresponding atom. Nevertheless, the development in this subsection is for
one given molecular system $\mathcal{M}$, so we omit its appearance for
density or orbital representation. Typically, the numbers of basis $B$ and $M$
increase linearly with the number of electrons $N$, _i.e._ , $O(B)=O(M)=O(N)$.

####  A.3.1 KSDFT under Atomic Basis

For the SCF iteration in Eq. (42), using the expansion of orbitals in Eq.
(46), it becomes: $\sum_{\beta}\mathbf{C}_{\beta
i}^{(k)}\hat{F}^{(k)}\eta_{\beta}(\mathbf{r})=\varepsilon_{i}^{(k)}\sum_{\beta}\mathbf{C}_{\beta
i}^{(k)}\eta_{\beta}(\mathbf{r}),\;\forall i=1,\cdots,N$. Integrating each
function equation with basis function $\eta_{\alpha}(\mathbf{r})$ gives:
$\sum_{\beta}\mathbf{C}_{\beta
i}^{(k)}\langle\eta_{\alpha}|\hat{F}^{(k)}|\eta_{\beta}\rangle=\varepsilon_{i}^{(k)}\sum_{\beta}\mathbf{C}_{\beta
i}^{(k)}\langle\eta_{\alpha}|\eta_{\beta}\rangle$, which can then be
formulated as a generalized eigenvalue problem in matrix form:

$\displaystyle\mathbf{F}^{(k)}\mathbf{C}^{(k)}=\mathbf{S}\mathbf{C}^{(k)}\bm{\upvarepsilon}^{(k)},$
(48) $\displaystyle\text{where}\quad\mathbf{F}^{(k)}_{\alpha\beta}:={}$
$\displaystyle\langle\eta_{\alpha}|\hat{F}^{(k)}|\eta_{\beta}\rangle\stackrel{{\scriptstyle\text{Eq.~{}\eqref{eqn:fock-
opr}}}}{{=}}\underbrace{\langle\eta_{\alpha}|\hat{T}|\eta_{\beta}\rangle}_{=-\frac{1}{2}\int\eta_{\alpha}(\mathbf{r})\nabla^{2}\eta_{\beta}(\mathbf{r})\,\mathrm{d}\mathbf{r}=:\mathbf{T}_{\alpha\beta}}+\underbrace{\langle\eta_{\alpha}|V_{\textnormal{eff}}^{(k)}|\eta_{\beta}\rangle}_{=:(\mathbf{V}_{\textnormal{eff}}^{(k)})_{\alpha\beta}},$
(49) $\displaystyle\mathbf{S}_{\alpha\beta}:={}$
$\displaystyle\langle\eta_{\alpha}|\eta_{\beta}\rangle,\quad\bm{\upvarepsilon}^{(k)}:=\begin{pmatrix}\varepsilon_{1}^{(k)}&&\\\
&\ddots&\\\ &&\varepsilon_{N}^{(k)}\end{pmatrix}.$ (50)

Here, $\mathbf{F}^{(k)}$ is called the Fock matrix, and $\mathbf{S}$ is the
overlap matrix of the orbital basis.

To show the expression of the Fock matrix, we first give the expression of the
density defined by the orbital coefficients from Eq. (24) and Eq. (46):

$\displaystyle\rho_{\mathbf{C}}(\mathbf{r}):={}$
$\displaystyle\rho_{[\Phi_{\mathbf{C}}]}(\mathbf{r})=\sum_{i=1}^{N}\bigg{\lvert}\sum_{\alpha}\mathbf{C}_{\alpha
i}\eta_{\alpha}(\mathbf{r})\bigg{\rvert}^{2}=\sum_{i=1}^{N}\sum_{\alpha\beta}\mathbf{C}_{\alpha
i}\mathbf{C}_{\beta i}\eta_{\alpha}(\mathbf{r})\eta_{\beta}(\mathbf{r})$ (51)
$\displaystyle={}$
$\displaystyle\sum_{\alpha\beta}\mathbf{\Gamma}_{\alpha\beta}\eta_{\alpha}(\mathbf{r})\eta_{\beta}(\mathbf{r}),$
(52)

where we have defined the density matrix corresponding to the orbital
coefficients:

$\displaystyle\mathbf{\Gamma}:=\mathbf{C}\mathbf{C}^{\top},\quad\mathbf{\Gamma}_{\alpha\beta}:=\sum_{i=1}^{N}\mathbf{C}_{\alpha
i}\mathbf{C}_{\beta i}.$ (53)

Note that Eq. (24) requires orthonormal orbitals, hence Eq. (52) requires
$\mathbf{C}$ to satisfy the corresponding orthonormality constraint shown in
Eq. (67) below. Orbital coefficient solutions $\mathbf{C}^{(k)}$ in SCF
iterations satisfy this constraint as explained later.

In the derivation of SCF iteration, $V_{\textnormal{eff}}^{(k)}$ is taken as
$V_{\textnormal{eff}[\rho_{[\Phi^{(k-1)}]}]}$, _i.e._ ,
$\hat{F}^{(k)}=\hat{F}_{[\rho_{[\Phi^{(k-1)}]}]}$. This allows explicit
calculation of $\mathbf{V}_{\textnormal{eff}}^{(k)}$ as
${\mathbf{V}_{\textnormal{eff}}}_{\mathbf{C}^{(k-1)}}$ based on Eq. (39) and
$\mathbf{F}^{(k)}$ as $\mathbf{F}_{\mathbf{C}^{(k-1)}}$, for which we
introduce the following series of definitions:

$\displaystyle\mathbf{F}_{\mathbf{C}}:=[\langle\eta_{\alpha}|\hat{F}_{[\rho_{\mathbf{C}}]}|\eta_{\beta}\rangle]_{\alpha\beta}=\mathbf{T}+{\mathbf{V}_{\textnormal{eff}}}_{\mathbf{C}}\quad\text{($\mathbf{T}$
defined in Eq.~{}\eqref{eqn:fock-mat})},$ (54)
$\displaystyle{\mathbf{V}_{\textnormal{eff}}}_{\mathbf{C}}:=[\langle\eta_{\alpha}|V_{\textnormal{eff}[\rho_{\mathbf{C}}]}|\eta_{\beta}\rangle]_{\alpha\beta}={\mathbf{V}_{\textnormal{H}}}_{\mathbf{C}}+{\mathbf{V}_{\textnormal{XC}}}_{\mathbf{C}}+\mathbf{V}_{\textnormal{ext}},$
(55) where
$\displaystyle({\mathbf{V}_{\textnormal{H}}}_{\mathbf{C}})_{\alpha\beta}:=\langle\eta_{\alpha}|V_{\textnormal{H}[\rho_{\mathbf{C}}]}|\eta_{\beta}\rangle\stackrel{{\scriptstyle\text{Eqs.~{}(\ref{eqn:Veff-
fn}, \ref{eqn:den-denmat})
}}}{{=}}\iint\eta_{\alpha}(\mathbf{r})\eta_{\beta}(\mathbf{r})\frac{\sum_{\gamma\delta}\mathbf{\Gamma}_{\gamma\delta}\eta_{\gamma}(\mathbf{r}^{\prime})\eta_{\delta}(\mathbf{r}^{\prime})}{\left\lVert\mathbf{r}-\mathbf{r}^{\prime}\right\rVert}\,\mathrm{d}\mathbf{r}^{\prime}\mathrm{d}\mathbf{r}$
(56) $\displaystyle\hskip
42.67912pt=\sum_{\gamma\delta}\tilde{\mathbf{D}}_{\alpha\beta,\gamma\delta}\mathbf{\Gamma}_{\gamma\delta}=(\tilde{\mathbf{D}}\bar{\mathbf{\Gamma}})_{\alpha\beta},$
(57) $\displaystyle\hskip 45.52458pt\text{where
}\tilde{\mathbf{D}}_{\alpha\beta,\gamma\delta}:=(\eta_{\alpha}\eta_{\beta}|\eta_{\gamma}\eta_{\delta}),\text{
$\bar{\mathbf{\Gamma}}$ is the vector of flattened $\mathbf{\Gamma}$},$ (58)
$\displaystyle({\mathbf{V}_{\textnormal{XC}}}_{\mathbf{C}})_{\alpha\beta}:=\langle\eta_{\alpha}|V_{\textnormal{XC}[\rho_{\mathbf{C}}]}|\eta_{\beta}\rangle=\int
V_{\textnormal{XC}[\rho_{\mathbf{C}}]}(\mathbf{r})\eta_{\alpha}(\mathbf{r})\eta_{\beta}(\mathbf{r})\,\mathrm{d}\mathbf{r},$
(59)
$\displaystyle(\mathbf{V}_{\textnormal{ext}})_{\alpha\beta}:=\langle\eta_{\alpha}|V_{\textnormal{ext}}|\eta_{\beta}\rangle=-\sum_{a=1}^{A}Z^{(a)}\int\frac{\eta_{\alpha}(\mathbf{r})\eta_{\beta}(\mathbf{r})}{\left\lVert\mathbf{r}-\mathbf{x}^{(a)}\right\rVert}\,\mathrm{d}\mathbf{r}.$
(60)

In defining $\tilde{\mathbf{D}}$, we have used the notation of Coulomb
integral for brevity:

$\displaystyle(f|g):=\int\frac{f(\mathbf{r})g(\mathbf{r}^{\prime})}{\left\lVert\mathbf{r}-\mathbf{r}^{\prime}\right\rVert}\,\mathrm{d}\mathbf{r}\mathrm{d}\mathbf{r}^{\prime}.$
(61)

In practice, integrals in $\mathbf{S}$, $\mathbf{T}$,
$\mathbf{V}_{\textnormal{ext}}$, and the Coulomb integral $\tilde{\mathbf{D}}$
can be calculated analytically under Gaussian-Type Orbitals (GTO) as the basis
$\\{\eta_{\alpha}\\}_{\alpha=1}^{B}$. The integral in
${\mathbf{V}_{\textnormal{XC}}}_{\mathbf{C}}$ is conducted numerically on a
quadrature grid, as typically used in a DFT calculation. After convergence,
the electronic energy can be calculated from the orbital coefficients
$\mathbf{C}$ by:

$\displaystyle
E(\mathbf{C}):=E[\Phi_{\mathbf{C}}]\stackrel{{\scriptstyle\text{Eq.~{}\eqref{eqn:engopt-
orb}}}}{{=}}T_{\textnormal{S}}(\mathbf{C})+\underbrace{E_{\textnormal{H}}(\mathbf{C})+E_{\textnormal{XC}}(\mathbf{C})+E_{\textnormal{ext}}(\mathbf{C})}_{=:E_{\textnormal{eff}}(\mathbf{C})},$
(62) where $\displaystyle
T_{\textnormal{S}}(\mathbf{C}):=\sum_{i=1}^{N}\langle\phi_{i,\mathbf{C}}|\hat{T}|\phi_{i,\mathbf{C}}\rangle=\sum_{\alpha\beta}\mathbf{\Gamma}_{\alpha\beta}\mathbf{T}_{\alpha\beta}=\bar{\mathbf{\Gamma}}^{\top}\bar{\mathbf{T}},$
(63) $\displaystyle
E_{\textnormal{H}}(\mathbf{C}):=E_{\textnormal{H}}[\rho_{\mathbf{C}}]=\frac{1}{2}\sum_{\alpha\beta\gamma\delta}\mathbf{\Gamma}_{\alpha\beta}\tilde{\mathbf{D}}_{\alpha\beta,\gamma\delta}\mathbf{\Gamma}_{\gamma\delta}=\frac{1}{2}\bar{\mathbf{\Gamma}}^{\top}\tilde{\mathbf{D}}\bar{\mathbf{\Gamma}},$
(64) $\displaystyle
E_{\textnormal{XC}}(\mathbf{C}):=E_{\textnormal{XC}}[\rho_{\mathbf{C}}]=E_{\textnormal{XC}}\Big{[}\sum_{\alpha\beta}\mathbf{\Gamma}_{\alpha\beta}\eta_{\alpha}\eta_{\beta}\Big{]},$
(65) $\displaystyle
E_{\textnormal{ext}}(\mathbf{C}):=E_{\textnormal{ext}}[\rho_{\mathbf{C}}]=\sum_{\alpha\beta}\mathbf{\Gamma}_{\alpha\beta}(\mathbf{V}_{\textnormal{ext}})_{\alpha\beta}=\bar{\mathbf{\Gamma}}^{\top}\bar{\mathbf{V}}_{\textnormal{ext}},$
(66)

where $\bar{\mathbf{\Gamma}}$, $\bar{\mathbf{T}}$,
$\bar{\mathbf{V}}_{\textnormal{ext}}$ are the vectors of flattened
$\mathbf{\Gamma}$, $\mathbf{T}$, $\mathbf{V}_{\textnormal{ext}}$,
respectively. The term $E_{\textnormal{XC}}[\rho_{\mathbf{C}}]$ is again
calculated by numerically integrating the defined density by $\mathbf{C}$ on a
quadrature grid.

##### Computational Complexity

Note that the construction of ${\mathbf{V}_{\textnormal{H}}}_{\mathbf{C}}$
from Eq. (57) and the evaluation of $E_{\textnormal{H}}(\mathbf{C})$ from Eq.
(64) require $O(B^{4})=O(N^{4})$ complexity. Even when using density fitting
which decreases the complexity to $O(N^{2})$, the complexity in each SCF
iteration of KSDFT is $O(N^{3})$ since the complexity of density fitting
itself is $O(N^{3})$ (see Supplementary Sec. A.4.1).

##### Orbital Orthonormality

Under the atomic basis, orbital orthonormality
$\langle\phi_{i}|\phi_{j}\rangle=\delta_{ij}$ becomes
$\sum_{\alpha\beta}\mathbf{C}_{\alpha i}\mathbf{C}_{\beta
j}\langle\eta_{\alpha}|\eta_{\beta}\rangle=\delta_{ij}$, or in matrix form,

$\displaystyle\mathbf{C}^{\top}\mathbf{S}\mathbf{C}=\mathbf{I}.$ (67)

As mentioned after Eq. (40), only the normalization constraints need to be
taken care of, as the orbitals are eigenstates of an Hermitian operator hence
are already orthogonal if non-degenerate. This property transmits to the
matrix form of the problem:444 A direct verification in the matrix form: for
$i\neq j$,
$(\mathbf{C}^{\top}\mathbf{S}\mathbf{C})_{ij}=\mathbf{C}_{:i}^{\top}\mathbf{S}\mathbf{C}_{:j}\stackrel{{\scriptstyle\text{Eq.~{}\eqref{eqn:roothaan-
eq}}}}{{=}}\mathbf{C}_{:i}^{\top}\frac{1}{\varepsilon_{j}}\mathbf{F}\mathbf{C}_{:j}\stackrel{{\scriptstyle\text{$\mathbf{F}$
is Hermitian
(symmetric)}}}{{=}}\frac{1}{\varepsilon_{j}}(\mathbf{F}\mathbf{C}_{:i})^{\top}\mathbf{C}_{:j}\stackrel{{\scriptstyle\text{Eq.~{}\eqref{eqn:roothaan-
eq}}}}{{=}}\frac{\varepsilon_{i}}{\varepsilon_{j}}(\mathbf{S}\mathbf{C}_{:i})^{\top}\mathbf{C}_{:j}\stackrel{{\scriptstyle\text{$\mathbf{S}$
is Hermitian
(symmetric)}}}{{=}}\frac{\varepsilon_{i}}{\varepsilon_{j}}(\mathbf{C}^{\top}\mathbf{S}\mathbf{C})_{ij}$,
which indicates
$\big{(}1-\frac{\varepsilon_{i}}{\varepsilon_{j}}\big{)}(\mathbf{C}^{\top}\mathbf{S}\mathbf{C})_{ij}=0$,
thus $(\mathbf{C}^{\top}\mathbf{S}\mathbf{C})_{ij}=0$ assuming non-degeneracy
$\varepsilon_{i}\neq\varepsilon_{j}$.

$\displaystyle\mathbf{C}_{:i}^{\top}\mathbf{S}\mathbf{C}_{:i}=1,\quad\forall
i=1,\cdots,N.$ (68)

Fulfilling this constraint only needs to normalize each eigenvector
$\tilde{\mathbf{C}}^{(k)}_{:i}$ of the problem in Eq. (48) to form
$\mathbf{C}^{(k)}$; explicitly,
$\mathbf{C}^{(k)}_{:i}=\tilde{\mathbf{C}}^{(k)}_{:i}/\sqrt{{\tilde{\mathbf{C}}^{(k)}_{:i}}^{\top}\mathbf{S}\tilde{\mathbf{C}}^{(k)}_{:i}}$.

##### Relation to Direct Gradient Derivation

The matrix form of the optimality equation under basis hence the SCF iteration
problem Eq. (48) can also be derived directly from Eq. (36) by taking the
gradient of the energy function of coefficients:
$E(\mathbf{C}):=E[\Phi_{\mathbf{C}}]=E\big{[}\\{\sum_{\alpha}\mathbf{C}_{\alpha
i}\eta_{\alpha}\\}_{i=1}^{N}\big{]}$. Its gradient is related to the variation
of the functional of orbitals by integral with the basis:

$\displaystyle\big{(}\nabla_{\mathbf{C}}E(\mathbf{C})\big{)}_{\alpha
i}=\int\frac{\updelta
E[\Phi_{\mathbf{C}}]}{\updelta\phi_{i}}(\mathbf{r})\big{(}\nabla_{\mathbf{C}}\phi_{i,\mathbf{C}}(\mathbf{r})\big{)}_{\alpha
i}\,\mathrm{d}\mathbf{r}=\int\frac{\updelta
E[\Phi_{\mathbf{C}}]}{\updelta\phi_{i}}(\mathbf{r})\eta_{\alpha}(\mathbf{r})\,\mathrm{d}\mathbf{r}.$
(69)

The variation is given by Eq. (38), which is
$2\hat{F}_{[\rho_{\Phi_{\mathbf{C}}}]}\phi_{i,\mathbf{C}}(\mathbf{r})=2\sum_{\beta}\mathbf{C}_{\beta
i}\hat{F}_{[\rho_{\Phi_{\mathbf{C}}}]}\eta_{\beta}(\mathbf{r})$, which turns
the gradient into matrix form:

$\displaystyle\nabla_{\mathbf{C}}E(\mathbf{C})=2\mathbf{F}_{\mathbf{C}}\mathbf{C}.$
(70)

For the orbital orthonormality constraint, as mentioned, only the
normalization constraints require explicit treatment. By introducing Lagrange
multiplier $\varepsilon_{i}$ for each constraint in Eq. (68) and taking the
gradient for the corresponding Lagrange term gives
$\nabla_{\mathbf{C}}\sum_{i=1}^{N}\varepsilon_{i}\big{(}\mathbf{C}_{:i}^{\top}\mathbf{S}\mathbf{C}_{:i}-1\big{)}=2\mathbf{S}\mathbf{C}\bm{\upvarepsilon}$.
This leads to the optimality equation in matrix form:

$\displaystyle\mathbf{F}_{\mathbf{C}}\mathbf{C}=\mathbf{S}\mathbf{C}\bm{\upvarepsilon}.$
(71)

By constructing the corresponding fixed-point iteration, Eq. (48) is derived.

##### Accelerating and Stabilizing SCF Iteration

As mentioned, $\mathbf{F}^{(k)}$ and $\mathbf{V}_{\textnormal{eff}}^{(k)}$ may
be taken differently from $\mathbf{F}_{\mathbf{C}^{(k-1)}}$ and
${\mathbf{V}_{\textnormal{eff}}}_{\mathbf{C}^{(k-1)}}$ for more stable and
faster convergence. The direct inversion in the iterative subspace (DIIS)
method [86, 87] is a popular choice for this. In DIIS, the Fock matrix
$\mathbf{F}^{(k)}$ in the eigenvalue problem Eq. (42) for each SCF iteration
$k$ is taken as a weighted mixing of the vanilla Fock matrices
$\mathbf{F}_{\mathbf{C}^{(k^{\prime})}},k^{\prime}<k$ in previous steps:

$\displaystyle\mathbf{F}^{(k)}:=\sum_{k^{\prime}=0}^{k-1}\pi^{(k)}_{k^{\prime}}\mathbf{F}_{\mathbf{C}^{(k^{\prime})}},$
(72)

where $\\{\pi^{(k)}_{k^{\prime}}\\}_{k^{\prime}=0}^{k-1}$ are the weights that
are positive and normalized
$\sum_{k^{\prime}=0}^{k-1}\pi^{(k)}_{k^{\prime}}=1$. Due to the normalization,
the kinetic part $\mathbf{T}$ of the matrix remains the same, so it agrees
with the form in Eq. (49) (or Eq. (41) in operator form), where:

$\displaystyle\mathbf{V}_{\textnormal{eff}}^{(k)}:=\sum_{k^{\prime}=0}^{k-1}\pi^{(k)}_{k^{\prime}}{\mathbf{V}_{\textnormal{eff}}}_{\mathbf{C}^{(k^{\prime})}}.$
(73)

####  A.3.2 OFDFT under Atomic Basis

To solve the optimization problem of OFDFT in Eq. (19), it is unnatural to
construct a fixed-point SCF iteration process from its variation in Eq. (45).
Hence, direct gradient-based density optimization is conducted. For this, the
energy functional of density function in Eq. (19) needs to be converted into a
function of density coefficients using the basis expansion of density function
in Eq. (47):

$\displaystyle
E(\mathbf{p}):=E[\rho_{\mathbf{p}}]=E\Big{[}\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}\Big{]}=T_{\textnormal{S}}(\mathbf{p})+\underbrace{E_{\textnormal{H}}(\mathbf{p})+E_{\textnormal{XC}}(\mathbf{p})+E_{\textnormal{ext}}(\mathbf{p})}_{=E_{\textnormal{eff}}(\mathbf{p}):=E_{\textnormal{eff}}[\rho_{\mathbf{p}}]},$
(74) where $\displaystyle
T_{\textnormal{S}}(\mathbf{p}):=T_{\textnormal{S}}[\rho_{\mathbf{p}}]=T_{\textnormal{S}}\Big{[}\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}\Big{]},$
(75) $\displaystyle
E_{\textnormal{H}}(\mathbf{p}):=E_{\textnormal{H}}[\rho_{\mathbf{p}}]=\frac{1}{2}\iint\frac{\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}(\mathbf{r})\sum_{\nu}\mathbf{p}_{\nu}\omega_{\nu}(\mathbf{r}^{\prime})}{\lVert\mathbf{r}-\mathbf{r}^{\prime}\rVert}\,\mathrm{d}\mathbf{r}\mathrm{d}\mathbf{r}^{\prime}=\frac{1}{2}\mathbf{p}^{\top}\tilde{\mathbf{W}}\mathbf{p},$
(76) $\displaystyle
E_{\textnormal{XC}}(\mathbf{p}):=E_{\textnormal{XC}}[\rho_{\mathbf{p}}]=E_{\textnormal{XC}}\Big{[}\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}\Big{]},$
(77) $\displaystyle
E_{\textnormal{ext}}(\mathbf{p}):=E_{\textnormal{ext}}[\rho_{\mathbf{p}}]=\int\sum_{\mu}\mathbf{p}_{\mu}\omega_{\mu}(\mathbf{r})V_{\textnormal{ext}}(\mathbf{r})\,\mathrm{d}\mathbf{r}=\mathbf{p}^{\top}\mathbf{v}_{\textnormal{ext}},$
(78)
$\displaystyle\text{where}\quad\tilde{\mathbf{W}}_{\mu\nu}:=(\omega_{\mu}|\omega_{\nu}),\quad(\mathbf{v}_{\textnormal{ext}})_{\mu}:=\langle\omega_{\mu}|V_{\textnormal{ext}}\rangle.$
(79)

The Coulomb integral notation $(\omega_{\mu}|\omega_{\nu})$ is defined in Eq.
(61). Recall that we have omitted the dependency of density basis
$\\{\omega_{\mu}\\}_{\mu}$ hence of the functions _e.g._
$T_{\textnormal{S}}(\mathbf{p})$ on the molecular structure $\mathcal{M}$ in
this subsection. Integrals for $\tilde{\mathbf{W}}$ and
$\mathbf{v}_{\textnormal{ext}}$ can be calculated directly [88, 89] under
Gaussian-Type Orbitals (GTO) as the basis $\\{\omega_{\mu}\\}_{\mu}$, using
software libraries _e.g._ libcint [90] in PySCF [79]. The term
$E_{\textnormal{XC}}[\rho_{\mathbf{p}}]$ is calculated by numerically
integrating the defined density $\rho_{\mathbf{p}}$ on a quadrature grid as
typically used in a DFT calculation. In our M-OFDFT,
$T_{\textnormal{S}}(\mathbf{p})$ is calculated directly from the coefficient
$\mathbf{p}$ using the deep-learning model
$T_{\textnormal{S},\theta}(\mathbf{p},\mathcal{M})$.

To carry out direct optimization using a learned KEDF model
$T_{\textnormal{S},\theta}(\mathbf{p})$, the gradient of the electronic energy
in Eq. (74) is required, which is given by:

$\displaystyle\nabla_{\mathbf{p}}E_{\theta}(\mathbf{p})=\nabla_{\mathbf{p}}T_{\textnormal{S},\theta}(\mathbf{p})+\underbrace{\tilde{\mathbf{W}}\mathbf{p}+\nabla_{\mathbf{p}}E_{\textnormal{XC}}(\mathbf{p})+\mathbf{v}_{\textnormal{ext}}}_{\nabla_{\mathbf{p}}E_{\textnormal{eff}}(\mathbf{p})}.$
(80)

This gradient is then used to update the density coefficient after projected
onto the linear subspace of normalized densities, following Eq. (3).

##### Relation to Derivation as Integral of the Variation with Basis

The gradient $\nabla_{\mathbf{p}}E(\mathbf{p})$ can also be derived by the
relation between gradient and variation that we have already seen in Eq. (69):