Learning Equivariant Non-Local Electron Density Functionals

## Abstract

Abstract The accuracy of density functional theory hinges on the approximation of non-local contributions to the exchange-correlation (XC) functional.
To date, machine-learned and human-designed approximations suffer from insufficient accuracy, limited scalability, or dependence on costly reference data.
To address these issues, we introduce Equivariant Graph Exchange Correlation (EG-XC), a novel non-local XC functional based on equivariant graph neural networks.
EG-XC combines semi-local functionals with a non-local feature density parametrized by an equivariant nuclei-centered point cloud representation of the electron density to capture long-range interactions.
By differentiating through a self-consistent field solver, we train EG-XC requiring only energy targets.
In our empirical evaluation, we find EG-XC to accurately reconstruct ‘gold-standard’ CCSD(T) energies on MD17.
On out-of-distribution conformations of 3BPA, EG-XC reduces the relative MAE by 35 % times 35 percent 35\text{\,}\mathrm{\char 37\relax} start_ARG 35 end_ARG start_ARG times end_ARG start_ARG % end_ARG to 50 % times 50 percent 50\text{\,}\mathrm{\char 37\relax} start_ARG 50 end_ARG start_ARG times end_ARG start_ARG % end_ARG .
Remarkably, EG-XC excels in data efficiency and molecular size extrapolation on QM9, matching force fields trained on 5 times more and larger molecules.
On identical training sets, EG-XC yields on average 51 % times 51 percent 51\text{\,}\mathrm{\char 37\relax} start_ARG 51 end_ARG start_ARG times end_ARG start_ARG % end_ARG lower MAEs.

## 1 Introduction

Kohn-Sham Density Functional Theory (KS-DFT) is the backbone of computational material and drug discovery (Jones, ). It is a quantum mechanical method to approximate the ground state energy of an ${N_{\text{el}}}$-electron system by finding the electron density $\rho\in\mathcal{D}_{{N_{\text{el}}}}=\left\{\rho:\mathbb{R}^{3}\to\mathbb{R}_{
+}|\int_{\mathbb{R}^{3}}\rho(r)dr={N_{\text{el}}}\right\}$ that minimizes the energy functional $E:\mathcal{D}_{{N_{\text{el}}}}\to\mathbb{R}$.
This functional maps electron densities to energies and is composed of

$$ $\displaystyle E[\rho]=T[\rho]+V_{\text{ext}}[\rho]+V_{\text{H}}[\rho]+E_{\text {XC}}[\rho]$ (1) $$

$T:\mathcal{D}_{{N_{\text{el}}}}\to\mathbb{R}$ is the kinetic energy functional, $V_{\text{ext}}:\mathcal{D}_{{N_{\text{el}}}}\to\mathbb{R}$ the external potential due to positively charged nuclei, $V_{\text{H}}:\mathcal{D}_{{N_{\text{el}}}}\to\mathbb{R}_{+}$ the Coulomb energy between electrons, and, finally, $E_{\text{XC}}:\mathcal{D}_{{N_{\text{el}}}}\to\mathbb{R}_{-}$ the *exchange-correlation* (XC) functional (Cramer, ).
While $T,V_{\text{ext}}$ and $V_{\text{H}}$ permit analytical computations, the exact form of $E_{\text{XC}}$ remains unknown, and its approximation frequently dominates DFT’s error (Kim et al., ).

Machine learning has emerged as a promising data-driven approach to approximate $E_{\text{XC}}$ (Kulik et al., ; Zhang et al., ).
Many such ML functionals adopt the classical approach and define $E_{\text{XC}}$ as the integral of a learnable XC energy density ${\epsilon_{\text{XC}}}:\mathbb{R}^{d_{\text{mGGA}}}\to\mathbb{R}$ (Perdew, ):

$$ $\displaystyle E_{\text{XC}}[\rho]=\int_{\mathbb{R}^{3}}\rho(r)\epsilon_{\text{ XC}}\left({\bm{g}}(r)\right)dr$ (2) $$

where ${\bm{g}}:\mathbb{R}^{3}\to\mathbb{R}^{d}$ are properties of the electron density $\rho$ (Dick & Fernandez-Serra, ; Nagai et al., ).
If ${\bm{g}}$ only depends on density quantities from its infinitesimal neighborhood, e.g., $\rho(r),\nabla\rho(r),\ldots$, the resulting functionals are called *semi-local*.
While this integrates well into existing quantum chemistry code, long-range interactions exceed the functional class (Kaplan et al., ).
In contrast, existing *non-local* functionals either scale poorly computationally (Zhou et al., ) or require costly reference data (Margraf & Reuter, ; Bystrom & Kozinsky, ).
The critical challenge in *non-local* XC functionals lies in efficiently capturing long-range interactions.

To this end, we present Equivariant Graph Exchange Correlation (EG-XC) transferring the success of equivariant graph neural networks (GNNs) (Batzner et al., ) to learning the non-local XC functional.
In EG-XC, we compress the electron density to equivariant point embeddings via a convolution of the electron density with an equivariant kernel at the nuclear positions.
On this point cloud, the equivariant GNN efficiently captures long-range information.
Given the final embeddings, we define a non-local feature density ${{\bm{g}}_{\text{NL}}}:\mathbb{R}^{3}\to\mathbb{R}^{d}$ that we feed into the XC energy density ${\epsilon_{\text{XC}}}$ from Equation 2.
Compared to previous approaches, our finite point cloud embeddings are neither dependent on the nuclear charge nor the basis set but purely derived from the density $\rho$.
To train on energies alone, we differentiate through the minimization of Equation 1 (Li et al., ).

In our experimental evaluation, EG-XC improves upon the learnable semi-local XC-functional on molecular dynamic trajectories, extrapolation to both out-of-distribution conformations and increasingly larger molecules.
In particular, we find EG-XC to reduce errors of the semi-local functional by a factor of 2 to 3, on par with accurate ML force fields combined with DFT calculations ($\Delta$-ML).
In extrapolation to unseen structures, we find EG-XC to accurately reproduce out-of-distribution potential energy surfaces unlike force fields (incl. $\Delta$-ML) while reducing the MAE by $35\text{\,}\mathrm{\char 37\relax}$ to $50\text{\,}\mathrm{\char 37\relax}$ compared to the next-best tested method.
Finally, we find EG-XC to have an excellent data efficiency, achieving similar accuracies to the best force fields with 5 times less data.
To summarize, we demonstrate that GNN-driven functionals push the frontier of non-local XC functionals and provide a promising path towards accurate and scalable DFT calculations.

## 2 Background

Notation. To denote functionals, i.e., functions mapping from functions to scalars, we write $F[f]$ where $F$ is the functional and $f$ the function.
We use superscripts in brackets^(t) to indicate sequences and regular superscripts of bold vectors, e.g., ${\bm{x}}^{l}$, to indicate the $l$-th irreducible representation of the SO(3) group.
We generally use $r\in\mathbb{R}^{3}$ to denote points in the 3D Euclidean space.
A notable exception is the finite set of nuclei positions $\{R_{1},\ldots,R_{N_{\text{nuc}}}\}$ for which we will use capital letters.
The norm of a vector is given by $\|{r}\|=\|r\|_{2}$.
For directions, i.e., unit length vectors, we use $\widehat{r}=\frac{r}{\|{r}\|}$.

Kohn-Sham density functional theory is the foundation of our work.
Here, we provide a very brief introduction.
For more details, we refer the reader to Appendix A or Lehtola et al. (). In KS-DFT, the electron density $\rho$ is represented by a set of orthogonal orbitals $\phi_{i}:\mathbb{R}^{3}\to\mathbb{R}$, $\phi_{i}(x)=C_{i}^{T}\bm{\chi}(x)$, that are defined as linear combinations of a basis set of atomic orbitals $\chi_{\mu}:\mathbb{R}^{3}\to\mathbb{R}$:

$$ $\displaystyle\rho(r)=\bm{\phi}(r)^{T}\bm{\phi}(r)=\bm{\chi}(r)^{T}CC^{T}\bm{ \chi}(r)=\bm{\chi}(r)^{T}P\bm{\chi}(r)$ (3) $$

where $P=CC^{T}\in\mathbb{R}^{{N_{\text{bas}}}\times{N_{\text{bas}}}}$ is the so-called density matrix.
Given the representation in a finite basis set, one can compute the kinetic energy $T$, external potential $V_{\text{ext}}$, and electron-electron repulsion energies $V_{\text{H}}$ analytically.
Unfortunately, such analytical expressions are generally unavailable for $E_{\text{XC}}[\rho]$ as in Equation 2.
Thus, one relies on numerical integration (Lehtola et al., ).

To minimize Equation 1, one typically uses the self-consistent field (SCF) method that we outline in greater detail in Appendix B.
The SCF method is an iterative two-step optimization where one first computes the so-called Fock matrix $F=\frac{\partial E}{\partial P}$ based on the current coefficients $C$.
Then, the optimal coefficients $C\in\mathbb{R}^{{N_{\text{bas}}}\times{N_{\text{el}}}}$ are obtained by solving the generalized eigenvalue problem

$$ $\displaystyle FC=SCE$ (4) $$

with $S_{\mu\nu}=\int_{\mathbb{R}^{3}}\chi_{\mu}(r)\chi_{\nu}(r)dr$ being the overlap matrix of the atomic orbitals and $E$ being a diagonal matrix.
The procedure is repeated until convergence of $P$.

Equivariance allows defining symmetries of functions.
A function $f:{\mathcal{X}}\to{\mathcal{Y}}$ is equivariant to a group ${\mathcal{G}}$ iff $\forall g\in{\mathcal{G}},x\in{\mathcal{X}},f\left(G_{g}^{\mathcal{X}}x\right)
=G_{g}^{\mathcal{Y}}f(x)$ where $G_{g}^{\mathcal{X}},G_{g}^{\mathcal{Y}}$ are the representations of $g$ in the domain ${\mathcal{X}}$ and codomain ${\mathcal{Y}}$, respectively.
Invariance is a special case of equivariance where $G^{\mathcal{Y}}_{g}=1$, i.e., $f\left(G_{g}^{\mathcal{X}}x\right)=f(x)$, for all $g\in{\mathcal{G}},x\in\mathcal{X}$.

