Learning Equivariant Non-Local Electron Density Functionals
Authors: Nicholas Gao
Eike Eberhard
Stephan Günnemann
Venue: ICLR 2025 Spotlight (2025)
Keywords: Density Functional Theory
DFT
Functional
Exchange Correlation
XC
Equivariance
Graph Neural Network
Electron Density
Kohn-Sham DFT
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose an equivariant graph neural network-driven exchange correlation functional for Kohn-Sham DFT.

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 (GNNs). Where previous works relied on semi-local functionals or fixed-size descriptors of the density, we compress the electron density into an SO(3)-equivariant nuclei-centered point cloud for efficient non-local atomic-range interactions. By applying an equivariant GNN on this point cloud, we capture molecular-range interactions in a scalable and accurate manner. To train EG-XC, we differentiate through a self-consistent field solver 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% to 50%. 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% lower MAEs.

Published as a conference paper at ICLR 2025
LEARNING EQUIVARIANT NON-LOCAL ELECTRON DENSITY FUNCTIONALS
Nicholas Gao', Eike Eberhard", Stephan Günnemann (n. gao, e eberhard, s. quennemann ) @m de Department of Computer Science & Munich Data Science Institute Technical University of Munich
ABSTRACT
The accuracy of density functional theory hinges on the approximation of non- local contributions to the exchange-correlann correlaaion (XC) functional. To date, machine- learned and human- -denaned approximations suffer from insufficient accuracy, limited scalability, dependence 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 (GNNs). Where previous works relied on semi-local functionals or fixed-size descriptors of the density, we compress the electron density into an (Quii-qquivariant cloud for efficient non-local atomic-range interactions. By applying an equivariant GNN on this point cloud, we capture molecular range interactions in scalable and accurate manner. To train EG-XC, we differentiate through self-consistent field solver requiring only energy targets. In our empirical evaluation, we find EG-XC to accurately reconstruct 'gold-standard' CCSD(T) energies MD17. On out-of-distribution conformations of 3BPA, EG-XC reduces the relative MAE by 35 % 50 %. Remarkably, EG-XC excels in data efficiency and molecular size extrapolation on QM9, matching force fields trained on times more and larger molecules. On identical training sets, EG-XC yields on average 51 % lower MAEs.
muclei-cnntered point
INTRODUCTION
Kohn- -ham Density Functional Theory (KS-DFT) the backbone of computational material and drug discovery (Jones] 2015). It is a quantum mechanical method to approximate the ground state energy an Na-electron system by finding the electron density p E DN, {p: F3 Rgs p(r)dr Na) that minimizes the energy functional E DNa + E. This functional maps electron densities to energies and is composed of
Elp) Tpp] Veu(pp Vuilp) Exclp]
T DN. - R the kinetic energy functional, Vex : DN,, - IR the external potential due to positively charged nuclei, V/= DN,, + E the Coulomb energy between electrons, and, finally, Exc DN,, + R_ the exchange-correlation (XC) functional (Cramer) 2004). While T, Vou and VH permit analytical computations, the exact form of Exc remains unknown, and its approximation frequently dominates DFT's error (Kim et al. 2013).
Machine learning has emerged promising data-driven approach to approximate Exc (Kulik et al., 2022 Zhang et al. 2023). Many such ML functionals adopt the classical approach and define Exe as the integral of a learnable XC energy density cxc: Rudua R (Perdew 2001):
Excle] [ p(r)exc (g(r)) dr
where g _ + Ed are properties of the electron density P DDick & Sernandez-Serce 202 Nagai ct al. 2022 If g only depends on density quantities from its infinitesimal neighborhood, e.g., p(r), V pp(r).... ..., the resulting functionals are called semi-local. While this integrates well into exist- ing quantum chemistry code, non-local interactions like Van der Waals forces exceed the functional
Equull contribution; corresponding authors.

---

Published as a conference paper at ICLR 2025
2023 In contrast, non-local functionals can capture such interactions by depend- class (Kaplan et al.
ing on multiple points in space simultaneously, c.g.. Exc(p] = ʃوga ((")))(()) g(r") ((r")) drrdr'. Non-locality can be further separated into atomic-range interactions, which depend on the electron density around the nucleus, and molecular-range interactions, which depend on the electron density within typical molecular length scales. But, existing non-local functionals either scale poorly com- putationally (Zhou et al., (2019) or require costly reference data (Margraf & Reuter] [2021] [ystrom & Kozinsky 2022). The critical challenge in XC functionals lies in efficiently captureing non-local molecular-range interactions at scale.
To this end, we propose to leverage equivariant graph neural networks (GNNs) to learn non-local XC functionals through our Equivariant Graph Exchange Correlation (EG-XCI To enable computation- ally efficient non-local interactions, we propose two key innovations: (1) We compress the electron density to an density with an 50(3)-equivarian kernel at the nuclear positions. (2) We use -Ouiieqquvaariat GNNs on this point cloud representation to efficiently capture molecular-range information. Finally, we use the embeddings to define a non-local feature density 9NL : IR3 + E for the XC energy density exc from Equation [ 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 p.
atomic-range point cloud representation by convolving the electron
To train on energies alone, we differentiate through the minimization of Equation [
[Li et al.! (021).
In our experimental evaluation, EG-XC improves upon the learnable semi-local XC-functional on molecular dynamic trajectories, extrapolation both out-of-distribution conformations and increasingly larger molecules. In particular, we find EG-XC to reduce errors of the semi-local functional by factor of 2 to 3, on par with accurate ML force fields combined with DFT calculations (A-ML). In extrapolation to unseen structures, we find EG-XC to accurately reproduce out-of- distribution potential energy surfaces unlike force fields (incl. A-ML) while reducing the MAE by 35 % to 50% % 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 times less data. To summarize, we demonstrate that GNN-driven functionals push the frontier of non-local XC functionals and provide a promising path toward accurate and scalable DFT calculations.
2 BACKGROUND
Notation. To denote functionals, i.e., functions mapping from functions to scalars, we write F[f] where Fi is the functional and f the function. We use superscripts in (rackets(") to indicate sequences and regular superscripts of bold vectors, e.g., x', to indicate the 1-th irreducible representation of the SO(3) group. We generally use r E I to denote points in the 3D Euclidean space. A notable exception is the finite set of nuclei positions [RI,.... RNn.. } for which we will use capital letters. The norm of vector given by |||||| = ||r||2- For directions, i.e., unit length vectors, we use +-
Kohn-Sham density functional theory is the foundation of our work. Here, we provide very brief introduction. For more details, we refer the reader to Appendix or Lehtola et al. (2020). In KS-DFT, the electron density p is represented by a set of orthogonal orbitals di : R that are defined as linear combinations of a basis set of atomic orbitals X₁= R3 R:
+ R. (.(x) = C.² x(x).
p(r) = (())) (T)()(r))))))))))r)) x(r) x(r) = x(r)
(3)
where P= CC T.N the so-called density matrix. Given the representation in a finite basis set, one can compute the kinetic energy T, external potential Vau. and electron-electron repulsion energies V/H analytically. Unfortunately, such analytical expressions are generally unavailable for Exclp] as in Equation Thus, one relies on numerical integration (Lchtola et al. 2020
To minimize Equation I one typically uses the self-consistent field (SCF) method that we outline in greater detail in Appendix B The SCF method is iterative two-step optimization where one first
computes the so-called Fock matrix F = aP
DE based on the current coefficients C. Then, the optimal
coefficients C² E I8 X.a are obtained by solving the generalized eigenvalue problem
FC SCE
(4)
We provide the source code http:///gthhbbccmmeeschhhuuggg
2

---

Published as a conference paper at ICLR 2025
with Spu = Sas X. being the overlap matrix of the atomic orbitals and E being diagonal matrix. The procedure is repeated until convergence of P.
Equivariance allows defining symmetries of functions. A function f : X y is equivariant to group & iff Vg E G E Kf (G. x) = f(x) where G Gý are the representations of g in the domain X' and codomain }/ respectively. Invariance is special case of equivariance where Gg = 1, i.e., f (c x) fxx), for all g 9,x X.
The real spherical harmonics y R3 - TR N., are a well-known example of SO(3)-
equivariant functions they transform under rotation according to the Wigner 1) matrices D'2
(()++))((++11 VR (0(3), with Da being the identity and D being 3D-rotation matrices:
(' = D'rY'(())
(S)
Such 21- + 1-dimensional equivariant functions h4 : R + g2 c.g., h4 = Y can be recombined using the tensor product
(01 2 ha)an
Σ Σ
Cladaida h
mo.mi,m ma ma h/ng
11 TR m2
where C' e + are the Clebsch-Gorte coefficients. As the Wigner D-matrices orthogonal, it follows that the inner product of two l-equivariant functions yields invariant one
bf (Dir)) T h2 (Dhr) h((r)² DIT (r) N((r)"
While, we are interested in E(3))invarien functions, using such equivariant intermediate representa- tions improves accuracy, expands the function class, and improves data efficiency (Schütt et al. 20177
Gasteiger et al. 2021 Batatia et al. 2022).
3 RELATED WORK
Learnable density fuucctionl approximations (DFAs) of the XC functional have a long history. In their seminal paper. Kohn & Sham (1963) proposed an Exc fitted to a reference calculation. In this first parameterization exc only depends on (()) = p(r). To improve accuracy, the density features g have been successively refined by including gradients and physical quantities such as the kinetic energy density T(r) = Σ Na [r) The resulting semi-local functionals are also known as vetasGGAs with (r) IMGGG (r) = (r),, pV(()). [r)) (Perdew (001). Many MI. functionals use this parameterization with cxc frequently being product of an MLP with a classical functional (Nagai al.) [022] Kulik et al.) [022]] Zhang et al., (023). An appealing aspect of learning semi-local energy densities rather than Excle @ directly is the known implementation of physical constraints (Sun et al., 2015 Kaplan et al., 2023), which inspired works by Diik & Fernandez
Serra 2022);; Nagai et al. (2022). While augmenting ImGGA with hand-crafted non-local features improves the accuracy, this comes at the cost of exact constraints (Nagai et al. 2020 Kirkpatrick
al. 2021). Similarly, such physical biasses are absent in kernel methods Margraf & Reuter (2021 );
Bystrom & (202)), Additionally, these methods require reference densities that are rarely
available (Szabo & Ostlund 2012). Alternative grid-based CNNs converge slowly in the number of grid points compared to specialized spherical integration grids (Zhou et al.) 2019 Treutler & Ahlrichs 1995). Lastly, machine-learned functionals have been used in other DFT contexts as well,
e.g., for energy corrections on SCF-converged energies (Chen & Yang 2024) or to model the kinetic energy functional in orbital-free free DFT (Snyder et al.) 2012 et al. 2023 Zhang et al. 2024 Remme et al. 2023). Here, we propose to extend semi-local models with expressive fully learnable non-local features obtained from standard integration grids. This allows us 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 cheap but inaccurate alternative quantum mechanical calculations by directly approximating the potential energy surface from data without solving the electronic structure problem (Unke al. 2021). Contemporary force fields generally rely on graph neural networks (GNNs) (Schütt et al., 20117 representing molecules in terms of graphs with atoms as nodes. The advent of SO -qqii-aqt models led to significant improvements in accuracy
3

---

Published as a conference paper at ICLR 2025
(I) Nuclei-ceniered oquivariant
(2) Equivaziant message passing
(3) Non-local rewnighted mcta-GGAA
A
gNL
(xc
(0)
JR² dr
(4) Gimpph readout
N Nuuc
Exclpl
p(r) (((()r))
(TFi
|
CNL
Figure I: Illustration of Equivariant Graph Exchange Correlation (EG-XC)'s four components: (1)
We obtain a finite atomic-range point cloud representation ( by convolving the electron density P with radial filters I& R+ - R and spherical harmonics Y!. R + R at the nuclear position. (2) The embeddings (0) are updated using equivariant message passing to obtain molecular effects in H (3) We define a non-local feature density 9NL IR3 + R from which we derive the exchange correlation density exc IR + R. (4) We add graph readout of H to learn additional corrections. To obtain Exc(p], we integrate exc and add it to the global graph readout.
and data efficiency, closing the gap to full-fidelity quantum mechanical calcualtions (Batzner et al. 2022, Batatia al.) 2022). While they effectively aim to accomplish the same goal as DFT at fraction of the cost, their out-of-distribution accuracy is a common problem (Stocker et al. 2022). 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
With Eguivariant Graph Exchange Correlation (EG-XC), propose efficient non-local approxi- mation to the unknown exchange-correlatio functional Exc that maps electron densities (positive integrable functions in R33 scalar energies. To this end, EG-XC consists of four components
as illustrated in Figure i (1) Nuclei-centered equivariant embeddings (0) compress the electron density p into finite point cloud representation by equivariantly integrating the density around nuclei enabling atomic-range interactions. (2) Equivariant message pussing on this point cloud includes molecular-range information in HTT) Based on the point cloud embeddings, we define Exc the sum two readouts: (3) A non-local reweighted meta-Gol GGG based our non-local feature R Fd and (4) graph readout of the invariant nuclei-centered features (T)0.
density gNI.
Exclp] / p(r) YNI. (gNI. (r), ImGGI (r)) EmiGA (ImGGI (r)) dr
ENI. H (1))
(8)
local weights
meta-GGA
graph readout
reweighted GGA
(1) Nuclei-centered equivariant embeddings (()) E [R] x... 20 +1 Nux x
J-
reduce the computational scaling by mapping the continuous density to finite per-nucleus represen- tation. Further, these enable using equivariant GNNs to learn molecular-range interactions efficiently. The positions of the nuclei R, R' naturally lend themselves centroids for such embeddings as they represent peaks of the electron density (Kato 1957). Similar to Remme et al. (2023), we R 2(+1
perform this reduction by convoluting the electron density p with equivariant filters S% 1 at the nuclear positions RG:
H 10 (p. 5%) (Re)
(9)
4

---

Published as a conference paper at ICLR 2025
where Pi R + R4 is the soft- soft-partitioned electron density associated with the i-th embedding
(x;(r)
p:(r) = p((r)
(10)
2j=1 Now
og (r)
||r R ?
ax(r) = exp
(11)
X?
with À E R2 being free parameter. Such partitioning prevents overcounting if nuclei are close,
similar to quadrature methods (Becke 1988). Like equivariant GNNs (Batzner et al. 20022. we define
the filters S as product of spherical harmonics Y R + R2++1 and radial filters F R4 + R:
(r) = A (Irrl)) (f)
(12)
where 0 S / S 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 purely radial representation. Put together, compute the nuclei-centeres equivariant embeddings
(O) H
p.(r) R (r- Rs) dr R3
ik
Naua
(13)
- R (ry-Ri
with set of Nguad standard integration points and weights {(r), wy)) }} j=1 E [R3 x R.] Naw Treutlen
& Ahlrichs 1995). Importantly, while these embeddings centered the nuclei, they do not the nuclear charges, as ML force fields (Schütt et al., 2017). 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 's derivative would not exist and, thus, not alter the DFT calculation, effectively yielding force field.
We follow Schütt et al. (2021) and parametrize the radial filters F3 combination of sin, cos with a fixed polynomial envelope function u R4 R from Gasteiger al. (2022) for a cutoff e R4:
(x(r) sim
rark
(2)
if k even,
(14)
Fx (k+1)
u (#) if & odd.
We found such frequency filters yield more stable results than Bessel functions as the latter ones are very sensitive close to the centroids (Gasteiger et al. 2022). This sensitivity might be beneficial in force fields but is implicitly given in DFT due to the higher density close the nuclei.
(2) Equivariant message passing allows for the propagation and updating of the equivariant electron density features (0) to capture molecular-range dependencies within the electron density like Van- der-Wis Waals forces (Frank et al. 2022). The equivariance these features is important as it allows us feature density gNL in the next step. We use Batzner et al. 22022'ss 's NequlP to perform message passing with the SO(3)-equivariant convolution from Thomas et al. (2018) to iteratively update the embeddings
to define non-radial-symment
(()) over T steps:
H!Ch (((-)) EquiMLP (EquiLin
(H(")) + EquiConv (H rR), EquiLin (H())
(15)
where EquiLin Fr
Fd. is equivariant dense layer that mixes features of the same order
EquiLin(X)) )
Concat
Σ W
W
(16)
1=0
with each W' E [ddd being a learnable weight matrix. EquiMLP Fd fum, - Fd is the equivariant
equivalent of standard MLP, where we use the invariant embeddings | and an activation function IR + R to gate the equivariant parts > 0:
EquiMLP(X) EquiLin (ActEquiiiiii
(17)
ActEquiLin X Concat (o(w"X x"), [EquiLin (X)) o o (
(18)
5

---

Published as a conference paper at ICLR 2025
Lastly, define the convolution EquiConv F. xd
R Nux
F-Nuexd
via the tensor product:
Nmw
EquiConv(A (X, R
Σ
(LP RellModdd
==1 1
(19)
(@quiMLP(az.' 60 yl, R. Re
To accelerate the computation of the tensor product, we use Passaro & Zitmick (2023) 's efficient equivariant convolutions. For more details, we refer the reader to Batzner et al. (2022).
(3) Non-local reweighted meta-GGA. As the majority of the XC energy can be captured by semi- local functionals (Goerigk et al. 2017), starting with a machine-leanne meta-GGA CMGGA IR Redano - accounts for most of the XC energy. To correct the meta- meta-GGA for non-local effects, we first define
{ (H(D)
non-local feature density 9NL
R + Rd based on the non-local embeddings
(=0"
(NL(F)s = L
Now
as
T lus
Rg
T
H r(Ir
R&ll)² (h
(H!0) (90)
(20)
1=1
angular
radial
(8) 1 + Rd are MLPs mapping to radial weights. While the inner product between the where wi spherical harmonics Y and our equívariant embeddings H expresses angular changes, the inner product of radial basis functions I and wi 9NL and standard meta-GGA inputs ImGGA R Raaaos we derive the non-local correction ON.." RR duass >> IR to the exchange energy density EMGGA:
(t)1 allows for radial changes. From this feature density
(xc(r) = "NL (F), (mm,,xr)))())) CMGGA (gmoca(F))
(21)
In practice, we implement 'N. as an MLP. To obtain the final readout, we integrate the exchange- correlation density over the electron density Jgs exc(r)p(r)dr. Like Equation 13 we evaluate the integral with standard integration grids.
(4) Graph readout. In addition to the meta meta-GGA- GGA-basedd readout, we add global graph readout of the embeddings H to capture the remaining non-local effects. We use an MLP eNT. IRd + I
== on top of the invariant embeddings (1 0) to obtain the final exchange-correlation energy:
M
Exclpl 1/c (xc(r)p(r)dr
ENI. (H!00)
(22)
Linitations. While EG-XC demonstrates significant improvements 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 Vest (Kohn & Sham 1965). Second, the non- non-local nature of our functional permits no known way enforce most physical constraints (Ka plan et al. (023). Though, the importance of these constraints is still debated (Kirkpatrick al. 2021). see Appendix O These missing constraints enable the correction of basis set errors through XC functional. While this may lead to unphysical matches berween 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. 2023). Fourth, to handle open-shell systems, e.g., in chemical reactions, one would need to extend the equivariant embeddings to include spin information. Lastly, while more data efficient and better in extrapolation, running KS-DFT' is more expensive than a surrogate force field.
semi-local ML. functional,
5 EXPERIMENTS
In the following, we compare EG-XC to various alternative methods of learning potential energy surfaces across several settings. In particular, focus the following tasks: interpolating accurate

---

Published as a conference paper at ICLR 2025
Table 1: Test set MAE on the CCSD(T) MD17 dataset in mbh- (best, second)
Force field
A-ML PaiNN
KS-DFT
Molecule
SchNet
PaiNN
NequlP
SchNet
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
energy surfaces, extrapolation to unseen conformations, and extrapolation to larger molecules. Additionally, we present an ablation study on EG-XC's components. For runtime complexity analysis, we refer the reader to Appendix M and for runtime measurements to Appendix N
Methods. To accurately position EG-XC, we compare to methods of varying computational costs: force fields, A-ML, i.e., combining force fields with KS-DFT calculations, and learnable XC- functional (Dick & Fernandez-Serra (02I). While force fields are orders of magnitude cheaper by bypassing quantum mechanical calculations, they lack prior physical knowledge. For the A-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., 2021). As force fields, we test increasingly expressive models based on their use SO(3) irreducible representations: SchNet (1 (1=00 0) (Schürt et al.) 2017), PaiNN (1 1) (Schütt et al. (021), and NequlP (1 (0=2 2) (Batzner et al. 2022). Finally, we compare with EG-XC's learnable semi-local XC-functional (20TT. As A-M methods and learnable XC-functionals require DFT computation for each structure, they are at the same computational cost as EG-XC. In Appendix L we provide additional A-ML calculations with learnable XC functionals base method.
(Dick & Fernandez-Serm
Setup. To train the XC functionals, we follow Li et al. (2021) 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 per- formance; we outline the changes in Appendix G We list EG-XC's hyperparameters in Appendix
Loss. We follow Dick & Fernandez-Serral (2021) and minimize the mean squared error between the converged SCF energy and the target energy over the last Iines E N4 steps
L -=Σ 1=1-1m
Σ
(Enageet,n Ešcr( Run Zn)))²
(23)
with Enagat E IR beign the target energy and ESCF (R3 x N+) + R the SCF energy after the i-th iteration. The index n runs over the dataset. The parameter gradients are obtained through the total derivative of the loss function with respect to the parameters of the XC functional dC backpropagating through the SCF iterations. Importantly, for the parameters updates exist, the XC functional must be at least twice differentiable as we compute the partial derivative #Fxclo to a1
do by
construct the ith Fock matrix (see Appendix B and the total derivative for the updates dC 24 We obtain a Cme-smooth XC functionals by using the SILU activation function (Hendrycks & Gimpel 2023).
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 faster method. Hence, we compare these methods on the revised MD17 7 dataset, which contains precise 'gold-stane goold-stndarrd' CCSD(T) (CCSD for aspirin) reference energies for conformations of five molecules along the trajectory of a molecular dynamic (MD) simulation (Chmiela et al.) 2018 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. (2017), separate model is fitted per molecule. Since CCSD(T) calculations inherently account for non-loce 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 provide additional A-ML data with more accurate DFT functional and basis sets.
7

---

Published as a conference paper at ICLR 2025
Table 2: Structural extrapolation on the 3BPA dataset. All methods were trained on the 300K training set. All numbers are relative MAE in m The last three rows refer to the potential energy surfaces.
Force field
A-ML PaiNN
KS-DFT
Test set
SchNet
PaiNN
NequlP
SchNet
NequlP
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
3 120
3.84
1.78
2.25
2.53
1.25
1.09
0.75
0.35
150
4.64
1.89
2.64
2.03
0.84
0.88
0.61
0.23
3 180
4.97
1.92
3.03
1.79
1.06
0.73
0.56
0.20
NequlP
A-NequIP
Dick
EG-XC
Target
Figure 2: Two- tro-ddmensional slice of the potential energy surface of 3BPA with the dihedral 3 120% Pure force fields like NequlP struggle to recover the shape this When paired with DFT calculations, shape but introduces additional extrema. Learnable XC functionals like Dick Fernandez-Serra (2021) and EG-XC demonstrate significantly better reproduction of the target energy surface.
energy surface. that the energy surface moves closer to the target
T lists lists the mean absolute error (MAE) for all methods on the MD17 test sets. We find that force fields generally struggle reconstruct the energy surface of the MD17 dataset accurately when trained solely on energies. In contrast, A-ML methods and learnable XC-functionals generally achieve chemical accuracy at kcal mol a 1.6 Fi. The DFT reference calculations systematically reduce the le-be-learned energy surface delta from 6.2 F., to 4.9 m F Compared to the EG-XC's base semi-local Dick & Fernandez-Serral (202T), EG-XC reduces errors by factor of 2 to 4. Among
all methods, EG-XC yields the lowest error on 3 of the molecules. In Appendix H we test the performance on reduced training set of just 50 samples. There we find that gap between KS-DFT methods and force fields ones to widen significantly supporting the data efficiency of learnable KS-DFT methods. We hypothesize the good performance of force fields on MD17 being due to the homogeneous nature of the dataset where the gap between the training and test set is small. In such settings, physical inductive bias may be required perform well the test Conversely, expect the performance of force fields to degrade in extrapolation settings where the training and test set differ significantly. We investigate this hypothesis in the following experiments.
Extrapolating structures. In MD simulations, one often encounters structures f'rr outside the training set, e.g., due higher temperatures environmental changes (Stocker al. 2022). To investigate the extrapolation to unseen structures, we use the 3BPA dataset (Kovács et al. 2021). 3BPA contains various geometric configurations of the molecule 3-(benzyloxy - 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-o--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 B97X XC functional and the 6-31G(d) basis set and, thus, include non-local interactions.
The training set
As constant offsets the potential energy surface do not affect the system dynamics, we evaluate the relative MAE, i.e., E Epeed Euaget median Epre Eusga)]]. We list the relative MAE for all methods and test sets in Table 2 for the absolute MAE, we refer the reader to Appendix Across all test sets, EG-XC results in 35 % 51 % lower errors than the next best-tested alternative. On the
8

---

Published as conference paper at ICLR 2025
Heavy atoms in test.
6
8
6
9
8
high
SehNet
15.3
19.4
21.7
24.4
27.1
17.2
20.7
24.3
27.7
5.1
8.4
11.5
3.1
5.0
PaiNN
15.1
17.5
20.1
22.3
13.0
23.5
4.3
1.8
NequiP
11.5
15.7
19.0
22.3
25.3
5.2
7.0
9.2
11.2
2.6
4.3
5.8
1.4
2.4
MAE (mE,)
N-SchNeet
14.0
16.8
19.8
23.7
7.5
10.9
14.7
5.4
8.2
1.7
A A-PaiNN
10.1
12.8
15.9
19.6
23.5
3.1
4.1
5.4
7.0
1.3
2.1
2.7
1.0
1.8
A-NequlP
13.6
15.6
17.8
20.0
22.2
2.4
3.4
5.1
6.9
1.3
2.1
2.9
0.9
1.5
Dick
3.2
4.1
4.4
4.9
5.3
2.7
2.7
2.9
3.1
2.1
2.3
2.6
2.0
2.1
EG-XC
4.3
5.7
7.1
8.7
10.4
1.1
1.6
2.2
3.1
0.7
1.0
1.3
0.8
1.0
low
4 (48)
5 (174)
(776)
7
(3884)
Maximum number of heavy atoms in train (training set size)
Figure MAE in mÊ, on QM9 size extrapolation. Each row represents 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 heavy atoms listed above.
far OOD 1200K samples, EG-XC is the only method achieving chemical accuracy 1.6 E, relative MAE of 1. 1.40 mE To illustrate the qualitative improvement of EG-XC's energy surfaces, we plot the most OOD potential energy slice at 3 120" Figure for NequlP, A-NequIP, Dick & Fernandez-Serra (221 EG-XC and the target. The remaining methods and corresponding energy surfaces are plotted Appendix K 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 A-NequIP closer to the target surface, the energy shape includes additional extrema not present target surface. In contrast, both XC functionals accurately reproduce the energy surface. In line with the results on MD17, we find support to the hypothesis that learnable XC functionals are better suited for extrapolation to unseen structures than force fields.
Extrapolation to larger molecules. Gathering reference data for large compounds is costly and expensive with accurate quantum chemical calculations like CCSD(T) scaling O(Na )) in the number of electrons No. 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. 2014) into subsets of increasing size based the number of heavy atoms, i.c., QM9(S) all QM9 molecules with most heavy atoms. For each structures, i.e., QM9 QM)(S). For each training set, split the structures 90% 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-3 (221()) basis set and, thus, contain non-local interactions through h the exact Hartree exchange. As the dataset contains only few molecules with fluorine, force fields (including A-ML) could not yield accurate energies if none or few are in the training set. Thus, we omitted all molecules that include fluorine.
5, 6, 7]. we train a separate model and test on the remaining
We visualize the MAE for each combination the number of heavy atoms in the test set and the maximum number of heavy atoms in the training set in Figure Across all combinations of training and test sets, we find learnable XC functionals yielding the lowest errors with a preference for Dick & Fernandez-Serral (2021) 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 least 33 more accurate than the competing methods on test molecules with 9 heavy atoms. Compared to Dick & Fermandez-Serral (202T), EG-XC reduces the MAE by at least 2 QM9(6) and QM9(7).

---

Published as a conference paper at ICLR 2025
Table Ablation study on the 3BPA dataset.
300K
600K
1200K
8 120
120
B 120
no mGGA
7.00
12.89
25.85
10.99
11.16
10.82
no graph readout
0.97
1.38
2.30
0.77
0,63
0.57
no GNN
0.60
0.87
1.59
0.57
0.54
0.53
EG-XC
0.42
0.73
1.39
0.35
0.23
0.20
Overall, the QM9 results support the hypothesis that learnable XC functionals are better suited for extrapolation to larger molecules than force fields.
Ablations. To highlight the importance EG-XC's components, we perform an ablation study on the 3BPA dataset. We compare the full EG-XC model to variants without the mGGA, the graph readout, and the equivariant GNN. For the model without GNN, we use an equivariant MLP defined in Equation that that we apply rode-wise. We report the relative MAE on all test sets in Table One sees that the mGGA is crucial component for the performance of EG-XC. The model without graph readout approximately doubles EG-XC's error close to the performance of Dick & Fernandez-Sern (202T). Without the equivariant message passing, we observe that EG-XC can still reduce errors significantly, indicating that our equivariant convolution is well-suited encode the electron density. Otherwise, we would not observe an improvement here as a charge embedding could only correct a constant error, which is accounted for in the relative MAE metric.
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 basis- independent reduction of the electron density into a finite point cloud representation, enabling equivariant graph neural networks to capture molecular-range information. Based on this point cloud embedding, we have parameterized a non-local feature density used for the reweighing of a semi-local rechange-correlatin 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 % to 50 % 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 % lower errors 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 scchange-correlation function- als compared to machine-learned force fields. We hypothesize that a significant portion of EG-XC's generalization 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 range of future research. Thanks to the basis-set independence, EG-XC can transfer 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 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. 2024 Gao &
Günnemann, 2024a) and fast force field (Batzner al. 2022
10

---

Published as a conference paper at ICLR 2025
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.
REFERENCES
Ilyes Batatia, David P. Kovacs, Gregor Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. Advances in Neural Information Processing Systems, 35:1 1423 1436, December 2022.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Korbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. Eg)-eaauivarian graph neural networks for data-efficient and accurate interatomic potentials. Nature Communicutions, 13(1):2453, May 2022. doi: 10. 10038//41467-0222299999
ISSN 2041-1723
A. D. Becke. A multicenter numerical integration scheme for polyatomic molecules. The Journal of Chemical Physics, 88(4):2547-255) February 1988. ISSN 0021-9606. 1089-7690. doi:
1411063/14544033
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
Kyle Bystrom and Boris Kozinsky. CIDER: An Expressive, Nonlocal Feature Set for Machine Learn- ing Density Functionals with Exact Constraints. Journal of Chemical Theory and Computation, 18 (4):2180-2192 April 2022. ISSN 1549-9618. doi: 10.1021/acc.jett 1c00904.
Zehua Chen and Weitao Yang. Development of a machine learning finite-range nonlocal density functional. The Journal of Chemical Physics, 160(1):01405,, January 2024. ISSN 0021 doi:
10.1063/5.0179144
Lixue Cheng, P. Bernát Szabó, Zeno Schätzle, Derk Kooi, Jonas Köhler, Klaas J. H. Giesbertz, Frank Noé, Jan Hermann, Paola Gori-Giorgi, and Adam Foster. Highly Accurate Real-space Electron Densities with Neural Networks, September 2024.
Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learedd force fields. Nature Communications, 9(1): 3887, September 2018. ISSN 2041-1723. doi: 10. 10.10388/44447--18888884 8-06169--.......
Christopher J. Cramer. Essentials of Computational Chemistry: Theories and Models. Wiley, Chichester, West Sussex, England; Hoboken, NJ, 2nd ed edition, 2004. ISBN 978-0-470-09188 778-0-470-09181-4
Sebastian Dick and Marivi Fernandez-Serra Highly accurate and constrained density functional obtained with differentiable programming. Physical Review B, 104(16)L161109, October 2021. doi: 10. 103/PhysRevv.. 104.L161199.
Thorben Frank, Oliver Unke, and Klaus-Robed Müller. So3krates: Equivariant attention for interac- tions on arbitrary length-scales in molecular systems. Advances in Neural Information Processing Systems, 35:29400-2941 2022.
Nicholas Gao and Stephan Günnemann. Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions. In International Conference on Learning Representations, April 2022.
Nicholas Gao and Stephan Günnemann. Generalizing Neural Wave Functions. In International Conference on Machine Learning, February 2023a.
Nicholas Gao and Stephan Günnemann. Sampling- free Inference for Ab-Initio Potential Energy Surface Networks. In The Eleventh International Conference on Learning Representations, February 2023b.
11

---

Published as a conference paper at ICLR 2025
Nicholas Gao and Stephan Günnemann. Neural Pfaffians: Solving Many Many-Electron Schrödinger Equations. In The ehihty-eighth Anmual Conference on Neural Information Processing Systems, September 2024a.
Nicholas Gao and Stephan Günnemann. Representing Electronic Wave Functions with Sign Equivariant Neural Networks. In ICLR 2024 Workshop on A14DiferrntialEquatios In Science, March 2024b.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal Directional Graph Neural Networks for Molecules. In Advances in Neural Information Processing Systems, May 2021.
Johannes Gasteiger, Chandan Yeshwanth, and Stephan Günnemann. Directional Message Passing on Molecular Graphs via Synthetic Coordinates, April 2022.
Lars Goerigk, Andreas Hansen, Christoph Bauer, Stephan Ehrlich, Asim Najibi, and Stefan Grimme. A look at the density functional theory zoo with the advanced GMTKN55 database for general main group thermochemistry, kinetics and noncovalent interactions. Physical Chemistry Chemical Physics, 19(48):32184-32) 342115,, 2017. ISSN 1463-9076, 1463-9084. doi: 00.103//CCCP44913.
Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs), June 2023.
R. O. Jones. Density functional theory: Its origins, rise to prominence, and future. Reviews of Modern Physics, 87(3):897-923, August 2015, ISSN 0034-6861, 1539-0756. doi: 10.1 10.1100/RewModPhy 87.897.
Aaron D. Kaplan, Mel Levy, and John P. Perdew. The Predictive Power of Exact Constraints and Appropriate Norms in Density Functional Theory. Annual Review of Physical Chemistry, 74(1): 193-218, 2023, doi: 10. 46//nnnrev--hysccheee www.ppy--ccccccmm--0000000
Tosio Kato. On the eigenfunctions of mann--pariccn systems in quantum mechanics. Communications Pure and Applied Mathematics, 10(2):151-177, 1957. ISSN 1097-0312, doi: 10. 10.1002/cpa. 3160100201.
Min-Cheol Kim, Eunji Sim, and Kieron Burke. Understanding and Reducing Errors in Density Functional Calculations. Physical Review Letters, 11 111(7):073003, August 2013. ISSN 0031-9007, 1079-7114, 14. doi: 10.1 1..100/yyyseel.ett.. 1110033033.
James Kirkpatrick, Brendan McMorrow, David H. P. Turban, Alexander L. Gaunt, James S. Spencer, Alexander G. D. G. Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, Lara Román Castellanos, Stig Petersen, Alexander W. R. Nelson, Pushmeet Kohli, Paula Mori- Sánchez, Demis Hassabis, and Aron J. Cohen. Pushing the frontiers of density funetionals by solving the fractional electron problem. Science, 37466573) December 2021. doi: 10.1 (0.1116/scienceeaab655111
W. Kohn and L. J. Sham. Self-Consistent Equations Including Exchange and Correlation Effects. Physical Review, 140(4))A1113--1113 November 1965. ISSN 0031-899X. doi: 10.1103/ PhysRev. 140.A1133.
Arthur Kosmala, Johannes Gasteiger, Nicholas Gao, and Stephan Günnemann. Ewald-based Long- Range Message Passing for Molecular Graphs. Proceedings of the 40th International Conference on Machine Learning, pp. 17544-17563. PMLR, July 2023.
Dávid Péter Kováes, Cas van der Oord, Jiri Kucera, Alice E. A. Allen, Daniel Cole, Christoph Ortner, and Gábor Csányi. Linear Atomic Cluster Expansion Foree Fields for Organic Molecules: Beyond RMSE. Journal of Chemical Theory and Computation, 17(12):696-777 77111, December 2021. ISSN 1549-9618. doi: 10. 10.1022/acs httaa/cs.ccc...00e....
David B. Krisiloff, Caroline M. Krauter, Francis J. Ricci, and Emily A. Carter. Density Fitting and Cholesky Decomposition of the Two-Electron Integrals in Local Multireference Configuration Interaction Theory. Journal of Chemical Theory and Computation, 11(11):5242-522 November 2015. ISSN 1549-9618. doi: 10. 10.1021/aa http///cccctc.b000000
12

---

Published as a conference paper at ICLR 2025
H. J. Kulik, T. Hammerschmidt, J. Schmidt, S. Botti, M. A. L. Marques, M. Boley, M. Scheffler, M. Todorović, P. Rinke, C. Oses, A. Smolyanyuk, S. Curtarolo, A. Tkatchenko, A. P. Bartók,
S. Manzhos, M. Ihara, T. Carrington, Behler, O. Isayev, M. Veit, A. Grisafi, J. Nigam, M. Ceriotti, K. T. Schütt, J. Westermayr, M. Gastegger, R. J. Maurer, B. Kalita, K. Burke, R. Nagai, R. Akashi,
O. Sugino, J. Hermann, F. Noé, S. Pilati, C. Draxl, M. Kuban, S. Rigamonti, M. Scheidgen, M. Esters, D. Hicks, C. Toher, P. V. Balachandran, I. Tamblyn, S. Whitelam, C. Bellinger, and
L. M. Ghiringhelli. Roadmap on Machine learning in electronic structure. Electronic Structure, doi: 10. 10.8/88/2 1066-1055//52.5
(2):023004, August 2022. ISSN 2516-1075,
Susi Lehtola, Frank Blockhuys, and Christian Van Alsenoy. An Overview of Self-Consistent Field Calculations Within Finite Basis Sets. Molecules, 25(5):1218, January 2020. ISSN 1420-3049. doi: 1033390moleuulee22551222..
Li Li, Stephan Hoyer, Ryan Pederson, Ruoxi Sun, Ekin D. Cubuk, Patrick Riley, and Kieron Burke. John-Sham Equations as Regularizer: Building Prior Knowledge into Machine-Learned Physics. Physical Review Letters, 126(3):036401, January 2021. ISSN 0031-9007, 1079-7114. doi:
10.1 103/PhysRevl.eet. 126.036401.
Johannes T. Margraf and Karsten Reuter. Pure non-local machine-learned density functional theory for electron correlation. Nature Communications, 12(1):344, January 2021. ISSN 2041-1723. doi: 10.1038/s4 467-0200- -20-22010....
Wenhui Mi, Kai Luo, S. B. Trickey, and Michele Pavanello. Orbital-Free Density Functional Theory: An Attractive Electronic Structure Method for Large-Seale crst--Pinciipes Simulations. Chemical Reviews, 2c00758.
November 2023. ISSN 0009-2665, doi: lww.///cc.eemrrre.
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, May 2020. ISSN 2057-3960. doi: 10. 10.10388/4115440020000011
Ryo Nagai, Ryosuke Akashi, and Osamu Sugino. Machine-learinng- based exchange correlation functional with physical asymptotic constraints. Physical Review Research, 4(1):013106, February 2022. doi: 10. Res
Saro Passaro and C. Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In Proceedings of the 40th International Conference Machine Learning, volume 202 of ICML.223, pp. 27420-27438, Honolulu, Hawaii, USA, July 2023. JMLR.org.
John P. Perdew. Jacob's ladder of density functional approximations for the exchange-correlation energy. In AIP Conference Proceedings, volume 577, pp. 1-20, Antwerp (Belgium), 2001. AIP. doi: 10.1063/1.13900
P. Pulay. Improved SCF convergence acceleration. Journal of Computational Chemistry, 3(4): 556-560, 1982. ISSN 1096-987X. doi: 10.1002/jec.554003441..
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022, December 2014. ISSN 2052-4463. doi: (038/sdata.20144
Roman Remme, Tobias Kaczun, Maximilian Scheurer, Andreas Dreuw, and Fred A. Hamprecht. KineticNet: Deep learning transferable kinetic energy functional for orbital-free density functional theory. The Journal of Chemical Physics, 159() 2023.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In Proceedings of the 38th International Conference on Machine Learning, pp. 9377-9388. PMLR, July 2021.
Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Müller. SchNet: A continuous-file filter convolutional neural network for modeling quantum interactions. arXiv: arXiv:/706.00 08566 (physics, stat], December 2017.
13

---

Published as a conference paper at ICLR 2025
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, June 2012. ISSN 0031-9007, 1079-7114. doi: 10. 111100/PhyyReLLLtt.0082550000000
Sina Stocker, Johannes Gasteiger, Florian Becker, Stephan Günnemann, and Johannes T. Margraf. How Robust are Modern Graph Neural Network Potentials in Long and Hot Molecular Dynamics Simulations? Preprint, Chemistry, April 2022.
Jianwei Sun, Adrienn Ruzsinszky, and John P. Perdew. Strongly Constrained and Appropriately Normed Semilocal Density Functional. Physical Review Letters, 115())036402, July 2015. ISSN 0031-9007, 1079-7114, 14. doi: 10.1 N0.siee/hysseeveett. 15.036402.
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, Sebastian Wouters, and Garnet Kin-Lic Chan. PySCF: The Python-based simulations of chemistry framework. WIREs Computational Molecular Science, 1340, 2018. ISSN 1759-0884. doi: 10. 10.1002/wems.13
Attila Szabo and Neil S. Ostlund. Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory. Courier Corporation, 2012.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and mnnsllation-qqivarrin neural networks for 3D point clouds. arXiv:1802.08219 [cs], May 2018.
Oliver Treutler and Reinhart Ahlrichs. Efficient molecular numerical integration schemes. The Journal of Chemical Physics, (02(1):346-33) January 1995. ISSN 0021-9606. doi: 10.
Oliver T. Unke, Stefan Chmiela, Huziel E. Sauceda, Michael Gastegger, Igor Poltavsky, Kristof Schüit,, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine Learning Force Fields. Chemi- cal Reviews, 121(166)10142-10186 August 2021. ISSN 0009-2665. doi: 10. 10c.///cc.cemmrvv. 0c01111.
O. Vahtras, J. Almlöf, and M.W. Feyereisen. Integral approximations for LCAO-SCF calculations. Chemical Physics Letters, 213(5-6):514-513 October 1993. ISSN 00092614. doi: 10.1016/ 0009-2614493)89151-17
Simon Wengert, Gábor Csányi, Karsten Reuter, and Johannes T. Margraf. Data-efficient machine learning for molecular crystal structure prediction. Chemical Science, pp. 10. D0.103..D0S0076660 2021. ISSN 2041-6520, 2041-6539,
doi: 10. (0.1039/D0SC0576G.
Max Wilson, Nicholas Gao, Filip Wudarski, Eleanor Rieffel, and Norm M. Tubman. Simulations of state-of-the-art fermionic neural network wave functions with diffusion Monte Carlo, March 2021.
Max Wilson, Saverio Moroni, Markus Holzmann, Nicholas Gao, Filip Wudarski, Tejs Vegge, and Arghya Bhowmik. Neural network ansatz for periodic wave functions and the homogencous electron gas. Physical Review B. 107(2))2231139, June 2023. doi: 10. B...PhyySReBB.07.2331111.3
He Zhang, Siyuan Liu, Jiacheng You, Chang Liu, Shuxin Zheng, Ziheng Lu, Tong Wang, Nanning Zheng, and Bin Shao. Overcoming the barrier of orbital-free density functional theory for molecular systems using deep learning. Nature Computational Science, ((3):210-222 March 2024. ISSN 2662-8457. doi: 10. 0001444333888-2-000000
Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang. Haiyang Yu, YuQing Xie, Xiang Fu, Alex Strasser, Shenglong Xu, Yi Liu, Yuanqi Du Alexandra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung, Minkai Xu, Chaitanya K. Joshi, Simon V. Mathis, Kamyar Azizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitmik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Günnemann, Jure Leskovec. Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems, November 2023.
14

---

Published as a conference paper at ICLR 2025
Yi Zhou, Jiang Wu, Shuguang Chen, and GuanHua Chen. Toward the Exact Exchange Correlation Potential: A DineesDimesii Convolutional Neural Network Construct. The Journal of Physical Chemistry Letters, (022):7264-7269) November 2019. ISSN 1948-7185, 1948-7185. doi: 10.102//acs.ppcdeen.M223333
15

---

Published as a conference paper at ICLR 2025
A KOHN-SHAM DENSITY FUNCTIONAL THEORY
In Kohn-Sham density functional theory (KS-DFT) (Kohn Sham 1965). one attempts to approxi- mate the ground state energy of molecular system by finding the electron density that minimizes the energy functional Equation 1. Compared to ab-initio methods like Gao & Günnemann (2022) 2023b 2024b); Wilson et al. 2021 2023), KS-DFT requires the empiniall approximation of the XC functional Exc but scales more favorably with system size. Here, we largely compress Lehtola et al. (2020) to provide 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 functions, are denoted by acting upon the ket, e.g. VIJ) is the derivative of f. Further, we use the braket notation to denote the inner product of two functions f and as (flg) f f We will use Einstein's 's summation convention, where indices that only appear on one side of an equation are summed over.
The electron density p an Na system is associated with an antisymmetric wave function of fictitious non con-interacting system consisting of Na orthonormal orbitals Φ; R33 - R. These orbitals are constructed as linear combinations of Now. finite basis set Xy RR3 - R such that (s(r) = (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
(r) 1 det (os( (r,)),==1.
(24)
VN!
p(r) = / (r, (N) (T, IN )d)2
(25)
= (()) = x(r)
= x(r)
where C' e Mis XNa are the orbital coefficients and P= CT No. Nos 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 1965)
1 Tlpl - (dx (dx -p (Xp (x.l CusCuiT jwr Pau Twus
(26)
making KS-DFT significantly more accurate than contemporary orbital-free DFT.
The external potential in molecular systems is defined by the nuclei positions R and charges Z Vea[p] = - Ir R.
prr)Zn, dr = CCuiCus /
X (r(r) (T) dr = CuiCon Vau = Puu Vuew
(27)
r R.
The Hartree energy is a mean-mieeee field approximation to the dicctron-llcctrn interaction 1 Vale| 2 //
p(r)p(r')
r |--rl drdr's -1 C
--- xx Xu (-)))xx)))
= PPuPPaeJauke
(28)
Finally, the exchange-correlate and has to be approximated. This is the main problem we are addressing in this work.
energy is a functional of the electron density, which is unknown
To find the ground state, the parameters C are optimized to minimize Equation T 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 Howey Tyw + Vie and the electron- electron repulsion tensor JawAk once and reuse them throughout the optimization.
B SELF-CONSISTENT FIELD METHOD
In this work, we use the Self-Consisten Field (SCF) method (Kohn & Sham 1965 2020 to find the coefficients Ò that minimize the energy functional Elp(r; C)] in Eq.
Lehtola et al.
Ö Ō= arg min E arg min Pau + P Prnd + Exc
(29)
s.t. (. ōij-
(30)
16

---

Published as a conference paper at ICLR 2025
We express the orthogonality constraint above via the overlap matrix S E R RNa Nex
(o)())) = =bi = SCSC SC== =1
(31)
Suy (xulxe)
(32)
This constraint is enforced by introducing the Lagrange multipliers Edg and we define the loss
- (C SC Eij.
(33)
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 == ap. C' as
8E
we can write the gradient of the loss w.r.t.
2FC - 2SCE - Oc
ac
(34)
= =C=
(35)
where the Fock matrix is given by
0Exclp]
Fur Heurer + Pm +
(36)
0P...
To evaluate the Exc and its derivative we discretize R33 with spherical integration grids around
the nuclei (Becke, 1 1988). Specifically, for of quadrature points {ri) / RN and
2Naus E 1 Name we first compute the density and mGGA features ImGGA = weights (wi))mn
(VVr:): ((r)) these points:
p(ri) x(n:) px(r))
(37)
[V 211x(r:) (ex(r.))VXx(xx))
(38)
((r)) 1/2 (x(ri)"PPxx(n
(39)
From this discretized input, we compute the scalar output Exc. e.g., Exclp]
ENaw w.p((r))emGGA (gMGGA (r:)) for an mGGA, and rely on automatic differentiation via back- propagation to compute the derivative w.r.t. P. The generalized eigenvalue problem in Equation 35 solved to find the coefficients C. We solve this by diagonalizing
VAVT, substituting C' XC with X VA- VA-1/2VT, and multiplying both sides from the left with X to obtain the ordinary eigenvalue problem
F XT ČE
(40)
= F=CE..
After solving this for Č, we recover the original coefficients C X"O via the eigenvectors associated with the Na lowest eigenvalues E. Since F depends on P = CCT, 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.
C SELF-CONSISTENT FIELD IMPLEMENTATION
We implement the SCF method from Appendix B fully differentiably in JAX (Bradbury et al. 2018) by following Lehtola et (2020). In particular, we precompute ((r) on the integration grid points (Trutler & Ahlrichs, 1995). and the core-Hamiltonian Hpv- Since the electron-repulsion integrals JuwAN scale as O NNas in compute and memory, we use the density-fitting approxi- mation (Vahtras et al. 1993) reduce the memory and compute scaling to O 2Nax) where Naux Nons 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 briefly summarize in Appendix E (Pulay] 1982). For precomputing the integrals and obtaining grid points, we use the PySCF (Sun et al. 2018). For the evaluation of the atomic orbitals on the integration grid points,
JAX implementation.
17

---

Published as a conference paper at ICLR 2025
Like Dick & Fernandez-Serm 2021). perform several SCF iterations with pre-defined XC- functional to obtain a good initial guess for the density. During training, we randomly interpolate between the precycled initial guess and standard initial guess via
Po 2 precycle + (1 (1-1
+
1 1+
Panndard
(41)
2
where 4(0, 1). At inference time, we fix - 1. This interpolation leads to varying initial densities and functions as a regularizer (Dick Fernandez- Serra 2021
D DENSITY FITTING
Here, we largely follow Krisiloff et al. (2015) and give brief introduction to the density-fitting approximation. For more details, we refer reader the original work. We want to approximate the two-electron integral tensor
JurAn= = xx(()xx(r)) -r- Xa (rr)))c("")drrr
1
(42)
To reduce the memory and compute scaling from O(Naas to O(Nbas Naux), we introduce an auxiliary basis set (x] Nan and approximate the two-electron integrals as
JuNAN 20 J J
(43)
where
Jwwn) / xult)xo(r) r/) (r)drar',
(44)
J1 = XP (x) pr- (r')drrrr'''
(45)
and the indices I, J run over the auxiliary basis set {xItx) >mm- Note that choosing {xx"x fxu N yields the exact two-electron integrals. One can further simplify by computing
N..
J1=1
( J( where - with j3 being the Choleksy decomposition of J. This reduces the contraction of the two-electron integrals with the density matrix to
(46)
VHrØe PaudduAn Biert (iness PAK-
(47)
E DIRECT INVERSION OF ITERATIVE SUBSPACE
The direct inversion of the iterative subspace (DIIS) method a common technique to accelerate the convergence of the SCF method. Here, we briefly summarize the method; for derivation and more details, we refer the reader to Pulay/ 1982). 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 F3 be the Fock matrix in the k-th iteration and P3 the density matrix. Then the error matrix is defined as
e/ = (s- (=((s+t) (s-+) T
(FeP,S SPFFASS*
(48)
where is the overlap matrix. We aim to find coefficients e that minimize the norm of a linear combination of error matrices
min {c)
Σ G
(49)
subject to the constraint Σ c; multipliers, leading to the línear system
This minimization problem can be solved using Lagrange
B TT
(50)
where B, Tr(e, e,). 1 is vector of ones, and À is Lagrange multiplier. Solving this system gives us the optimal coefficients e; which we use to obtain the extrapolated Fock matrix that solved in the next SCF iteration.
18

---

Published as a conference paper at ICLR 2025
Table 4: Hyperparameters for EG-XC.
Hyperparameter
Value
number of features per irrep
32
Imax
number of irreps
2
T
number layers
Radial filters
32
<mGGA Base semilocal functional
Dick Fernandez-Serre (2021)
Batch size
Joss
Number of steps compute loss
Parameter EMA
0.995
Optimizer
Adam
81
0.9
82
0.999
Basis set
6-31G(d)
Density fitting basis set SCF iterations
weigend 15
Precycle XC functional
LDA 15
Precycle iterations
Learning rate
MD17
1+ 1000
3BPA
1+ 1000
QM9
1+
F HYPERPARAMETERS
We list the hyperparameters for EG-XC in Table 4
G BASELINE CHANGES
For reference methods, we used the model hyperparameters from their respective works. For NequlP, used the default / 2. We extended patience schedules to ensure full convergence for all baselines. On the QM9(S) datasets, we first fitted linear model on the training set with the number 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.
H MD17 WITH 50 SAMPLES
We imitate the setting of Batatia et al. (002 and train all models on only 50 samples for each of the MD17 trajectories. We reduce the validation set 10 samples well. As 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-Sens set MAEs are listed in Table KS-DFT methods and A-ML are more successful learning from such few samples. Especially the challenging aspirin structures, EG-XC is the only method yielding accuracies within chemical accuracy.
(02T) and EG-XC. The resulting test
I DELTA ML WITH SCAN (6-31G) DFT
In Table 6 we present additional MD17 data with the SCAN functional (Sun et al. 2013) and the 6-31G(d) basis set. The results show that improving the reference DFT functional significantly improves A-ML approaches. It should be noted that the pure SCAN is already more accurate than the Dick & Fernandez-eeral 2027) specifically fitted functional Aspirin. We argue that this setting is
19

---

Published as a conference paper at ICLR 2025
Table 5: MD17 with only 50 samples.
Force field
A-ML PaiNN
KS-DFT
SchNet
PaiNN
NequlP
SchNet
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
Table 6: MAE for A-ML with SCAN and the 6-31G(d) basis set on MD17. For SCAN, report the relative MAE instead to account for mean shifts.
Molecule
SchNet
PaiNN
NequlP
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
an ideal case of A-ML as SCAN is well-suited for such close-t--qquiliriun structures, However, it should mentioned that such improvements are orthogonal one may also apply A-learning on top of EG-XC.
J 3BPA MEAN ABSOLUTE ERROR
In Table 7 we present the absolute MAE on the 3BPA dataset for all methods. It apparent that while EG-XC performs well, A A-NequIP achieves the lowest MAE. Thus, demonstrating that EG-XC's error is largly constant offset that would not affect actual MD simulations.
K 3ВРА POTENTIAL ENERGY SURFACES
We plot the potential energy surfaces for 120%, 150", and 180 in Figure 4 respectively. It is apparent that force fields struggle all energy surfaces while XC-functionals yield a similar surface structure to the target.
Figure 5 and Figure 6
Table 7: Absolute MAE in mE for structural extrapolation on the 3BPA dataset. All methods are trained on the 300K training set.
Force field
A-ML PaiNN
KS-DFT
Test set
SchNet
PaiNN
NequlP
SchNet
NequlP
Dick
EG-X8 XC
300K 600K
5.15
2.91
3.85
2.38
1.14
0.81
0.96
0.42
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
120
19.98
4.16
13.77
2.54
1.26
1.09
1.58
1.04
8 150
20.13
6.64
15.53
2.59
0.89
0.88
1.60
1.02
180
20.25
9.01
17.35
1.93
1.15
0.77
1.70
1.04
20

---

Published as a conference paper at ICLR 2025
SchNet
PaiNN
NequlP
A-SchNet
A A-PaiNN
A-NequIP
E
Dick
EG-XC
Target
E
E
Figure 4: Energy surfaces of the 3BPA dataset 120"
Table 8: Relative MAE in mA, on the 3BPA dataset with A-MI. on top of learnable XC functionals.
Dick PaiNN
EG-XC PaiNN
Test set
SchNet
NequlP
SchNet
NequlP
300K
0.62
0.30
0.34
0.42
0.29
0.30
600K
1.01
0.56
0.62
0.71
0.52
0.56
1200K
1.85
1.24
1.38
1.38
1.15
1.26
3 120"
0.39
0.28
0.32
0.31
0.19
0.22
8 150"
0.41
0.35
0.27
0.24
0.19
0.22
8 180
0.45
0.37
0.25
0.20
0.15
0.21
21

---

Published as a conference paper at ICLR 2025
SchNet
PaiNN
NequIP
E
E
E
A-SchNet
A A-PaiNN
A-NequIP
E
Dick
EG-XC
Target
E
Figure 5: Energy surfaces of the 3BPA dataset at 3 = 150".
22

---

Published as a conference paper at ICLR 2025
SchNet
PaiNN
NequIP
E
E
E
A-SchNet
A-PaiNN
A-NequIP
Dick
EG-XC
Target
E
Figure 6: Energy surfaces of the 3BPA dataset at 3 = 180".
23

---

Published as a conference paper at ICLR 2025
L A-LEARNING ON LEARNABLE XC FUNCTIONALS
As one could see in the experiments of the main body, A-ML is a powerful technique to boost the accuracy of foree fields by integrating quantum mechanical calculations. One may improve this procedure further by relying on learnable XC functionals as baseline method instead of fixed functionals. Here, we investigate this approach by training A-ML on top of Dick & F'ernandezSerraa (2021) and EG-XC on the 3BPA dataset.
The results are presented in Table One can see that A-ML on top of such accuracy XC functionals improves accuracies significantly beyond the A learning with LDA functional in the main body. Across all force field models, we find EG-XC to yield lower errors on all test sets than the base mGGA functional from Dick & Fernandez-Serra (202T). This supports the hypothesis that the improvements learned by EG-XC are not identical with those A-ML.
M RUNTIME COMPLEXITY
In the following, we ignore the cost of precomputing the necessary integral tensors, the core Hamilto- nian Horre the density fitting tensor B (see. [.) the evaluation of the atomic orbitals on the quadrature grid, and the initial guess Po for the density as these do not contribute significantly to the overall runtime. A single EG-XC forward pass consists of the following steps:
1. O Naus )) Compute the electron density features JmGGA from the density matrix P. 2. ON (Nquad(laxx)) Compute the nuclear-centered point cloud embeddings H.
3. ONNic dd 3T) Of mux T): Equivariant message passing on the point cloud embeddings.
4. O max ): Compute the non-local feature density gNL on the nuclear grid points.
5. (N L): Compute the XC energy density exc from the non-local feature density.
6. ONNmed Graph readout to obtain the XC energy.
Note that due to the exponential decay in the partitioned density Pts we truncate the summations in Equation 13 and Equation 20 to avoid O quad Nmme) scaling. Additionally, single SCF requires the following steps:
1. ONNass 2 Naux): Compute the electron repulsion integrals.
2. ONNoss 2) Compute the Fock matrix.
3. O( 5) + O(T3): DIS extrapolation.
4. O(Nas 9): Solve the generalized eigenvalue problem.
5. O(Nas 3) Compute the density matrix.
For hybrid functionals like B3LYP wB97X, one has to add the cost of the exact exchange term with density fitting O(Nras" 3 N.ax).
For complete SCF calculation, we iterate these steps until convergence which typically requires & 15 iterations. In practice, computing the electron density features dominates the runtime for semi-local functionals. For hybrid functionals, the cost of the exact exchange term dominates the runtime.
N RUNTIME COMPARISON
Here, we compare the runtime of DFT calculation with different XC functionals, namely
LDA: a local XC functional,
SCAN: a mGGA XC functional,
B3LYP: hybrid XC functional,
24

---

Published as a conference paper at ICLR 2025
wB97X
2.5
B3LYP
EG-XC
2.0
Dick2021
Runline (s)
SCAN
1.5
LDA
1.0
0.0
2
10
14
18
22
26
30
34
36
HN
Figure 7: Runtime scaling of KS-DFT methods with different XC functionals.
Table Comparison between Dick & Femaandez-Sery 2021
and Nagai al. 2020 on 3BPA.
Test
Dick
Nagai
300K
0.96
0.91
600K
1.36
1.24
1200K
2.27
2.09
= 120
0.75
0.85
150
0.61
0.68
8 = 180
0.56
0.66
BB9X: range-separated hybrid XC functional,
Dick & Fernandez-Sery (2021): machine-leamed mGGA XC functional,
EG-XC
As test system, we use increasing cumulene chains C, with E G-d) basis set. We 15 SCF cycles for each DFT calculation and repeat the calculations 10 times from which we report the minimum to account for other processes running on the same machine. All calculations were performed on a single NVIDIA A100 GPU with our JAX implementation. For the hybrid functionals, we use density fitting to fit the exact exchange term into GPU memory.
38}. We use the
The runtime scaling for each method can be found in Figure While DFT calculations are approxi- mately 2 times slower than Dick & Fernandez-Sene (2021) on small structures, at the largest tested system size of Cas H4. 232 electrons, EG-XC only slower. Compared to hybrid functionals, we find B3LYP and wB97X being 43 % and 96 % slower than EG-XC, respectively. This highlights the efficiency of EG-XC's non-local interactions in large systems.
PHYSICAL CONSTRAINTS IN MGGA FUNCTIONALS
To investigate the importance of physical constraints on the mGGA, we train Dick & Fernandez. Serra (2021) and Nagai et al. (220) the 3BPA dataset. While Dick & Fernandez-eerra (2021)
implements many physical constraints, none of these are present in Nagai et al. (2020).
The relative MAE across all test sets is shown in Table 9 While Nagai al. (2000) yields lower errors on the MD trajectories at higher temperature, Dick & Fernandez Serra
(022 offers lower
errors the potential energy surfaces. These inconclusive results suggest that physical constraints do not necessarily improve the accuracy of XC functionals, even extrapolation tasks.
25


Reviews:
[{"review_id": "dCHgdfnaGX", "reviewer": "Reviewer_T8Ng", "summary": "This paper proposes to learn weighting coefficients for every point in 3D space to re-weight meta-GGA functional.", "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?\n\n- How is non-local defined? Will the re-weighting also make density gradients non local?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is basis-independent.\n- The method shows good results on structural extrapolation and data efficiency.", "weaknesses": "- Since the proposed method requires computing DFT, comparing it to ML force field might not be fair due to the higher computational cost.\n\n- The proposed ML functional still rely on the existing functional, i.e.,  meta-GGA.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SOzZZ85ljX", "reviewer": "Reviewer_6puv", "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.", "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.\n\n- 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.\n\n- 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?\n\n- 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?\n\n- 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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 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.\n\n- 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.\n\n- 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.\n\n- 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.", "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.\n\n- 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. \n\n- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dqAFT0LjWw", "reviewer": "Reviewer_NL7q", "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.", "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?\n2. 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?\n3. Has the model been assessed for generalizability on intermolecule non-covalent interactions and diverse chemical reactions?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. GG-XC significantly reduces errors in DFT calculations, achieving state-of-the-art performance on energy benchmarks.\n2. The method is highly data-efficient, achieving excellent results with much less training data.\n3. It demonstrates robust extrapolation capabilities, performing well even on out-of-distribution and larger molecular structures.", "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.\n2. 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.\n3. The lack of explicit enforcement of physical constraints may limit the accuracy in certain edge cases.\n4. The model’s reliance on nuclei-based embeddings may make it less suitable for non-atomic systems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qhMpcB8idD", "reviewer": "Reviewer_Xv1K", "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.", "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.\n\n2. Finally, a comparison of the computational complexity of different methods is necessary.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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.", "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. \n\nFurthermore, 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.\n\nFinally, a comparison of the computational complexity of different methods is necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)