The real spherical harmonics $Y^{l}:\mathbb{R}^{3}\to\mathbb{R}^{2l+1},l\in\mathbb{N}_{+},$ are a well-known example of SO(3)-equivariant functions as they transform under rotation according to the Wigner D matrices $D^{l}_{R}\in\mathbb{R}^{(2l+1)\times(2l+1)},\forall R\in SO(3)$, with $D^{0}_{R}=1$ being the identity and $D^{1}_{R}$ being 3D-rotation matrices:

$$ $\displaystyle Y^{l}(D^{1}_{R}x)=D^{l}_{R}Y^{l}(x).$ (5) $$

Such $2l+1$-dimensional equivariant functions $h_{i}:\mathbb{R}^{3}\to\mathbb{R}^{2l+1}$, e.g., $h_{i}=Y^{l}$, can be recombined using the tensor product

$$ $\displaystyle(h_{1}\otimes h_{2})^{l_{o}}_{m_{o}}=\sum_{l_{1},l_{2}}\sum_{m_{1 },m_{2}}C^{l_{o},l_{1},l_{2}}_{m_{o},m_{1},m_{2}}h^{l_{1}}_{m_{1}}h^{l_{2}}_{m _{2}},$ (6) $$

where $C\in\mathbb{R}^{(2l_{o}+1)\times(2l_{i}+1)\times(2l_{j}+1)}$ are the Clebsch-Gordan coefficients.
As the Wigner D-matrices are orthogonal, it follows that the inner product of two $l$-equivariant functions yields an invariant one

$$ $\displaystyle h^{l}_{i}\left(D^{1}_{R}r\right)^{T}h^{l}_{j}\left(D^{1}_{R}r \right)=h^{l}_{i}(r)^{T}D^{lT}_{R}D^{l}_{R}h^{l}_{j}(r)=h^{l}_{i}(r)^{T}h^{l}_ {j}(r).$ (7) $$

While, we are interested in $E(3)$-invariant functions, using such equivariant intermediate representations improves accuracy, expands the function class, and improves data efficiency (Schütt et al., ; Gasteiger et al., ; Batatia et al., ).

## 3 Related work

Learnable density functional approximations (DFAs) of the XC functional have a long history.
In their seminal paper, Kohn & Sham () proposed an ${\epsilon_{\text{XC}}}$ fitted to a reference calculation.
In this first parameterization ${\epsilon_{\text{XC}}}$ only depends on ${{\bm{g}}}(r)=\rho(r)$.
To improve accuracy, the density features ${{\bm{g}}}$ have been successively refined by including gradients and physical quantities such as the kinetic energy density $\tau(r)=\frac{1}{2}\sum_{i}^{{N_{\text{el}}}}|\nabla\phi_{i}(r)|^{2}$.
The resulting semi-local functionals are also known as meta-GGAs with ${{\bm{g}}}(r)={{\bm{g}}_{\text{mGGA}}}(r):=\left[\rho(r),|\nabla\rho(r)|,\tau(
r)\right]$ (Perdew, ).
Many ML functionals use this parameterization with ${\epsilon_{\text{XC}}}$ frequently being a product of an MLP with a classical functional (Nagai et al., ; Kulik et al., ; Zhang et al., ).
An appealing aspect of learning semi-local energy densities rather than $E_{\text{XC}}[\rho]$ directly is the known implementation of physical constraints (Sun et al., ; Kaplan et al., ), which inspired works by Dick & Fernandez-Serra (); Nagai et al. ().
While augmenting ${{\bm{g}}_{\text{mGGA}}}$ with hand-crafted non-local features improves the accuracy, this comes at the cost of exact constraints (Nagai et al., ; Kirkpatrick et al., ).
Similarly, such physical biasses are absent in kernel methods Margraf & Reuter (); Bystrom & Kozinsky ().
Additionally, these methods require reference densities that are rarely available (Szabo & Ostlund, ).
Alternative grid-based CNNs converge slowly in the number of grid points compared to specialized spherical integration grids (Zhou et al., ; Treutler & Ahlrichs, ).
Lastly, machine-learned functionals have been used in other DFT contexts as well, e.g., for energy corrections on SCF-converged energies (Chen & Yang, ) or to model the kinetic energy functional in orbital-free DFT (Snyder et al., ; Mi et al., ; Zhang et al., ).
Here, we propose to extend semi-local models with expressive fully learnable non-local features obtained from standard integration grids.
This allows us to leverage the physical biases of meta-GGAs while allowing efficient non-local interactions with fast integration.
Although this work focuses on the XC functional, the methods we present can be transferred to these applications as well.

Machine learning force fields
have long functioned as a cheap but inaccurate alternative to quantum mechanical calculations by directly approximating the potential energy surface from data without solving the electronic structure problem (Unke et al., ).
Contemporary force fields generally rely on graph neural networks (GNNs) (Schütt et al., ) representing molecules in terms of graphs with atoms as nodes.
The advent of SO(3)-equivariant models led to significant improvements in accuracy and data efficiency, closing the gap to full-fidelity quantum mechanical calcualtions (Batzner et al., ; Batatia et al., ).
While they effectively aim to accomplish the same goal as DFT at a fraction of the cost, their out-of-distribution accuracy is a common problem (Stocker et al., ).
In EG-XC, we transfer the success of equivariant GNNs to DFT, by combining their approximation power with the physical nature of DFT, resulting in an accurate and data-efficient method.

## 4 Equivariant Graph Exchange Correlation

Figure: Figure 1: Illustration of Equivariant Graph Exchange Correlation (EG-XC)’s four components: (1) We first compress the electron density into a finite point cloud representation ${\bm{\mathsfit{H}}}^{(0)}$ via a convolution with radial filters $\Gamma_{k}:\mathbb{R}_{+}\to\mathbb{R}$ and spherical harmonics $Y^{l}_{m}:\mathbb{R}^{3}\to\mathbb{R}$ evaluated at the nuclear position. (2) The embeddings ${\bm{\mathsfit{H}}}^{(0)}$ are updated using equivariant message passing to obtain long-range effects in ${\bm{\mathsfit{H}}}^{(T)}$. (3) We define a non-local feature density ${{\bm{g}}_{\text{NL}}}:\mathbb{R}^{3}\to\mathbb{R}^{d}$ from which we derive the exchange correlation density ${\epsilon_{\text{XC}}}:\mathbb{R}^{3}\to\mathbb{R}$. (4) We add a graph readout of ${\bm{\mathsfit{H}}}^{(T)}$ to learn additional corrections. To obtain $E_{\text{XC}}[\rho]$, we integrate ${\epsilon_{\text{XC}}}$ and add it to the global graph readout.
Refer to caption: x1.png

With Equivariant Graph Exchange Correlation (EG-XC), we propose an efficient *non-local* approximation to the unknown exchange-correlation functional $E_{\text{XC}}$ that maps electron densities (positive integrable functions in $\mathbb{R}^{3}$) to scalar energies.
To this end, EG-XC consists of four components as illustrated in Figure 1:
(1) *Nuclei-centered equivariant embeddings* ${\bm{\mathsfit{H}}}^{(0)}$ compress the electron density $\rho$ into a finite point cloud representation by equivariantly integrating the density around the nuclei.
(2) *Equivariant message passing* on this point cloud includes long-range information in ${\bm{\mathsfit{H}}}^{(T)}$.
Based on our final point cloud embeddings ${\bm{\mathsfit{H}}}^{(T)}$, we define $E_{\text{XC}}$ as the sum of two readouts:
(3) A *non-local reweighted meta-GGA* based on our non-local feature density ${{\bm{g}}_{\text{NL}}}:\mathbb{R}^{3}\to\mathbb{R}^{d}$ and
(4) a *graph readout* of the invariant nuclei-centered features ${\bm{\mathsfit{H}}}^{(T)0}$:

$$ $\displaystyle E_{\text{XC}}[\rho]=\underbrace{\int_{\mathbb{R}^{3}}\rho(r) \underbrace{{\gamma_{\text{NL}}}\left({{\bm{g}}_{\text{NL}}}(r),{{\bm{g}}_{ \text{mGGA}}}(r)\right)}_{\text{non-local weights}}\underbrace{\epsilon_{\text {mGGA}}\left({{\bm{g}}_{\text{mGGA}}}(r)\right)}_{\text{meta-GGA}}dr}_{\text{ non-local reweighted meta-GGA}}+\underbrace{\sum_{i=1}^{N_{\text{nuc}}}{ \epsilon_{\text{NL}}}\left({\bm{H}}^{(T)0}_{i}\right)}_{\text{graph readout}}.$ (8) $$

(1) Nuclei-centered equivariant embeddings ${\bm{\mathsfit{H}}}^{(0)}\in\left[\mathbb{R}^{1}\times\ldots\times\mathbb{R}^{
2{l_{\text{max}}}+1}\right]^{{N_{\text{nuc}}}\times d}\coloneqq{\mathcal{F}}_{
l_{\text{max}}}^{{N_{\text{nuc}}}\times d}$ reduce the computational scaling by mapping the continuous density to a finite per-nucleus representation.
Further, these enable using equivariant GNNs to learn long-range interactions efficiently.
The positions of the nuclei $R_{i}\in\mathbb{R}^{3}$ naturally lend themselves as centroids for such embeddings as they represent peaks of the electron density (Kato, ).
To perform this reduction, we evaluate a convolution of the electron density $\rho$ with equivariant filters $S^{l}_{k}:\mathbb{R}^{3}\to\mathbb{R}^{2l+1}$ at $R_{i}$:

$$ $\displaystyle H_{ik}^{(0)l}=\left(\rho_{i}*S^{l}_{k}\right)(R_{i})$ (9) $$

where $\rho_{i}:\mathbb{R}^{3}\to\mathbb{R}_{+}$ is the soft-partitioned electron density associated with the $i$-th embedding

$$ $\displaystyle\rho_{i}(r)=\rho(r)\frac{\alpha_{i}(r)}{\sum_{j=1}^{{N_{\text{nuc }}}}\alpha_{j}(r)},$ (10) $\displaystyle\alpha_{i}(r)=\exp\left(-\frac{\|{r-R_{i}}\|^{2}}{\lambda^{2}}\right)$ (11) $$

with $\lambda\in\mathbb{R}_{+}$ being a free parameter.
Such partitioning prevents overcounting if nuclei are close similar to quadrature methods (Becke, ).
Like equivariant GNNs (Batzner et al., ), we define the filters $S^{l}_{k}$ as a product of spherical harmonics $Y^{l}:\mathbb{R}^{3}\to\mathbb{R}^{2l+1}$ and radial filters $\Gamma_{k}:\mathbb{R}_{+}\to\mathbb{R}$:

$$ $\displaystyle S^{l}_{k}(r)=\Gamma_{k}(\|{r}\|)Y^{l}\left(\widehat{r}\right)$ (12) $$

where $0\leq l\leq{l_{\text{max}}}$ indexes the real spherical harmonics.
The spherical harmonics allow us to capture angular changes in the electron density that would otherwise be lost in a purely radial representation.
Put together, we compute the nuclei-centered equivariant embeddings

$$ $\displaystyle\begin{split}H_{ik}^{(0)l}&=\int_{\mathbb{R}^{3}}\rho_{i}(r) \Gamma^{l}_{k}(\|{r-R_{i}}\|)Y^{l}\left(\widehat{r-R_{i}}\right)dr\\ &\approx\sum_{j=1}^{N_{\text{quad}}}w_{j}\rho(r_{j})\alpha_{i}(r_{j})\Gamma_{k }(\|{r_{j}-R_{i}}\|)Y^{l}\left(\widehat{r_{j}-R_{i}}\right)\end{split}$ (13) $$

with a set of ${N_{\text{quad}}}$ standard integration points and weights $\{(r_{j},w_{j})\}_{j=1}^{N_{\text{quad}}}\in\left[\mathbb{R}^{3}\times\mathbb{
R}\right]^{N_{\text{quad}}}$ (Treutler & Ahlrichs, ).
Importantly, while these embeddings are centered at the nuclei, they do not embed the nuclear charges, as in ML force fields (Schütt et al., ), but the electron density around them.
This is important as the derivative with respect to the electron density affects the SCF procedure; see Appendix B.
A nuclear charge embedding’s derivative would not exist and, thus, not alter the DFT calculation, effectively yielding a force field.

We follow Schütt et al. () and parametrize the radial filters $\Gamma_{k}$ as a combination of $\sin,\cos$ with a fixed polynomial envelope function $u:\mathbb{R}_{+}\to\mathbb{R}$ from Gasteiger et al. () for a cutoff $c\in\mathbb{R}_{+}$:

$$ $\displaystyle\Gamma_{k}(r)=\begin{cases}\sin\left(\frac{r\pi k}{2c}\right)u \left(\frac{r}{c}\right)&\text{if }k\text{ even},\\ \cos\left(\frac{r\pi(k+1)}{2c}\right)u\left(\frac{r}{c}\right)&\text{if }k \text{ odd}.\\ \end{cases}$ (14) $$

We found such frequency filters to yield more stable results than Bessel functions as the latter ones are very sensitive close to the centroids (Gasteiger et al., ).
This sensitivity might be beneficial in force fields but is implicitly given in DFT due to the higher density close to the nuclei.

(2) Equivariant message passing allows for the propagation and updating of the equivariant electron density features ${\bm{\mathsfit{H}}}^{(0)}$ to capture long-range dependencies within the electron density like van-der-Waals forces (Frank et al., ).
The equivariance of these features is important as it allows us to define a non-radial-symmetric feature density ${{\bm{g}}_{\text{NL}}}$ in the next step.
We use Batzner et al. ()’s NequIP to perform message passing with the SO(3)-equivariant convolution from Thomas et al. () to iteratively update the embeddings ${\bm{\mathsfit{H}}}^{(t)}$ over $T$ steps:

$$ $\displaystyle{\bm{H}}_{i}^{(t+1)}$ $\displaystyle=\text{EquiMLP}\left(\text{EquiLin}\left({\bm{H}}_{i}^{(t)}\right )+\text{EquiConv}\left({\bm{\mathsfit{H}}}^{(t)},{\bm{R}}\right)_{i}\right)+ \text{EquiLin}\left({\bm{H}}^{(t)}_{i}\right),$ (15) $$

where $\text{EquiLin}:{\mathcal{F}}_{l_{\text{max}}}^{d}\to{\mathcal{F}}_{l_{\text{
max}}}^{d}$ is an equivariant dense layer that mixes features of the same order $l$

$$ $\displaystyle\text{EquiLin}({\bm{X}})_{k}=\text{Concat}\left(\left[\sum_{k^{ \prime}=1}^{d}W^{l}_{k^{\prime}k}X^{l}_{k^{\prime}}\right]_{l=0}^{{l_{\text{ max}}}}\right).$ (16) $$

with each ${\bm{W}}^{l}\in\mathbb{R}^{d\times d}$ being a learnable weight matrix.
$\text{EquiMLP}:{\mathcal{F}}_{l_{\text{max}}}^{d}\to{\mathcal{F}}_{l_{\text{
max}}}^{d}$ is the equivariant equivalent of a standard MLP, where we use the invariant embeddings $l=0$ and an activation function $\sigma:\mathbb{R}\to\mathbb{R}$ to gate the equivariant parts $l>0$:

$$ $\displaystyle\text{EquiMLP}({\bm{X}})$ $\displaystyle=\text{EquiLin}(\text{ActEquiLin}({\bm{X}})),$ (17) $\displaystyle\text{ActEquiLin}({\bm{X}})$ $\displaystyle=\text{Concat}\left(\sigma(W^{0}{\bm{X}}^{0}),\left[\text{EquiLin }({\bm{X}})^{l}\circ\sigma\left(W^{l}{\bm{X}}^{0}\right)\right]_{l=1}^{{l_{ \text{max}}}}\right).$ (18) $$

Lastly, we define the convolution $\text{EquiConv}:{\mathcal{F}}_{l_{\text{max}}}^{{N_{\text{nuc}}}\times d}
\times\mathbb{R}^{{N_{\text{nuc}}}\times 3}\to{\mathcal{F}}_{l_{\text{max}}}^{
{N_{\text{nuc}}}\times d}$ via the tensor product:

$$ $\displaystyle\begin{split}\text{EquiConv}({\bm{X}},\mathbf{R})^{l_{o}}_{t}= \sum_{s=1}^{N_{\text{nuc}}}\sum_{l_{i},l_{f}=0}^{l_{\text{max}}}&\text{MLP} \left(\Gamma(\|{R_{s}-R_{t}}\|)\right)_{l_{o},l_{i},l_{f}}\\ &\cdot\left(\text{EquiMLP}({\bm{x}}_{s})^{l_{i}}\otimes Y^{l_{f}}\left( \widehat{R_{s}-R_{t}}\right)\right)^{l_{0}}.\end{split}$ (19) $$

To accelerate the computation of the tensor product, we use Passaro & Zitnick ()’s efficient equivariant convolutions.
For more details, we refer the reader to Batzner et al. ().

(3) Non-local reweighted meta-GGA. As the majority of the XC energy can be captured by semi-local functionals (Goerigk et al., ), starting with a machine-learned meta-GGA $\epsilon_{\text{mGGA}}:\mathbb{R}^{d_{\text{mGGA}}}\to\mathbb{R}$ accounts for most of the XC energy.
To correct the meta-GGA for non-local effects, we first define a non-local feature density ${{\bm{g}}_{\text{NL}}}:\mathbb{R}^{3}\to\mathbb{R}^{d}$ based on the non-local embeddings $\left\{{\bm{\mathsfit{H}}}^{(t)}\right\}^{T}_{t=0}$:

$$ $\displaystyle{{\bm{g}}_{\text{NL}}}(r)_{k}=\sum_{i=1}^{N_{\text{nuc}}}\alpha_{ i}(r)\sum_{t=0}^{T}\sum_{l=0}^{{l_{\text{max}}}}\underbrace{Y^{l}\left( \widehat{r-R_{i}}\right)^{T}H_{ik}^{(t)l}}_{\text{angular}}\underbrace{\Gamma( \|{r-R_{i}}\|)^{T}{\bm{w}}^{(t)l}_{k}\left({\bm{H}}^{(t)0}_{i}\right)}_{\text{ radial}}$ (20) $$

where ${\bm{w}}_{k}^{(t)l}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ are MLPs mapping to radial weights.
While the inner product between the spherical harmonics $Y^{l}$ and our equivariant embeddings ${\bm{\mathsfit{H}}}^{(t)l}$ expresses angular changes, the inner product of radial basis functions $\Gamma$ and ${\bm{w}}_{k}^{(t)l}$ allows for radial changes.
From this feature density ${{\bm{g}}_{\text{NL}}}$ and standard meta-GGA inputs ${{\bm{g}}_{\text{mGGA}}}:\mathbb{R}^{3}\to\mathbb{R}^{d_{\text{mGGA}}}$, we derive the non-local correction ${\gamma_{\text{NL}}}:\mathbb{R}^{d+{d_{\text{mGGA}}}}\to\mathbb{R}$ to the exchange energy density $\epsilon_{\text{mGGA}}$:

$$ $\displaystyle{\epsilon_{\text{XC}}}(r)={\gamma_{\text{NL}}}\left({{\bm{g}}_{ \text{NL}}}(r),{{\bm{g}}_{\text{mGGA}}}(r)\right)\cdot\epsilon_{\text{mGGA}} \left({{\bm{g}}_{\text{mGGA}}}(r)\right).$ (21) $$

In practice, we implement ${\gamma_{\text{NL}}}$ as an MLP.
To obtain the final readout, we integrate the exchange-correlation density over the electron density $\int_{\mathbb{R}^{3}}{\epsilon_{\text{XC}}}(r)\rho(r)dr$.
Like Equation 13, we evaluate the integral with standard integration grids.

(4) Graph readout. In addition to the meta-GGA-based readout, we add global graph readout of the embeddings $\left\{{\bm{\mathsfit{H}}}^{(t)}\right\}^{T}_{t=0}$ to capture the remaining non-local effects.
We use an MLP ${\epsilon_{\text{NL}}}:\mathbb{R}^{d}\to\mathbb{R}$ on top of the invariant embeddings ($l=0$) to obtain the final exchange-correlation energy:

$$ $\displaystyle E_{\text{XC}}[\rho]=\int_{\mathbb{R}^{3}}{\epsilon_{\text{XC}}}( r)\rho(r)dr+\sum_{t=0}^{T}\sum_{i=1}^{M}{\epsilon_{\text{NL}}}\left({\bm{H}}^{ (t)0}_{i}\right).$ (22) $$

Limitations. While EG-XC demonstrates significant improvements over a semi-local ML functional, s. Section 5, it is not free of limitations.
First, as we rely on the nuclear positions to represent the electronic density, it is not truly universal, i.e., independent of the external potential $V_{\text{ext}}$ (Kohn & Sham, ).
Second, the non-local nature of our functional permits no known way to enforce most physical constraints (Kaplan et al., ).
These missing constraints enable the correction of basis set errors through the XC functional.
While this may lead to unphysical matches between densities and energies, our experiments suggest that this does not lead to overfitting energies.
Third, systems without nuclei, e.g., the homogenous electron gas, cannot be modeled with our approach.
In such cases, one may want to replace the real-space point cloud with a frequency representation (Kosmala et al., ).
Lastly, while more data efficient and better in extrapolation, running KS-DFT is more expensive than a surrogate force field.

## 5 Experiments

In the following, we compare EG-XC to various alternative methods of learning potential energy surfaces across several settings.
In particular, we focus on the following tasks: interpolating accurate energy surfaces, extrapolation to unseen conformations, and extrapolation to larger molecules.

Methods. To accurately position EG-XC, we compare to methods of varying computational costs.
In particular, we evaluate force fields, $\Delta$-ML, i.e., combining force fields with KS-DFT calculations, and a learnable XC-functional (Dick & Fernandez-Serra, ).
While force fields are orders of magnitude cheaper as they bypass quantum mechanical calculations altogether, they lack prior physical knowledge.
For the $\Delta$-ML methods, we shift all energies by DFT energies with LDA in the STO-6G basis.
This reduces the learning problem to the difference between the KS-DFT and the target energy (Wengert et al., ).
As force fields, we test increasingly expressive models based on their use of SO(3) irreducible representations: SchNet ($l=0$) (Schütt et al., ), PaiNN ($l=1$) (Schütt et al., ), and NequIP ($l=2$) (Batzner et al., ).
Finally, we compare with the same learnable semi-local XC-functional (Dick & Fernandez-Serra, ) used in EG-XC.
As $\Delta$-ML methods and learnable XC-functionals require a DFT computation for each structure, they are at the same computational cost as EG-XC.

Setup. To train the XC functionals, we follow Li et al. () and implement the SCF method differentiably.
This allows us to match the converged SCF energies directly to the target energies without needing ground truth electron densities.
We provide implementation details in Appendix C.
All methods are trained on energy labels only.
Force fields are trained with hyperparameters from their respective works with modifications to learning rate, batch size, and initialization to improve performance; we outline the changes in Appendix G.
We list EG-XC’s hyperparameters in Appendix F.

**Table 1: Test set MAE on the CCSD(T) MD17 dataset in $\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$. (best, second)**
|  | Force field | $\Delta$-ML | KS-DFT |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Molecule | SchNet | PaiNN | NequIP | SchNet | PaiNN | NequIP | Dick | EG-XC |
| Aspirin | 7.01 | 2.82 | 5.52 | 2.02 | 1.20 | 1.04 | 1.94 | 0.69 |
| Benzene | 0.40 | 0.16 | 0.09 | 0.13 | 0.11 | 0.02 | 0.39 | 0.10 |
| Ethanol | 1.41 | 0.89 | 0.95 | 0.93 | 0.42 | 0.25 | 0.85 | 0.21 |
| Malonaldehyde | 2.10 | 1.00 | 2.32 | 0.61 | 0.44 | 0.29 | 0.73 | 0.27 |
| Toluene | 1.80 | 1.10 | 1.87 | 0.44 | 0.31 | 0.13 | 0.38 | 0.20 |

Reproducing gold-standard accuracies.
The objective behind learning XC-functionals or ML force fields is to facilitate access to accurate potential energy surfaces by distilling highly accurate reference data into a faster method.
Hence, we compare these methods on the revised MD17 dataset, which contains precise ‘gold-standard’ CCSD(T) (CCSD for aspirin) reference energies for conformations of five molecules along the trajectory of a molecular dynamic (MD) simulation (Chmiela et al., ).
Each molecule has a training set of 1000 structures, which we split into 950 training and 50 validation structures.
Each test set contains an additional 500 structures (1000 for ethanol).
Following Schütt et al. (), a separate model is fitted per molecule.
Since CCSD(T) calculations inherently account for non-local effects, these datasets are well suited for investigating the ability to accurately interpolate multi-dimensional energy surfaces from a limited number of reference structures.
In Appendix I, we provide additional $\Delta$-ML data with a more accurate DFT functional and basis sets.

Table 1 lists the mean absolute error (MAE) for all methods on the MD17 test sets.
We find that force fields generally struggle to reconstruct the energy surface of the MD17 dataset accurately when trained solely on energies.
In contrast, $\Delta$-ML methods and learnable XC-functionals generally achieve chemical accuracy at $$1\text{\,}\mathrm{kcal}\text{\,}{\mathrm{mol}}^{-1}$\approx$1.6\text{\,}
\mathrm{m\text{$E$}_{\mathrm{h}}}$$.
The DFT reference calculations systematically reduce the to-be-learned energy surface delta from $6.2\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$ to $4.9\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$.
Compared to the EG-XC’s base semi-local Dick & Fernandez-Serra (), EG-XC reduces errors by a factor of 2 to 4.
Among all methods, EG-XC yields the lowest error on 3 of the 5 molecules.
In Appendix H, we test the performance on a reduced training set of just 50 samples.

**Table 2: Structural extrapolation on the 3BPA dataset. All methods were trained on the 300K training set. All numbers are relative MAE in $\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$. The last three rows refer to the potential energy surfaces.**
|  | Force field | $\Delta$-ML | KS-DFT |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Test set | SchNet | PaiNN | NequIP | SchNet | PaiNN | NequIP | Dick | EG-XC |
| 300K | 5.15 | 2.91 | 3.81 | 2.38 | 1.14 | 0.81 | 0.96 | 0.42 |
| 600K | 9.06 | 5.81 | 7.55 | 3.96 | 2.13 | 1.56 | 1.36 | 0.73 |
| 1200K | 18.33 | 14.14 | 17.30 | 6.84 | 5.97 | 3.30 | 2.27 | 1.39 |
| $\beta=$120\text{\,}\mathrm{\SIUnitSymbolDegree}$$ | 3.84 | 1.78 | 2.25 | 2.53 | 1.25 | 1.09 | 0.75 | 0.35 |
| $\beta=$150\text{\,}\mathrm{\SIUnitSymbolDegree}$$ | 4.64 | 1.89 | 2.64 | 2.03 | 0.84 | 0.88 | 0.61 | 0.23 |
| $\beta=$180\text{\,}\mathrm{\SIUnitSymbolDegree}$$ | 4.97 | 1.92 | 3.03 | 1.79 | 1.06 | 0.73 | 0.56 | 0.20 |

Figure: Figure 2: Two-dimensional slice of the potential energy surface of 3BPA with the dihedral $\beta=$120\text{\,}\mathrm{\SIUnitSymbolDegree}$$. Pure force fields like NequIP struggle to recover the shape of this out-of-distribution energy surface. When paired with DFT calculations, one can see that the energy surface moves closer to the target shape but introduces additional extrema. Learnable XC functionals like Dick & Fernandez-Serra () and EG-XC demonstrate significantly better reproduction of the target energy surface.
Refer to caption: x2.png

Extrapolating structures.
In MD simulations, one often encounters structures far outside the training set, e.g., due to higher temperatures or environmental changes (Stocker et al., ).
To investigate the extrapolation to unseen structures, we use the 3BPA dataset (Kovács et al., ).
3BPA contains various geometric configurations of the molecule 3-(benzyloxy)pyridin-2-amine.
The training set consists of 500 structures sampled from an MD simulation at room temperature (300K).
The test sets consist of MD trajectories at 300K, 600K, and 1200K to test in and out-of-distribution (OOD) performance.
Additionally, the dataset contains three 2-dimensional potential energy surface slices where two dihedral angles are varied while keeping one fixed.
The labels have been computed with the hybrid $\omega$B97X XC functional and the 6-31G(d) basis set and, thus, include non-local interactions.

As constant offsets of the potential energy surface do not affect the system dynamics, we evaluate the relative MAE, i.e., $\mathbb{E}\left[\left|E_{\text{pred}}-E_{\text{target}}-\textrm{median}(E_{
\text{pred}}-E_{\text{target}})\right|\right]$.
We list the relative MAE for all methods and test sets in Table 2; for the absolute MAE, we refer the reader to Appendix J.
Across all test sets, EG-XC results in $35\text{\,}\mathrm{\char 37\relax}$ to $51\text{\,}\mathrm{\char 37\relax}$ lower errors than the next best-tested alternative.
On the far OOD 1200K samples, EG-XC is the only method achieving chemical accuracy $1.6\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$ at an relative MAE of $1.40\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$.
To illustrate the qualitative improvement of EG-XC’s energy surfaces, we plot the most OOD potential energy slice at $\beta=$120\text{\,}\mathrm{\SIUnitSymbolDegree}$$ in Figure 2 for NequIP, $\Delta$-NequIP, Dick & Fernandez-Serra (), EG-XC and the target.
The remaining methods and corresponding energy surfaces are plotted in Appendix K.
It is evident that force fields fail to reproduce the target energy surface with no resemblance to the target surface.
While the reference DFT calculations move $\Delta$-NequIP closer to the target surface, the energy shape includes additional extrema not present in the target surface.
In contrast, both XC functionals accurately reproduce the energy surface.

Figure: Figure 3: MAE in $\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$ on QM9 size extrapolation. Each row represents a different method. The four groups indicate the maximum number of heavy atoms in the training set and its size. Each column represents the test subset with the number of heavy atoms listed above.
Refer to caption: x3.png

Extrapolation to larger molecules. Gathering reference data for large compounds is costly and expensive with accurate quantum chemical calculations like CCSD(T) scaling $O({N_{\text{el}}}^{7})$ in the number of electrons ${N_{\text{el}}}$.
Thus, extrapolation from small or medium-sized molecules to larger ones is critical.
Here, we simulate this setting by splitting the QM9 dataset (Ramakrishnan et al., ) into subsets of increasing size based on the number of heavy atoms, i.e., QM9($S$) are all QM9 molecules with at most $S$ heavy atoms.
For each $S\in\{4,5,6,7\}$, we train a separate model and test on the remaining structures, i.e., QM9$\setminus$QM9($S$).
For each training set, we split the structures 90%/10% into training and validation sets.
The QM9 dataset comprises 134k stable organic molecules with up to 9 heavy atoms.
The energies have been computed with the hybrid B3LYP XC functional with the 6-31G(2df,p) basis set and, thus, contain non-local interactions through the exact Hartree exchange.
As the dataset contains only few molecules with fluorine, force fields (including $\Delta$-ML) could not yield accurate energies if none or few are in the training set.
Thus, we omitted all molecules that include fluorine.

We visualize the MAE for each combination of the number of heavy atoms in the test set and the maximum number of heavy atoms in the training set in Figure 3.
Across all combinations of training and test sets, we find learnable XC functionals yielding the lowest errors with a preference for Dick & Fernandez-Serra () on the smallest QM9(4).
We hypothesize that this is due to the physical constraints enforced by the learnable XC functional, which are not present in the other methods, including EG-XC.
Starting with QM9(5), EG-XC’s errors are consistently the lowest.
Notably, EG-XC trained on QM9(6) yields lower MAE on the largest structures than the best alternative on QM9(7) with 5 times more samples and one-atom larger molecules.
On QM9(7), EG-XC is at least $33\text{\,}\mathrm{\char 37\relax}$ more accurate than the competing methods on test molecules with 9 heavy atoms.
Compared to Dick & Fernandez-Serra (), EG-XC reduces the MAE by at least $2\times$ on QM9(6) and QM9(7).

## 6 Discussion

We tackled the problem of learning efficient non-local exchange-correlation functionals to enhance the accuracy of density functional theory.
To this end, we have presented Equivariant Graph Exchange Correlation (EG-XC), an equivariant graph neural network approach.
We have introduced a basis-independent reduction of the electron density into a finite point cloud representation, enabling equivariant graph neural networks to capture long-range information.
Based on this point cloud embedding, we have parameterized a non-local feature density used for the reweighing of a semi-local exchange-correlation energy density.
Unlike other learnable non-local functionals, EG-XC can be trained by differentiating through the SCF solver such that the training requires only energy targets.

In our empirical evaluation, EG-XC improves upon previous machine learning-based methods of modeling potential energy surfaces.
On MD17, EG-XC accurately reconstructs ‘gold-standard’ potential energy surfaces, namely CCSD(T), within the DFT framework.
On the 3BPA, EG-XC reduces errors by $35\text{\,}\mathrm{\char 37\relax}$ to $50\text{\,}\mathrm{\char 37\relax}$ compared to the best-performing baseline.
On QM9, EG-XC demonstrates remarkable data efficiency, achieving similar accuracies with 5 times less data or up to $33\text{\,}\mathrm{\char 37\relax}$ lower errors at the same amount of data compared to the best baseline.
Finally, EG-XC extrapolates well to unseen conformations and larger molecules on 3BPA and QM9, respectively.

Overall, these results strongly underline the data efficiency of learning exchange-correlation functionals compared to machine-learned force fields.
We hypothesize that a significant portion of EG-XC’s generalization is thanks to including the physically constrained semi-local model that captures most of the exchange-correlation energy and biases EG-XC to approximately fulfill the same constraints.
On top of these, our non-local contributions transfer the accuracies of graph neural network-based force fields to functionals while maintaining a similar data efficiency.

Future work.
The simple implementation of EG-XC through standard deep learning libraries and the integration with the self-consistent field method open the door to a range of future research.
Thanks to the basis-set independence, EG-XC can transfer to non-atom-centered basis sets like plane waves, as they are common in periodic systems, or approximate the kinetic energy functional in orbital-free DFT.
Further, while we trained EG-XC solely on energies, multimodal training with other DFT-computable observables like electron densities or atomic forces could further improve accuracy and generalization.
Another path to improving generalization may be integrating known physical constraints to the non-local part of EG-XC.
Finally, accurate functionals like EG-XC may bridge the gap between sparse, accurate quantum mechanical calculations (Cheng et al., ; Gao & Günnemann, ) and fast force field (Batzner et al., ).

Acknowledgements. We are grateful to Leo Schwinn and Jan Schuchardt for their invaluable manuscript feedback, and to Arthur Kosmala and Johannes Margraf for insightful discussions on density functional theory.
Funded by the Federal Ministry of Education and Research (BMBF) and the Free State of Bavaria under the Excellence Strategy of the Federal Government and the Länder.

## Appendix A Kohn-Sham density functional theory

In Kohn-Sham density functional theory (KS-DFT) (Kohn & Sham, ), one attempts to approximate the ground state energy of a molecular system by finding the electron density that minimizes the energy functional Equation 1.
Compared to ab-initio methods like Gao & Günnemann (, , , ); Wilson et al. (, ), KS-DFT requires the empirical approximation of the XC functional $E_{\text{XC}}$ but scales more favorably with system size.
Here, we largely compress Lehtola et al. () to provide a brief introduction to the topic.
We will neglect electron spin for simplicity and focus on a discretized DFT scheme based on atom-centered basis sets.
Operators, i.e., functions mapping from functions to functions, are denoted by acting upon the ket, e.g., $\nabla|f\rangle$ is the derivative of $f$.
Further, we use the braket notation to denote the inner product of two functions $f$ and $g$ as $\langle f|g\rangle=\int g(x)f(x)dx$.
We will use Einstein’s summation convention, where indices that only appear on one side of an equation are summed over.

The electron density $\rho$ of an ${N_{\text{el}}}$ system is associated with an antisymmetric wave function of a fictitious non-interacting system consisting of ${N_{\text{el}}}$ orthonormal orbitals $\phi_{i}:\mathbb{R}^{3}\to\mathbb{R}$. These orbitals are constructed as linear combinations of ${N_{\text{bas}}}$ a finite basis set $\chi_{\mu}:\mathbb{R}^{3}\to\mathbb{R}$ such that $\phi_{i}(r)=C_{\mu i}\chi_{\mu}(r)$. For systems without periodic boundary conditions, these basis sets typically consist of nuclei-centered functions called atomic orbitals. The electron density and pseudo-wave-function are given by

$$ $\displaystyle\Psi(\mathbf{r})$ $\displaystyle=\frac{1}{\sqrt{N!}}\det\left[\phi_{i}(r_{j})\right]_{i,j=1}^{N},$ (23) $\displaystyle\begin{split}\rho(r)&=\int\Psi(r,r_{2},...,r_{N})^{*}\Psi(r,r_{2} ,...,r_{N})dr_{2}...dr_{N}\\ &=\bm{\phi}(r)^{T}\bm{\phi}(r)=\bm{\chi}(r)^{T}C^{T}C\bm{\chi}(r)=\bm{\chi}(r) ^{T}P\bm{\chi}(r),\end{split}$ (24) $$

where $C\in\mathbb{R}^{{N_{\text{bas}}}\times{N_{\text{el}}}}$ are the orbital coefficients and $P=C^{T}C\in\mathbb{R}^{{N_{\text{bas}}}\times{N_{\text{bas}}}}$ is the so-called density matrix.
This construction ensures that the electron density belongs to the class of densities expressible by antisymmetric wave functions and enables the exact computation of the kinetic energy (Kohn & Sham, )

$$ $\displaystyle T[\rho]=-\frac{1}{2}\langle\phi_{i}|\nabla^{2}|\phi_{i}\rangle=- \frac{1}{2}C_{\mu i}C_{\nu i}\langle\chi_{\mu}|\nabla^{2}|\chi_{\nu}\rangle=C_ {\mu i}C_{\nu i}T_{\mu\nu}=P_{\mu\nu}T_{\mu\nu},$ (25) $$

making KS-DFT significantly more accurate than contemporary orbital-free DFT.

The external potential in molecular systems is defined by the nuclei positions $\mathbf{R}$ and charges $\mathbf{Z}$

$$ $\displaystyle V_{\text{ext}}[\rho]=-\int\frac{\rho(r)Z_{n}}{|r-R_{n}|}dr=-C_{ \mu i}C_{\nu i}\int\frac{\chi_{\mu}(r)Z_{n}\chi_{\nu}(r)}{|r-R_{n}|}dr=C_{\mu i }C_{\nu i}V_{\mu\nu}=P_{\mu\nu}V_{\mu\nu}.$ (26) $$

The Hartree energy is a mean-field approximation to the electron-electron interaction

$$ $\displaystyle V_{\text{H}}[\rho]$ $\displaystyle=\frac{1}{2}\iint\frac{\rho(r)\rho(r^{\prime})}{|r-r^{\prime}|} drdr^{\prime}=C_{\mu i}C_{\nu i}C_{\lambda j}C_{\kappa j}\iint\chi_{\mu}(r) \chi_{\lambda}(r^{\prime})\frac{1}{|r-r^{\prime}|}\chi_{\nu}(r)\chi_{\kappa}(r ^{\prime})drdr^{\prime}$ $\displaystyle=C_{\mu i}C_{\nu i}C_{\lambda j}C_{\kappa j}J_{\mu\nu\lambda \kappa}=P_{\mu\nu}P_{\lambda\kappa}J_{\mu\nu\lambda\kappa}.$ (27) $$

Finally, the exchange-correlation energy is a functional of the electron density, which is unknown and has to be approximated. This is the main problem we are addressing in this work.

To find the ground state, the parameters $C$ are optimized to minimize Equation 1. This is typically achieved iteratively in a self-consistent fashion; see Appendix B.
The construction via Gaussian-type basis functions (GTOs) greatly simplifies the integrals appearing in the calculation allowing us to analytically precompute the so-called core Hamiltonian ${H_{\text{core}}}_{\mu\nu}=T_{\mu\nu}+V_{\mu\nu}$ and the electron-electron repulsion tensor $J_{\mu\nu\lambda\kappa}$ once and reuse them throughout the optimization.

## Appendix B Self-consistent field method

In this work, we use the Self-Consistent Field (SCF) method (Kohn & Sham, ; Lehtola et al., ) to find the coefficients $\hat{C}$ that minimize the energy functional $E[\rho(r;C)]$ in Eq. 1

$$ $\displaystyle\hat{C}=\arg\min_{C}E=\arg\min_{C}P_{\mu\nu}{H_{\text{core}}}_{ \mu\nu}+P_{\mu\nu}P_{\lambda\kappa}J_{\mu\nu\lambda\kappa}+E_{\text{XC}}$ (28) $\displaystyle\text{s.t. }\forall_{ij}\langle\phi_{i}|\phi_{j}\rangle=\delta_{ ij}.$ (29) $$

We can express the orthogonality constraint above via the overlap matrix $S\in\mathbb{R}^{{N_{\text{bas}}}\times{N_{\text{bas}}}}$

$$ $\displaystyle\forall_{ij}\langle\phi_{i}|\phi_{j}\rangle=\delta_{ij}$ $\displaystyle\iff C^{T}SC=I,$ (30) $\displaystyle S_{\mu\nu}$ $\displaystyle\coloneqq\langle\chi_{\mu}|\chi_{\nu}\rangle.$ (31) $$

This constraint is enforced by introducing the Lagrange multipliers $E_{ij}$, and we define the loss

$$ $\displaystyle\mathcal{L}=E+(C^{T}SC-I)_{ij}E_{ij}.$ (32) $$

Since the matrix $E$ is symmetric and the rotation of the basis does not change the energy, we can choose $E$ to be diagonal.
Using the Fock matrix $F\coloneqq\frac{\partial E}{\partial P}$, we can write the gradient of the loss w.r.t. $C$ as

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial C}=2FC-2SCE$ $\displaystyle\mathop{=}\limits^{!}0$ (33) $\displaystyle\implies FC$ $\displaystyle=SCE.$ (34) $$

This is the generalized eigenvalue problem we need to solve to find the coefficients $C$.
We solve this by diagonalizing $S=V\Lambda V^{T}$, substituting $C=X\tilde{C}$ with $X=V\Lambda^{-1/2}V^{T}$, and multiplying both sides from the left with $X^{T}$ to obtain the ordinary eigenvalue problem

$$ $\displaystyle\begin{split}\underbrace{X^{T}FX}_{\tilde{F}}\tilde{C}&= \underbrace{X^{T}SX}_{I}\tilde{C}E\\ \implies\tilde{F}&=\tilde{C}E.\end{split}$ (35) $$

After solving this for $\tilde{C}$, we recover the original coefficients $C=X^{T}\tilde{C}$ via the eigenvectors associated with the ${N_{\text{el}}}$ lowest eigenvalues $E$.
Since $F$ depends on $P=CC^{T}$, this problem is linearized by alternating optimization steps of $F(C)$, $C$, s.t. we need to iterate until convergence commonly referred to as self-consistence in this context.

## Appendix C Self-consistent field implementation

We implement the SCF method from Appendix B fully differentiably in JAX (Bradbury et al., ) by following Lehtola et al. ().
In particular, we precompute $\bm{\chi}(r)$ on the integration grid points (Treutler & Ahlrichs, ), and the core-Hamiltonian $H_{\mu\nu}$.
Since the electron-repulsion integrals $J_{\mu\nu\lambda\kappa}$ scale as $O({N_{\text{bas}}}^{4})$ in compute and memory, we use the density-fitting approximation (Vahtras et al., ) to reduce the memory and compute scaling to $O({N_{\text{bas}}}^{2}{N_{\text{aux}}})$ where ${N_{\text{aux}}}\ll{N_{\text{bas}}}^{2}$ is the size of the auxiliary basis set.
We explain the procedure in more detail in Appendix D.
Additionally, we improve the convergence by implementing the direct inversion of the iterative subspace (DIIS) method, which we briefly summarize in Appendix E (Pulay, ).
For precomputing the integrals and obtaining grid points, we use the PySCF (Sun et al., ).
For the evaluation of the atomic orbitals on the integration grid points, we use our own JAX implementation.

Like Dick & Fernandez-Serra (), we perform several SCF iterations with a pre-defined XC-functional to obtain a good initial guess for the density.
During training, we randomly interpolate between the precycled initial guess and a standard initial guess via

$$ $\displaystyle P_{0}=\frac{1+\alpha}{2}P_{\text{precycle}}+\left(1-\frac{1+ \alpha}{2}\right)P_{\text{standard}}$ (36) $$

where $\alpha\sim\mathcal{U}(0,1)$. At inference time, we fix $\alpha=1$. This interpolation leads to varying initial densities and functions as a regularizer (Dick & Fernandez-Serra, ).

## Appendix D Density fitting

Here, we largely follow Krisiloff et al. () and give a brief introduction to the density-fitting approximation.
For more details, we refer the reader to the original work.
We want to approximate the two-electron integral tensor

$$ $\displaystyle J_{\mu\nu\lambda\kappa}=\int\chi_{\mu}(r)\chi_{\nu}(r^{\prime}) \frac{1}{|r-r^{\prime}|}\chi_{\lambda}(r)\chi_{\kappa}(r^{\prime})drdr^{\prime}.$ (37) $$

To reduce the memory and compute scaling from $O({N_{\text{bas}}}^{4})$ to $O({N_{\text{bas}}}^{2}{N_{\text{aux}}})$, we introduce an auxiliary basis set $\{\chi_{\mu}^{\text{aux}}\}_{\mu=1}^{{N_{\text{aux}}}}$ and approximate the two-electron integrals as

$$ $\displaystyle J_{\mu\nu\lambda\kappa}\approx\hat{J}_{(\mu\nu)I}\tilde{J}^{-1}_ {IJ}\hat{J}_{(\lambda\kappa)J}$ (38) $$

where

$$ $\displaystyle\hat{J}_{(\mu\nu)I}$ $\displaystyle=\int\chi_{\mu}(r)\chi_{\nu}(r)\frac{1}{|r-r^{\prime}|}\chi_{I}^{ \text{aux}}(r)drdr^{\prime},$ (39) $\displaystyle\tilde{J}_{IJ}$ $\displaystyle=\int\chi_{I}^{\text{aux}}(r)\frac{1}{|r-r^{\prime}|}\chi_{J}^{ \text{aux}}(r^{\prime})drdr^{\prime},$ (40) $$

and the indices $I,J$ run over the auxiliary basis set $\{\chi_{I}^{\text{aux}}\}_{I=1}^{{N_{\text{aux}}}}$.
Note that choosing $\{\chi_{I}^{\text{aux}}\}_{I=1}^{{N_{\text{aux}}}}=\{\chi_{\mu}\cdot\chi_{\nu}
\}_{\mu,\nu=1}^{{N_{\text{bas}}}}$ yields the exact two-electron integrals.
One can further simplify by computing

$$ $\displaystyle B_{(\mu\nu)I}=\hat{J}_{(\mu\nu)J}\tilde{J}_{IJ}^{-\frac{1}{2}}$ (41) $$

where $\tilde{J}^{-\frac{1}{2}}=\tilde{J}^{-1}\tilde{J}^{\frac{1}{2}}$ with $\tilde{J}^{\frac{1}{2}}$ being the Choleksy decomposition of $\tilde{J}$.
This reduces the contraction of the two-electron integrals with the density matrix to

$$ $\displaystyle V_{\text{H}}[\rho]=P_{\mu\nu}J_{\mu\nu\lambda\kappa}P_{\lambda\kappa}$ $\displaystyle\approx P_{\mu\nu}B_{(\mu\nu)I}B_{(\lambda\kappa)I}P_{\lambda \kappa}.$ (42) $$

## Appendix E Direct inversion of iterative subspace

The direct inversion of the iterative subspace (DIIS) method is a common technique to accelerate the convergence of the SCF method.
Here, we briefly summarize the method; for a derivation and more details, we refer the reader to Pulay ().
The DIIS method aims at finding a linear combination of previous Fock matrices that minimizes the norm of the error matrix.
This accelerates the convergence of the SCF method by extrapolating a new Fock matrix from previous ones.
Let $F_{k}$ be the Fock matrix in the $k$-th iteration and $P_{k}$ the density matrix.
Then the error matrix is defined as

$$ $\displaystyle e_{k}=\left(S^{-\frac{1}{2}}\right)^{T}(F_{k}P_{k}S-SP_{k}F_{k}) S^{-\frac{1}{2}}$ (43) $$

where $S$ is the overlap matrix.
We aim to find coefficients $c_{i}$ that minimize the norm of a linear combination of error matrices

$$ $\displaystyle\min_{\{c_{i}\}}\left\|\sum_{i=1}^{m}c_{i}e_{i}\right\|^{2}$ (44) $$

subject to the constraint $\sum_{i=1}^{m}c_{i}=1$.
This minimization problem can be solved using Lagrange multipliers, leading to the linear system

$$ $\displaystyle\begin{pmatrix}{\bm{B}}&\mathbf{1}\\ \mathbf{1}^{T}&0\end{pmatrix}\begin{pmatrix}{\bm{c}}\\ \lambda\end{pmatrix}=\begin{pmatrix}\mathbf{0}\\ 1\end{pmatrix}$ (45) $$

where $B_{ij}=\text{Tr}(e_{i}^{T}e_{j})$, $1$ is a vector of ones, and $\lambda$ is a Lagrange multiplier.
Solving this system gives us the optimal coefficients $c_{i}$ which we use to obtain the extrapolated Fock matrix that is solved in the next SCF iteration.

## Appendix F Hyperparameters

**Table 3: Hyperparameters for EG-XC.**
| Hyperparameter | Value |
| --- | --- |
| $d$  number of features per irrep | 32 |
| ${l_{\text{max}}}$ number of irreps | 2 |
| $T$  number of layers | 3 |
| Radial filters | 32 |
| $\epsilon_{\text{mGGA}}$ Base semilocal functional | Dick & Fernandez-Serra () |
| Batch size | 1 |
| Number of steps to compute loss | 3 |
| Parameter EMA | 0.995 |
| Optimizer | Adam |
| $\beta_{1}$ | 0.9 |
| $\beta_{2}$ | 0.999 |
| Basis set | 6-31G(d) |
| Density fitting basis set | weigend |
| SCF iterations | 15 |
| Precycle XC functional | LDA |
| Precycle iterations | 15 |
| Learning rate |  |
| MD17 | $\frac{0.01}{1+\frac{1}{1000}}$ |
| 3BPA | $\frac{0.01}{1+\frac{1}{1000}}$ |
| QM9 | $\frac{0.001}{1+\frac{1}{1000}}$ |

We list the hyperparameters for EG-XC in Table 3.

## Appendix G Baseline changes

For reference methods, we used the model hyperparameters from their respective works.
For NequIP, we used the default $l=2$.
We extended patience schedules to ensure full convergence for all baselines.
On the QM9($S$) datasets, we first fitted a linear model on the training set with the number of atoms and shifted all labels by this prediction.
Further, we initialized all models with zero in the last layer.
These steps improved generalization between 10 and 100 times on the smaller datasets.

## Appendix H MD17 with 50 samples

We imitate the setting of Batatia et al. () and train all models on only 50 samples for each of the MD17 trajectories.
We reduce the validation set to 10 samples as well.
As we found KS-DFT methods to be more learning rate sensitive here, we ablated the initial learning rate with 0.01 and 0.001 and report the lower ones for Dick & Fernandez-Serra () and EG-XC.
The resulting test set MAEs are listed in Table 4.
KS-DFT methods and $\Delta$-ML are more successful learning from such few samples.
Especially on the challenging aspirin structures, EG-XC is the only method yielding accuracies within chemical accuracy.

**Table 4: MD17 with only 50 samples.**
|  | Force field | $\Delta$-ML | KS-DFT |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | SchNet | PaiNN | NequIP | SchNet | PaiNN | NequIP | Dick | EG-XC |
| Aspirin | 8.94 | 9.46 | 8.75 | 9.55 | 3.27 | 3.19 | 1.98 | 1.23 |
| Benzene | 2.11 | 2.28 | 2.84 | 0.25 | 0.26 | 0.21 | 0.53 | 0.59 |
| Ethanol | 4.15 | 3.30 | 4.44 | 2.99 | 1.60 | 1.32 | 0.99 | 0.35 |
| Malonaldehyde | 5.47 | 3.99 | 5.05 | 2.81 | 1.77 | 1.36 | 0.78 | 0.54 |
| Toluene | 6.65 | 4.40 | 5.15 | 1.15 | 0.79 | 0.80 | 0.70 | 0.99 |

## Appendix I Delta ML with SCAN (6-31G) DFT

In Table 5, we present additional MD17 data with the SCAN functional (Sun et al., ) and the 6-31G(d) basis set.
The results show that improving the reference DFT functional significantly improves $\Delta$-ML approaches.
It should be noted that the pure SCAN is already more accurate than the Dick & Fernandez-Serra () specifically fitted functional on Aspirin.
We argue that this setting is an ideal case of $\Delta$-ML as SCAN is well-suited for such close-to-equilibrium structures.
However, it should be mentioned that such improvements are orthogonal as one may also apply $\Delta$-learning on top of EG-XC.

**Table 5: MAE for $\Delta$-ML with SCAN and the 6-31G(d) basis set on MD17. For SCAN, we report the relative MAE instead to account for mean shifts.**
| Molecule | SchNet | PaiNN | NequIP | SCAN |
| --- | --- | --- | --- | --- |
| Aspirin | 0.57 | 0.46 | 0.34 | 1.84 |
| Benzene | 0.08 | 0.07 | 0.02 | 1.27 |
| Ethanol | 0.26 | 0.17 | 0.08 | 1.04 |
| Malonaldehyde | 0.32 | 0.25 | 0.13 | 1.26 |
| Toluene | 0.21 | 0.17 | 0.08 | 1.92 |

## Appendix J 3BPA mean absolute error

In Table 6, we present the absolute MAE on the 3BPA dataset for all methods.
It is apparent that while EG-XC performs well, $\Delta$-NequIP achieves the lowest MAE.
Thus, demonstrating that EG-XC’s error is largly a constant offset that would not affect actual MD simulations.

**Table 6: Absolute MAE in in $\text{\,}\mathrm{m\text{$E$}_{\mathrm{h}}}$ for structural extrapolation on the 3BPA dataset. All methods are trained on the 300K training set.**
|  | Force field | $\Delta$-ML | KS-DFT |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Test set | SchNet | PaiNN | NequIP | SchNet | PaiNN | NequIP | Dick | EG-XC |
| 300K | 5.15 | 2.91 | 3.85 | 2.38 | 1.14 | 0.81 | 0.96 | 0.42 |
| 600K | 28.14 | 13.85 | 24.56 | 4.09 | 2.13 | 1.56 | 2.19 | 1.43 |
| 1200K | 93.50 | 49.14 | 81.00 | 7.05 | 6.01 | 3.30 | 6.32 | 4.39 |
| $\beta=$120\text{\,}\mathrm{\SIUnitSymbolDegree}$$ | 19.98 | 4.16 | 13.77 | 2.54 | 1.26 | 1.09 | 1.58 | 1.04 |
| $\beta=$150\text{\,}\mathrm{\SIUnitSymbolDegree}$$ | 20.13 | 6.64 | 15.53 | 2.59 | 0.89 | 0.88 | 1.60 | 1.02 |
| $\beta=$180\text{\,}\mathrm{\SIUnitSymbolDegree}$$ | 20.25 | 9.01 | 17.35 | 1.93 | 1.15 | 0.77 | 1.70 | 1.04 |

## Appendix K 3BPA potential energy surfaces

We plot the potential energy surfaces for $120\text{\,}\mathrm{\SIUnitSymbolDegree}$, $150\text{\,}\mathrm{\SIUnitSymbolDegree}$, and $180\text{\,}\mathrm{\SIUnitSymbolDegree}$ in Figure 4, Figure 5, and Figure 6, respectively.
It is apparent that force fields struggle on all energy surfaces while XC-functionals yield a similar surface structure to the target.

Figure: Figure 4: Energy surfaces of the 3BPA dataset at $\beta=$120\text{\,}\mathrm{\SIUnitSymbolDegree}$$.
Refer to caption: x4.png

Figure: Figure 5: Energy surfaces of the 3BPA dataset at $\beta=$150\text{\,}\mathrm{\SIUnitSymbolDegree}$$.
Refer to caption: x5.png

Figure: Figure 6: Energy surfaces of the 3BPA dataset at $\beta=$180\text{\,}\mathrm{\SIUnitSymbolDegree}$$.
Refer to caption: x6.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - The results on MD17 (Table 1) and QM9 (Figure 3) seem to be good but with mixed results. How to understand these difference in accuracy?

- How is non-local defined? Will the re-weighting also make density gradients non local?
soundness: 3
strengths: - The proposed method is basis-independent.
- The method shows good results on structural extrapolation and data efficiency.
summary: This paper proposes to learn weighting coefficients for every point in 3D space to re-weight meta-GGA functional.
weaknesses: - Since the proposed method requires computing DFT, comparing it to ML force field might not be fair due to the higher computational cost.

- The proposed ML functional still rely on the existing functional, i.e.,  meta-GGA.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - Contribution Clarification: Could the authors explicitly clarify GG-XC’s unique methodological contributions? For readers less familiar with recent advances in this field, a succinct outline of GG-XC’s innovative aspects compared to prior methods would enhance clarity.

- Learning Objective and Loss Function: How is the GG-XC's loss function specifically defined? If CCSD(T) energy and density are indeed the targets, please confirm this assumption and clarify how functional derivatives are computed efficiently.

- Non-Local Reweighting Design Choice: What justifies the inclusion of the graph readout term when non-local terms are already introduced in the integral preceding it? Is there evidence that including the readout term substantively enhances performance?

- Differentiation Through SCF: Could the authors provide a more detailed explanation of the SCF process, particularly the derivation of the exchange-correlation functional’s differential form with respect to density? Furthermore, how is differentiation handled in the context of the graph readout term?

- Computational Efficiency Analysis: What is GG-XC’s relative computational cost compared to traditional mGGA and hybrid methods? An analysis of bottlenecks, particularly in terms of graph processing and SCF convergence, would be helpful for readers.
soundness: 2
strengths: - Originality: GG-XC offers a novel methodology by integrating GNNs with DFT, particularly beneficial for non-local XC functional design. This approach advances ML-based DFT by capturing long-range interactions through an equivariant point cloud representation of electron density.

- Quality: The paper presents rigorous empirical evaluation across datasets (e.g., MD17, QM9, 3BPA), clearly showcasing GG-XC's data efficiency and accuracy. The use of CCSD(T) energies as a baseline adds to the robustness of the study.

- Clarity: The structural breakdown of GG-XC into nuclei-centered embeddings, equivariant message passing, non-local reweighting, and graph readout is well-organized. However, clearer explanations regarding specific design choices, such as graph readout and the integration of the SCF process, would strengthen accessibility for readers outside the immediate field.

- Significance: By bridging ML force fields and DFT, GG-XC has potential implications for more accurate and scalable quantum mechanical computations. The work promises improvements in data efficiency and applicability to larger molecules, essential for fields like materials science and drug discovery.
summary: The paper introduces the Global Graph Exchange Correlation (GG-XC) functional, an innovative approach leveraging equivariant graph neural networks (GNNs) to model non-local exchange-correlation (XC) contributions in density functional theory (DFT). GG-XC addresses limitations in scalability, accuracy, and data dependency faced by previous ML-based and human-designed approximations. Through a combination of semi-local functionals and non-local density features, GG-XC captures long-range interactions and is trained using only energy targets by differentiating through a self-consistent field (SCF) solver. The empirical results demonstrate GG-XC's ability to achieve competitive accuracy with significantly reduced data requirements, even in out-of-distribution scenarios.
weaknesses: - Ambiguity in Methodological Details: Key details around the learning objectives and loss functions are unclear. For instance, it is inferred that GG-XC utilizes CCSD(T)-derived energy and density for training, but an explicit statement on this is missing. Additionally, the training schemes ˙such as backpropagating through functional derivatives) need to be explicitly written. Please clarify the optimization procedure and loss function explicitly.

- Design Choices in Non-Local Reweighting: The inclusion of the graph readout term and the reasoning behind the design choice to learn meta-GGA coefficients in GG-XC lacks sufficient justification. The authors suggest that the graph readout captures remaining non-local effects but do not clearly explain why the non-local terms in the preceding convolution is insufficient. 

- Insufficient SCF Explanation: The SCF procedure, particularly the calculation of the exchange-correlation functional's differential form with respect to density, is inadequately detailed. Especially, the calculation of Fock operator is absent in the manuscript. A deeper explanation of the differentiation through the graph readout term, specifically how it integrates into the SCF iterations, is necessary.

- Computational Cost and Efficiency: The paper omits a thorough computational cost analysis. An assessment of GG-XC’s runtime compared to standard mGGA or hybrid methods would clarify the efficiency benefits. Additional experiments comparing GG-XC with conventional mGGA and hybrid methods regarding speed and computational bottlenecks would significantly improve the work. This reviewer's recommendation are the mGGA functional $\epsilon_{\text{mGGA}}$ used in the GG-XC and $\omega B97X$ functional as in the 3BPA dataset.

- Comparative Baselines: While delta-machine learning methods are included in the comparison, the paper contrasts GG-XC trained on higher-level mGGA data with models trained on lower-level LDA for delta learning. This might obscure performance differences, as delta-learning methods trained on comparable baseline data (mGGA) could offer a more direct comparison. As mentioned in the previous point, please use the mGGA functional for the baseline models.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Given the accessibility of force calculations through differential techniques, why were these not included in the model’s experiments, could this learned energy functional be used for MD simulations or geometry optimizations?
2. Considering the fundamental role of electron density in DFT, could the model be evaluated for density accuracy and density-derived properties to ensure broader alignment with core DFT objectives?
3. Has the model been assessed for generalizability on intermolecule non-covalent interactions and diverse chemical reactions?
soundness: 3
strengths: 1. GG-XC significantly reduces errors in DFT calculations, achieving state-of-the-art performance on energy benchmarks.
2. The method is highly data-efficient, achieving excellent results with much less training data.
3. It demonstrates robust extrapolation capabilities, performing well even on out-of-distribution and larger molecular structures.
summary: The paper introduces a new model called Global Graph Exchange Correlation (GG-XC) to improve the accuracy of exchange-correlation (XC) functionals in density functional theory (DFT). Traditional methods for approximating XC functionals struggle with capturing long-range electron interactions, leading to errors in DFT calculations. GG-XC addresses this by using graph neural networks (GNNs) to create a nuclei-centered representation of the electron density that captures long-range interactions. The approach leverages machine learning (ML) to reduce computational costs and improve data efficiency while maintaining accuracy on various energy prediction benchmarks.
weaknesses: 1. In comparison, the force field paper prioritizes force accuracy, critical for molecular dynamics simulations, while the learned XC functional paper focuses on density accuracy, which is essential for practical DFT applications. However, the authors of the equivariant model rely only on energy metrics for evaluation, potentially leading to an incomplete comparison between the works. Despite the accessibility of force calculations via differential techniques, the equivariant model omits these, which are crucial for molecular dynamics. Additionally, density, a foundational component in DFT that governs electron distribution and influences properties such as dipole moments and band gaps, is not evaluated in the equivariant model.
2. The lack of tests on intermolecular non-covalent interactions and chemical reactions is a considerable gap. The differentiable programming approach demonstrates generalizability by successfully benchmarking against intermolecular non-covalent interaction datasets and different chemical reactions. Expanding the equivariant model's validation to include these would indicate broader reliability.
3. The lack of explicit enforcement of physical constraints may limit the accuracy in certain edge cases.
4. The model’s reliance on nuclei-based embeddings may make it less suitable for non-atomic systems.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The claim that 'the equivariant GNN efficiently captures long-range information' is somewhat misleading, as many equivariant graph neural networks (EGNNs) are designed with the assumption that inter-atomic interactions are localized. The underlying message-passing scheme typically aggregates information only from atoms within a cutoff radius, which raises questions about their ability to effectively model long-range interactions.

2. Finally, a comparison of the computational complexity of different methods is necessary.
soundness: 4
strengths: By combining semi-local functionals with a non-local feature density, GG-XC significantly improves the capture of long-range interactions. Its impressive performance in reconstructing 'gold-standard' CCSD(T) energies highlights its effectiveness in enhancing computational accuracy.
summary: The introduction of the Global Graph Exchange Correlation functional is innovative, utilizing equivariant graph neural networks to model non-local contributions to the exchange-correlation functional. This approach is relatively uncommon in traditional density functional theory, showcasing the potential of integrating machine learning with quantum chemistry.
weaknesses: The article offers a limited background overview and would benefit from a more thorough discussion of certain concepts. For instance, the section on 'non-local functionals' lacks detail and should clarify their significance in density functional theory. 

Furthermore, the claim that 'the equivariant GNN efficiently captures long-range information' is somewhat misleading, as many equivariant graph neural networks (EGNNs) are designed with the assumption that inter-atomic interactions are localized. The underlying message-passing scheme typically aggregates information only from atoms within a cutoff radius, which raises questions about their ability to effectively model long-range interactions.

Finally, a comparison of the computational complexity of different methods is necessary.

Meta-review: The submission proposes a new way of modeling and learning the exchange-correlation (XC) functional for more accurate Kohn-Sham density functional theory (DFT) calculations. The work models the electron density as atom-centered equivariant coefficients, which is processed by equivariant GNNs. The authors highlighted the effectiveness in handling long-range electron interactions, hence better accuracy and generalizability.

Reviewers have shown appreciation on the effective handling of non-local effects, the advantageous performance on benchmark datasets, efficiency over CCSD(T), and the generalization/extrapolation experiments to extend the usefulness. Reviewers also raised a few insufficiencies, including details about model design and hyperparameter choices, clear definition and ablation study of long-range effects, more evaluation metrics, and influence of lack of physical constraints as promoted in related work. In the rebuttal, the authors provided more details in the model development, detailed explanation of long-range effects and a related empirical study, and evidence that incorporating physical constraints is still debatable. These seem to have sufficiently addressed authors doubts, and two reviewers even strengthened their support.

I share the same appreciation as the reviewers on the contributions. Although both atom-centered equivariant coefficients representation of density and training a functional model with optimality enforcement are not completely novel, it is worth exploring a way that can be more suitable for a machine learning model to process, generalize, and capturing nonlocal effects. The authors also have addressed major concerns in the rebuttal. I hence recommend an accept.

Nevertheless, I still have a few points that the authors may consider for a stronger paper:
* This could be an inspiring work for machine-learning XC functionals, but at the current stage, it would be more impactful to show the promise to solve problems that conventional DFT cannot solve well, e.g., recovering macroscopic phenomenon that currently requires simulation in the CC level of theory. This may also related to the mentioned limitations, which could be reasonable for the first step, but need to be overcome for advancing the field (e.g., handling open-shell and charged systems with vast off-equilibrium structure generalizability for reactions).
* It could be a reasonable request to include more evaluation metrics as mentioned by Reviewer NL7q, considering that DFT is an electronic structure method that derives more than energy and force. The authors may also consider showing bond-breaking or angle-rotation curves to verify reasonable and useful generalization over the conformation space.
* Technically, why the radial part is not taken as the commonly chosen Gaussians, which could enjoy analytical calculation of the integrals hence getting rid of the expense and discretization error of using integral grids.
* It seems necessary to write out in the main paper how the inference stage is carried out; particularly what variable is being optimized? If it is the density integrals, then how to compute the non-interacting kinetic energy term, and how to guarantee the positiveness of the density function?

Decision: Accept (Spotlight)