A Function-Centric Graph Neural Network Approach for Predicting Electron Densities
Authors: Manuel V. Klockow
Marc K. Ickler
Peter Lippmann
Fred A. Hamprecht
Venue: ICLR 2026 Poster (2026)
Keywords: Graph Neural Network
Electron Density
Charge Density
Density Functional Theory
Message Passing
Basis Overlap
Equivariance
Molecules
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Electronic structure predictions are relevant for a wide range of applications, from drug discovery to materials science. Since the cost of purely quantum mechanical methods can be prohibitive, machine learning surrogates are used to predict the results of these calculations. This work introduces the Basis Overlap Architecture (BOA), an equivariant graph neural network architecture based on a novel message passing scheme that utilizes the overlap matrix of the basis functions used to represent the predicted ground state electron density. BOA is evaluated on QM9 and MD density datasets, surpassing the previous state of the art in predicting accurate electron densities. Excellent generalization to larger molecules of up to nearly 200 atoms is demonstrated using a model trained only on QM9 molecules of at most 9 heavy atoms.

Published as a conference paper at ICLR 2026
A FUNCTION-CENTRU GRAPH NEURAL NETWORK APPROACH FOR PREDICTING ELECTRON DENSITIES
Manuel V. Klockow, Mare K. Ickler, Peter Lippmann & Fred A. Hamprecht Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University 69120 Heidelberg, Germany (manuel. k <lockow, marc. ickler, peter. piter.lippman lippmann, fred. haaphhh)))c)) @iwr. uni mmi-heedeelbe www.eeeeee.......
ABSTRACT
Electronic structure predictions are relevant for a wide range applications, from drug discovery to materials science. Since the cost of purely quantum mechanical methods can be prohibitive, machine learning surrogates are used to predict the results of these calculations. This work introduces the Basis Overlap Architecture (BOA), an equivariant graph neural network architecture based on a novel mes- sage passing scheme that utilizes the overlap matrix of the basis functions used represent the predicted ground state electron density. BOA is evaluated on QM9 and MD density datasets, surpassing the previous state of the art in predicting accurate electron densities. Excellent generalization to larger molecules of up to nearly 200 atoms is demonstrated using a model trained only on QM9 molecules of at most 9 heavy atoms.
INTRODUCTION
Accurate electronic structure predictions are crucial for the development of new catalysts, improved batteries, or more specific drugs. Today's gold standard for reasonably sized systems is Kohn-Sham density functional theory (KS-DFT) It accounts for significant fraction of worldwide supercom- puting time, and three of its cornersoone methods are amongst the ten most cited publications of all time and fields (Van Noorden, 2025). Still, its computational cost prohibits routine use on large sys- tems, or in very high throughput scenarios. In response, machine learning surrogates are developed to either circumvent speed up KS-DFT calculations.
These methods range from property prediction (directly predicting observables from molecular ge- ometry (Batatia et al., 2022; Ko et al., 2023; Kozinsky et al., 2023; Simeon & Fabritis,, 2023; Batatia al., 2025; Liao et al., 2023; Anstine et al., 2025; Wood al., 2025)) to incorporating more physi- cal knowledge into the model. Examples of the latter include Hamiltonian prediction (Zhang al., 2024a; Yuan et al., 2024; Wang et al., 2024; Qian et al., 2025; Luo et al., 2025), mimicking the self-consistent iterations of KS-DFT (Song & Feng, 2024), or predicting the electron density using orbital-free schemes (Remme et al., 2023; Zhang et al., 2024b; Remme et al., 2025). Directly pre- dicting the electron density, as done here, lies between these extremes. The ground state electron density is an observable of central interest as it uniquely determines all ground state molecular prop- erties in theory (Hohenberg & Kohn, 1964), and many properties of interest can be derived from the density in practice. Another use for electron densities is to reduce the number of self-consistent field iterations required in KS-DFT (Koker et al., 2024: Sunshine al., 2023: Elsborg et al., 2025). Even single KS-DFT diagonalization step can be enough to reach chemically accurate energies relative to self-consistent KS solution (Jørgensen Bhowmik, 2022; Li al., 2025).
Most previous work on directly predicting ground state electron densities can be broadly grouped into two categories based on the representation used: The first class is based on a representation of the density in a (typically atom-centered) basis (Cuevas-Zuviría & Pacios, 2021: 1; Rackers et al., 2023; del Rio et al., 2023; Elsborg et al., 2025; Qiao et al., 2022; 2020: Fu et al., 2024; Cheng & Peng, 2023: Kim & Ahn, 2024; Mitnikov & Jacobson, 2024; Febrer et al., 2025). A model would then predict the coefficients of an expansion of the ground state electron density in the given basis. This approach has the advantage of being relatively scalable, as the density only needs to be evaluated

---

Published as a conference paper at ICLR 2026
on volumetric grid at the end of the model, if all. The choice of basis is however crucial, and to achieve high accuracy a very large number of basis functions is typically required. This problem was partially addressed by Fu al. (2024), where virtual nodes are employed, and the basis functions adapted a fine-tuning step to mitigate the impact of a suboptimal basis. Virtual nodes placed in the molecule, e.g. at the midpoint of bonds, and additional localized basis functions are centered at these nodes. Another approach to address this problem was taken by Elsborg et al. (2025), where floating basis functions are employed. These floating basis functions are no longer centered at the atoms of the molecule and instead the basis function positions are predicted per molecule individually, enabling much more flexible representation of the density.
The second class of methods works directly with a representation of the electron density on vol- umetric grid (Jørgensen & Bhowmik, 2022; Li et al., 2025; Koker et al., 2024; Gong et al., 2019; Sunshine et al., 2023). While this approach avoids basis set related inaccuracies, it is typically much more memory intensive, due the large number of grid points required to accurately represent the density.
Contributions This work falls into the first category of methods, i.c. we represent the density in an atom-centered basis and the model predicts coefficients in this basis. However, in contrast to previous work, we choose to represent the density in a quadratic expansion of the basis functions, inspired by the internal representation of the density in KS-DFT calculations using a density matrix. A product of two atom-centered Gaussian-type basis functions, as employed in this work, will be centered between the two atoms, avoiding the need for floating orbitals or virtual nodes. For a benzene molecule, this is shown in Figure IC. This work avoids explicitly predicting coefficients for each pair of basis functions, effectively employing a low-rank representation of blocks of the density matrix, without constructing the full coefficient matrix at any point.
To effectively utilize this representation, we introduce the Basis Overlap Architecture (BOA), a novel equivariant message passing neural network. The full BOA architecture is shown in Figure 1A. The fundamental idea of BOA is to imbue the model with information about the underlying basis by interpreting, appropriate points in the model, the internal features as functions represented in the given basis. Most notably the message passing formulation utilizes the overlap between functions represented in atom-centered basis functions different atoms to facilitate communication between nodes. During message passing each message is transformed from the basis of the sending node to the basis of the receiving node, essentially giving the best fit of the message in the basis of the receiving node. This instills the model with information not only about the underlying basis, but also the geometry of the molecule, as the overlap between basis functions centered at different atoms depends their relative position. This novel approach to message passing is shown to be very effective for the task predicting the electron density.
The internal features of BOA are separated into node and edge features, with the bulk of the compu- tation happening in the node features. This separation avoids the high computational costs of basing the full procedure solely on edge features, while still allowing for rich representation needed to accurately predict the density. The flow of information between node and edge features is shown in Figure IB.
BOA is evaluated on electron densities generated from the widely used QM9 dataset (Jørgensen & Bhowmik, 2022; Li et al., 2025; Ruddigkeit et al., 2012; Ramakrishnan et al., 2014) and molec- ular dynamics (MD) dataset of small organic molecules (Cheng & Peng, 2023; Bogojeski et al., 2020; Brockherde al., 2017). On all evaluated datasets BOA outperforms previous state-of- the-art methods by a significant margin. The full code to reproduce these results is available at /ttps::/githae con/ monnsccai-- lab/boa.
METHODS
To achieve the desired inductive bias, internal features are interpreted as functions represented in an atom-centered basis. At each point, the node features are given by hamps where E ' is the node with atom type Za, N is the set of all nodes, m E {1,..., No) } is the feature channel, and (1,..., N } is the basis index. We use Gaussian-type ("GTO") basis functions, specifically a fully uncontracted version of radial part is predicted by the model; see Appendix A for details. The internal features hamp
(Weigend & Ahlrichs, 2005). A correction to the
2

---

Published as a conference paper at ICLR 2026
A
B
Edge Reelookiig
Nodo Bactelding
Nayu
Nrelburty
Yony
Hiss
Fasteg
C
Partal Charzed Suc
JuM
CmTo,
KTO..
Rodbal
25.2
Aal (Pa)
Beooch Ala
A-Oxl
Mr
Chureel
Figure I: The Basis Overlap Architecture (BOA). (A) The node embeddings are updated using BOA blocks, which contain function message passing step to facilitate communication between nodes. The edge features are modified using the edge update block, which uses the current edge and node features to calculate new edge features. The output of the BOA backbone consists of coefficients used to expand the density in atom-centered Gaussian-type basis functions (Sec. 1). In the partial channel mean the number of channels is reduced by taking the mean of groups of channels. M denotes the molecular geometry and rg are the grid positions with g E {1...., No} and N9 being the number of grid points. (B) shows the flow of information in the BOA message passing block. While the edge features are updated using the node features, there is no flow of information from the edge features features with superscript (I) and (r) respectively, corresponding to the two adjacent nodes. (C) shows benzene molecule from the MD dataset with carbon atoms (black) and hydrogen atoms (white). Shown as green spheres are the centers of the Gaussian parts of the products of the atom- centered basis functions. The smaller def2-SVP (Weigend & Ahlrichs, 2005) basis is used here for better visualization. A product of | 0 basis functions from two different hydrogen atoms shown in blue, with its center (yellow) lying the middle between the two atoms on which the original basis functions are centered. Since the products of basis functions well distributed in space, a highly accurate representation of the density can be achieved without the need for floating basis functions virtual nodes.
the node features. Every edge in the graph has two edge
understood as expansion coefficients of No functions of space
fi (r)
ham (r) -->>
hample Ia),
aEA
aEN
consisting a sum over functions ham (r) localized at the node positions r.. Each of these functions in turn is represented in the atom type specific basis w.t (r). The same principle holds for the edge features Pabmp and Fabmas where b) E & is directed edge in the graph, with & being the set of all edges. Superscripts I and denote left and right, respectively, for reasons that are clear from Eq. 2. The edge features are interpreted as functions localized at node a, while the r edge features interpreted as functions localized at node b.
3

---

Published as a conference paper at ICLR 2026
To guarantee rotational equivariance the choice of basis functions is essential. Choosing basis fune- tions where the angular part is given by spherical harmonics ensures that the basis function coeffi- cients transform according to irreducible representations of SO(3), making it possible to construct equivariant operations. BOA is fully equivariant under rotations and translations by construction (see App. H).
2.1 DENSITY REPRESENTATION
Previous work explored different ways to represent the density, either in basis or directly on the grid. We choose a representation in an atom-centered basis. In contrast to previous work, we however do not expand the density (Cuevas-Zuviría & Pacios, 2021; Rackers et al., 2023; del Rio et al., 2023; Qiao et al., 2022: Fu al., 2024; Cheng & Peng, 2023) or its square root (Mitnikov & Jacobson, 2024) directly as linear expansion of the basis functions p(r) LnEN En Papin (r-
r.). Instead, we choose quadratic expansion, inspired by the natural expansion of the density in squared orbital functions in KS-DFT:
No
p(r) = Σ
La (r)g (T)+ aEN
E (o,b)E&
Σ
Jabo
(r)g
(2)
where No is the number of function pairs per edge and Jabo gabo(r) (r), Jobo (r), expanded in the localized basis functions w/t (r) and (r): (r):
10 (r) and (r)
(r) are
( (r) = Σ
Pabon wp Za rrra) (al)(r) =
Σ
Pabon ft
w faslit Ib),
(3)
1
#
0a (r) = Σ Illiuui La (r- ra),
( = Σ
Zo (r- ra).
(4)
This representation can be rewritten by defining I Tatuu = No (
Iobouiabon + HatDallaa Jav. with dab =
(r)
if b and 0 otherwise, resulting in p(r) = =E(a (a,eeee E Fabuwl Zo ra) (r- ro). Each I Tabur be interpreted as a block of the full density matrix T,,.. where the indices μ and understood to run over all basis functions of all atoms. Using the full density matrix, the density is represented by p(r) = Ene T y
(r (r), where (() is the concatenation of the basis functions of all atoms. This is exactly the representation of the density used in KS-DFT calculations. The full density matrix I f never explicitly constructed in BOA, instead the functions gabo (r) and (r) (r) evaluated on the grid and the density is obtained using Eq. 2, avoiding the costly evaluation of all pairwise products of basis functions on the grid. Equation amounts to evaluating a density given by low-rank representation of each block aabur of the density matrix.
(1) The expansion coefficients Pabop. and Fabor are predicted from the molecule geometry M by BOA
as described in the following sections. The self-loop coefficients gap and g&y are used to represent an initial guess for the density, which added to the density predicted by BOA. Taking the smooth absolute value of one of the functions in each pair proved beneficial (see App. B).
2.2 INITIAL GUESS OF THE ELECTRON DENSITY
BOA uses an initial guess of the ground state electron density based on the atom types of the nodes. After pre-training the initial guess, only an offset to that guess is learned by the model. The learned guess amounts to one pair of edge functions as described in Section 2.1 per atom. Consequently, (11) (()) the initial guess for the density is given by EeE fa (r)ge (r) and is part of the predicted electron density as shown in Eq. The initial guess coefficients are pre-trained for 1000 steps, but are not fixed after that, i.c. the initial guess is refined during the full training process.
2.3 NODE AND EDGE EMBEDDINGS
The node and edge embeddings are used to set the initial node and edge features respectively. Both lookup tables, which map the atom type Zu of the node and the atom types of the nodes
4

---

Published as a conference paper at ICLR 2026
kum
Plm
Chos
Cliss
Orah
ham
Onan Mha
desh
Figure 2: 1D illustration of function message passing. Two nodes a and are shown, each with two feature functions (blue and orange) centered at their respective positions. These functions expanded in local basis built from monomials multiplied with Gaussians. The first line shows the original feature functions ham and h6m- The second line shows the messages sent from node and node to node a. These are the projection of the respective feature functions to the basis of node a. The third line shows the new features ham after message passing which are a weighted sum of all incoming messages.
connected by the edge (a, b) to a learned set features. To keep equivariance under rotations, only the coefficients corresponding to / = 0 basis functions set in the embeddings, while the rest is initialized to zero. The node features are therefore initialized as
Wh(n) Z.ma u (r) is an I2,1 = basis function
hamp
(5)
0, otherwise,
where W-(n) Zamp learned during training and 1z.H is labeling the irreducible representation of the basis function w (r).
The edge features are initialized similarly as
W (e.)) Z.Zemp: if w (r) is an 1Z,1 basis function
11) 9
Jabmp
(6)
0, otherwise,
W (c.r)
7, mpmpp if wf (r) (r) is an 1z,0 = 0 basis function
Jabmp
0, otherwise,
where ((..))
and W Woler) Za Zemp are again learned.
Za Zuma
2.4 BASIS OVERLAP MESSAGE PASSING
The fundamental principle behind the message passing mechanism introduced in this section that each feature channel should be interpretable as function represented in given basis. To respect this interpretation, a basis change is employed to transform each message from the basis of the sending node to the basis of the receiving node, concept similar to Trame-to-frame transitions employed in local canonicalization schemes (Lippmann et al., 2025). Inspired by the attention mechanism (Vaswani al., 2017), the messages weighted by attention weights, which calculated from the overlap integrals between the features of the sending and receiving node. The full message passing mechanism shown in Figure 3B and one dimensional illustration is given in Figure 2.
BOA operates using two different cutoff radii Te, Tmp E R, with re being a smaller cutoff defining the edges & used for the edge features and the edge update module defined in 2.8, and Tmp being a larger cutoff defining the edges Emp used in the message passing step.
The first step in transforming the features of node b into the basis of node a calculating the overlap integrals between these features and the basis functions of node a. This overlap Oabmp is given by
Oabmu = - dr w (I ra) ha)hum (r) = dr w (r ra)) Σ w (r rb home
(8)
Σ
W Wabh hime:
(9)
5

---

Published as a conference paper at ICLR 2026
A
Wow
Selera
W0
B
MLLP
(W"))
MN
MN
MIN
MN
C
M. M.
Gansecar: TUB
husp
MLP
9ooma
Pahmu
Figure 3: Subblocks of the BOA architecture. An encircled index in the graph indicates a con- traction over the corresponding dimension, i.e. multiplication of the inputs and summation over the given index. (A) The new edge features 3(0) and g(r) Jabmy are generated as a superposition of the old
edge features Gabmy. and gabmp. and the node features hamp and humpy The weights of this superpo- sition are generated by passing the overlap integrals the edge and node features through an MLP. After the MLP the weights normalized using the matrix normalization (MN) block and scaled by the predicted factors aab The resulting weight matrices are then used to linearly mix the edge and node features. (B) Incoming messages calculated by transforming the feature functions from the basis of the sending node to the basis of the receiving node. Attention weights ōahmn are calculated from the overlap between the feature functions of the nodes and used weight the messages. (C) A correction to the radial part of the basis functions is learned for each atom type separately by small MLP. The radius is passed through a Gaussian radial embedding (Passaro & Zitnick, 2023) before being fed into the MLP. The output e, of the MLP is then used to scale the radial part of the basis functions by + Cp (see App. A for details).
(1)
where W = f dr w Fa ra)w> r&) is the overlap matrix between the basis functions of node and node b. From these overlap integrals, the representation of the features of node in the basis of node the basis functions of node a, resulting in the message mabmy from node b node
be obtained by multiplying with the inverse of the overlap matrix (Waa) of
ju
Mabmya
Σ (Waa) pv Oabmv.
(10)
These messages Mabma are the coefficients in the basis of node a that best represent the features of node 6 in the least squares sense. A derivation of this result can be found in Appendix G.
To weight the messages, attention matrix ōabmn is calculated for each edge (a, b) E Emp Each entry in this attention matrix describes the weight of the message from channel of node b to channel m of node a. The attention matrix calculated from the overlap between the feature functions of the two nodes
-f dr ham hr)
aabmn =
(r) = dr Σ Bammwi7 7-
ra) himow
Zu (r rb)
(11)
JAM
Dary Wab hanu
hampoatmg-
(12)
100
This feature overlap matrix Qabmm is processed by a multi-layer perceptron to produce the atten- tion weights Gabmns which determine the contribution of messages from neighboring nodes. These

---

Published as a conference paper at ICLR 2026
weighted messages are then aggregated to update the features of node
hamp
Σ
Σ
Öabma Magbmyy
(13)
6 st. (a.
Since the edges Emp include self-loops, the original features of each node are included in the message passing step, which amounts residual connection weighed by the attention weights Gaamn.
2.5 NONLINEARITY
Searching for suitable nonlinearity for features that should be interpretable as functions, one faces similar problems as for nonlinearities applied in equivariant neural networks. Here features are grouped together to form tensors with known transformation behavior under rotations, so that the nonlinearity can not be applied independently each feature. our case, the features are not only grouped to form tensors, but also form larger groups that each represent function.
One solution in the case of equivariant networks are gated nonlinearities (Weiler et al., 2018), where first scalar features are computed to which some nonlinear function is applied. 'he resulting scalar features are then used to scale the tensors, i.e. each group of features forming a tensor is scaled with the same scalar. In a similar approach we first calculate scalar features
lamn ff dr s dr' hamm (r) (1')(r)
Σ hamy haniene
(14)
where Chs f dr f drr (w ra ra (r' -ar--ry -r)) ( (1|r- r'll) is the Coulomb matrix, which can be generated for Gaussian-type basis functions using the PySCF package (Sun al., 2020). The resulting scalar features lamn are flattened and passed through an MIP. The result
(
Wa = MLPz, z. (l,)
(15)
reshaped to the original shape of lama and then used to linearly transform the features of node a
hamp Wamn
(16)
Z, denotes the atom type of node a, i.c. separate MLPs are learned for each atom type.
2.6 L2 NORMALIZATION
The normalization also aims to respect the function nature of the features, the L2 norm of each of the per-atom channel functions used for normalization. The norm calculated as
Dam = S dr (r)²
Bama Wag hamo and the features are normalized by hamy = hamp// (Nam + c) where is chosen as 10-6 to avoid numerical issues.
2.7 LINEAR LAYERS
Standard equivariant linear layers, as implemented in the e3nn package (Geiger & Smidt, 2022), are applied to the features of each channel and node. This means that only tensors of the same type are mixed and biases are only applied to scalars to preserve cquivariance. The linear layers used depend on the atom type of each node, i.c. separate parameters are learned for each atom type. After the application of these equivariant linear layers, the channels are mixed by a weight matrix, which again depends on the atom type of the node. This results in
2a Kamp
=Σ> W W W.Zg.lo4) hanv mn
(17)
where mm is the weight matrix for the mixing of the channels and W.Zg.1o0) is the weight matrix of the equivariant linear layer. While the linear layer used here is fully equivariant, it arguably does not fully respect the function interpretation of the features. Initial experiments with linear layers that respect the function interpretation, i.e. by applying an integration f dr Wmn (r) (r) where Wmn (r) learned functions represented in some basis, showed strong instabilities in training. Since the linear layers shown here worked well in practice, we decided to use them in the final architecture, opting for stability over additional inductive bias.

---

Published as a conference paper at ICLR 2026
2.8 EDGE UPDATE
The edge functions Jabm (r), -atem (r) are updated after each BOA block. The flow of information is illustrated in Figure IB. The node features ham (r) are used together with the old edge features
Jabm (r), Fabm (r) to generate new edge features Jabm (r), ((r) (r). There is however no flow of
information from the edge features to the node features. The full edge update procedure is shown in Figure 3A. Intermediate features are generated using the SO(3)-invariant overlap integrals of edge and node features
ff dr ham (1))lion
(n) abmn
=
Σ hamp Wab honys
(18)
(c) abmn
f dr Jabm (1)
(1) (t) 9abn (r)
Σ pe
9
Submy pubabb
(1) Wab (r)
(19)
These intermediate features are flattened and passed through an MLP to obtain weights w (n.)) w (n.r)
w (e.1) abman
abmnn w (e.)) w (e.r) (n.1) (m.r) (e.n) (e.r). The weight matrices w (n.I) (n,r) obman
abmmn s(rr)
ab w (e.r)
ab ab
are normalized by
wG..) ab
= (w (w())))
/ (I
+ )
(20)
where (+--) denotes either (n, l), (n, r), (e, or (e,r) and 7/v)) is a learned scalar factor. The normalization is computed in the Frobenius norm I1-11y- Using these weights the new edge features
( Pabma
Σ
(e.) (lee)) ab w
abmnnaabpp + Σ
ab w abman hamp
(n.) (n.))
(21)
generated, where for the / features and b for the r features.
3 EXPERIMENTS
The performance of BOA is evaluated on two electron density datasets based on QM9 (Jørgensen & Bhowmik, 2022: Li et al., 2025: Ruddigkeit et al., 2012: Ramakrishnan et al., 2014) and on a dataset based on MD trajectories (Cheng & Peng, 2023; Bogojeski et al., 2020; Brockherde et al., 2017). These datasets provide ground state electron densities for given geometries on a volumetric grid. Two versions of BOA are evaluated, small version and large version, differing only in the number of grid points used to evaluate the loss during training, and the batch size. The small version uses 5000 grid points and batch size of 12 while the large version uses 6000 grid points and a batch size of 24. Using these settings, the small version fits on a single A100 GPU with 440GB of memory while the large version requires an H100 GPU with ~94B of memory. All hyperparameters and training details are listed in Appendix I.
Additionally, the generalization capabilities of BOA are evaluated larger molecules with up to almost 200 atoms from the QMugs dataset (Isert et al., 2022) not seen during training. Ground-truth ground state electron densities are generated using the same method as used for the QM9 PySCF dataset (Li et al., 2025), the dataset on which the evaluated models were trained.
3.1 EVALUATION ON SMALL MOLECULES
As in previous works, the predicted electron density p(r) is compared to the refer- ence electron density p(r) using the normalized mean absolute error NMAE(P, p) (ʃ dr lp(r) p(r)l) / (f dr (e(r))). The integration is approximated on the full regular grid by summing over all grid points.
The two QM9 based datasets differ in the way the reference electron density was calculated, either using VASP (Jørgensen & Bhowmik, 2022) or PySCF (Li et al., 2025). BOA surpasses the previous state of the art on both QM9 datasets, shown in Table 1. The split between training, validation and test set is taken from Fu et al. (2024) for the QM9 VASP dataset and from Li et al. (2025) for the QM9 PYSCF dataset.
8

---

Published as a conference paper at ICLR 2026
Table 1: Comparison of BOA with previous best methods on the QM9 charge density datasets. Two datasets based on QM9 are evaluated, differing in the reference electron density calculation method. Errors are reported as NMAE 1%]. For BOA the mean and standard error over three runs are reported for the small models. For the large models the mean and standard error over five runs are reported. Errors of DFDeepDFT, InfGCN, ChargE3Net, and SCDP are reproduced from Fu al. (2024). The ResNet results are taken from Li et al. (2025).
NMAE [%]]
VASP ground truth
PySCF ground truth
DFF (Jörgensen & Bhowmik, 2022)
0.284
n/a n/a n/a n/a n/a 0.14
InfGCN (Cheng & Peng, 2023)
0.869
ChargENNet (Koker et al., 2024)
0.196
SCDP (Fu et al., 2024)
0.178
ELECTRA (Elsborg al.. 2025)
0.177
ResNet (Li et al., 2025)
n/a
BOA small
0.1381 = 0.0003
0.13 + 0.01
BOA large
0.1339 = 0.0005
0.116 + 0.006
Additionally, the performance of BOA is evaluated on the MD) dataset. As described in Cheng & Peng (2023), the MD dataset is curated from two sources. The ethanol, benzene, phenol and resorcinol data is taken from Bogojeski et al. (2020), while the ethane and malonaldchyde data is taken from Brockherde et al. (2017). BOA outperforms all previous methods on all molecules but one, where it matches the state of the art, shown in Table 2. Especially promising is the fact that no additional hyperparameter tuning was needed to achieve these results. The model is exactly the same in the QM9 experiments, with the only change being reduction in training steps to 200,000. Since only a training and test set is provided for the MD dataset, 10% of the training set is randomly sampled and used validation set to choose the best model during training.
Comparison of BOA with other methods the MD charge density dataset. Errors are reported as NMAE [%]. For BOA the mean and standard error of the mean over three runs are reported. Errors of the other models (InfGCN (Cheng & Peng, 2023), GPWNO (Kim & Ahn, 2024), SCDP (Fu et al., 2024), ELECTRA (Elsborg et al., 2025)) are reproduced from Elsborg et al. (2025).
Table 2:
NMAE [%]]
ethanol
benzene
phenol
resorcinol
ethane
malonaldehyde
InfGCN
8.43
5.11
5.51
5.95
7.01
10.34
GPWNO
4.00
2.45
2.68
2.73
3.67
5.32
SCDP
2.34
1.13
1.29
1.35
2.05
2.71
ELECTRA
1.02
0.45
0.56
0.62
0.91
0.80
BOA small
0.710 + 0.004
0.361
0.003
0.56
0.03
0.371
0.004
0.772 1: 0.002
0.61 1: 0.01
3.2 GENERALIZATION TO LARGER MOLECULES
To evaluate the generalization performance of BOA larger molecules, the QMugs dataset (Isert et al., 2022) is used. Ground-truth ground state electron densities for molecules with up to almost 200 atoms are calculated using PySCF (Sun al., 2020) with the same settings as used for the QM9 PySCF dataset (Li et al., 2025). Two BOA models trained on only the QM9 PySCF dataset are evaluated on this new test set. One of the models uses smaller message passing cutoff of rmp 3A instead of Tmp GÅ and smaller edge feature cutoff of r. = 2Á instead of r. 3Å. Additionally, we evaluate the ResNet model (Li et al., 2025) on this dataset for comparison. The time-scaling behavior and the NMAE all three models are shown in Figure 4. While the BOA model with the standard cutoffs performs worse than the ResNet model, the NMAE of the BOA model with smaller cutoffs stays roughly constant over the evaluated molecule sizes and outperforms the ResNet model. This shows that to achieve good generalization from small to large molecules, it is beneficial to limit the field of view of the model. This may be because the contents of a large field of view differ significantly between small and large molecules, introducing significant distribution shift, making

---

Published as a conference paper at ICLR 2026
it difficult for the model to generalize. Using smaller field of view enables BOA to generalize well to large molecules despite being trained on only small molecules of up to 9 heavy atoms.
The time scaling behavior shown in Figure 4A also shows second advantage of using smaller cutoffs. The BOA model with smaller cutoffs is significantly faster than both the ResNet model and the BOA model with larger cutoffs when evaluating larger molecules.
100
30A (standard cutaf)) RecNeet
B0A benall outodl)
1.8
B
BOA istundard cutaf))
BDA bmall cutof()
BecNiet
80
1.4
Time [4]
60
NMAE (96)
1.2
1.0
0.0
40
60
80
100
120
140
160
180
200
40
60
80
100
120
140
190
180
200
Number Atoms
Number Alomm
Figure 4: Generalization of BOA to larger molecules. Two BOA models trained on the QM9 PySCF data are evaluated larger molecules from the QMugs dataset (Isert al., 2022). One of the models uses a smaller message passing cutoff of Tmp = 3A instead of rmp = 6Å and smaller edge feature cutoff of re = 2Å instead of models compared to the ResNet model (Li et al., 2025). (B) shows the NMAE [%] of both models on the QMugs dataset over the number of atoms in the molecules. While the BOA model with the larger cutoffs performs worse than the ResNet model, the BOA model with smaller cutoffs outperforms
3,. (A) shows the time scaling behavior of both
the ResNet model with accuracy staying roughly constant over the evaluated molecule sizes.
4 DISCUSSION
Electronic structure calculations play fundamental role in computational chemistry, with myriad of practical applications. Accelcrating these calculations has the potential of enabling new applica- tions which are currently not possible, and of cutting costs for what is already feasible. BOA takes step towards this goal by introducing a novel architecture, permitting predictions of the electron den- sity with unprecedented accuracy. These results enabled by novel message passing mechanism treating the internal features as functions additional to a representation of the density in a quadratic basis expansion. The excellent generalization capabilities of BOA to larger molecules enables its application to problems where full electronic structure calculations are especially costly.
While BOA already achieves state-of-the-art accuracy on the evaluated datasets, there are several avenues for future work. Both the QM9 and the MD dataset contain only relatively small organic molecules. Training on larger more diverse datasets will be needed to enable applicability wider range of practical applications. Generalization over a large part of the periodic table would be desirable, which might require changes the architecture. Currently, BOA uses separate parameters for each atom type, which could become infeasible for more diverse datasets. A unified basis set for all atom types could be employed in the future to mitigate this problem.
An efficient representation of the density is crucial for dasis-set-based models like BOA. Another avenue for future work is therefore to improve the basis set used. BOA currently uses a fixed basis set of uncontracted Gaussian-type basis functions To achieve more flexibility. also the exponents of the Gaussian basis functions could be learned, done in (Fu al., 2024). In BOA, the overlap and Coulomb matrices used internally could be adapted during training, which could be achieved using differentiable quantum chemistry packages like PySCFAD (Zhang& & Chan, 2022).
base and leams radial correction factor.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This study has received funding from the Klaus Tschira Stiftung gGmbH (Simplaix project) and the Wildcard program from the Carl Zeiss Stiftung. Also this work is supported by Deutsche Forschungsgemmiinccaatt (DFG) under Germany's 's Excellence Strategy EXC-2181/1 390900948 (the Heidelberg STRUCTURES Excellence Cluster). The authors acknowledge support by the state of Baden- Wailln-Wtttttebbergg through bwHPC and the German Research Foundation (DFG) through grant INST 35/1597-1 FUGG.
REFERENCES
Dylan M. Anstine, Roman Zubatyuk, and Olexandr Isayev. AIMNet2: neural network potential to meet your neutral, charged, organic, and elemental-organt needs. Chemical Science, 16(23);
10228-10244, June 2025. ISSN 2041-6539. doi: 10. 10.10399/DSC008522H URL https:// pubs rsc. org/ ///c///et//ttiiicciicccoo wew.leeaannnn///22//222 /ttpg//000//////255555/5/3 The Royal Society of Chemistry.
Publisher:
llyes 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:11423-11466, Decem- ber 2022. URL. https: /prooceee///////eeeeee ngs. new.reiissnngnnnnuooeo...o...... hash/ /ttp////c...l1aa1111e100022a.... wtw.o//eeedebbbbbbbbbbbbbeebeeee cersstrracc-ooffeenneee
llyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kováes, Janosh Riebesell, Xavier R. Advincula, Mark Asta, Matthew Avaylon, William Baldwin, Fabian Berger, Noam Bernstein, Arghya Bhowmik, Filippo Bigi, Samuel M. Blau, Vlad Cărare, Michele Ceriotti, Sang- gyu Chong, James P. Darby, Sandip De, Flaviano Della Pia, Volker I. Deringer, Rokas Flijošius, Zakariya El-Machachi, Fabio Falcioni, Edvin Fako, Andrea C. Ferrari, John L. A. Gardner, Mikolaj J. Gawkowski, Annalena Genreith-Schrieee Janine George, Rhys E. A. Goodall, Jonas Grandel, Clare P. Grey, Petr Grigorev, Shuang Han, Will Handley, Hendrik H. Heenen, Kersti Hermansson, Christian Holm, Cheuk Hin Ho, Stephan Hofmann, Jad Jaafar, Konstantin S. Jakob, Hyunwook Jung, Venkat Kapil, Aaron D. Kaplan, Nima Karimitari, James R. Kermode, Panagio- tis Kourtis, Namu Kroupa, Jolla Kullgren, Matthew C. Kuner, Domantas Kuryla, Guoda Liepuo- niute, Chen Lin, Johannes Margraf, mran-Bogdan Magdäu, Angelos Michaelides, Moore, Aakash A. Naik, Samuel P. Niblett, Sam Walton Norwood, Niamh O'Neill, Christoph Ortner, Kristin A. Persson, Karsten Reuter, Andrew S. Rosen, Louise A. M. Rosset, Lars L. Schaaf, Christoph Schran, Benjamin X. Shi, Eric Sivonxay, Tamás K. Stenczel, Viktor Svahn, Christopher Sutton, Thomas D. Swinburne, Jules Tilly, Cas van der Oord, Santiago Vargas, Eszter Varga-Umbries Tejs Vegge, Martin Vondrák, Yangshuai Wang. William C. Witt, Thomas Wolf, Fabian Zills, and Gábor Csányi. A foundation model for atomistic materials chemistry, September 2025. URL http: / arxiv. org/ abs 24001440000 C0096. arXiv:2401.0m 1.0096 [physics].
Harry
Mihail Bogojeski, Leslie Vogt-Maranto, Mark F. Tuckerman, Klaus-Robert Müller, and Kieron Burke. Quantum chemical accuracy from density functional approximations via ma- chine learning. Nature Communications, 11(I):5223, October 2020. ISSN 2041-1723, doi: 14110038/441467-000---9993--- s41467-020-19093 Publisher: Nature Publishing Group.
URL https:/
nature.......... com/ articles/
Felix Brockherde, Leslie Vogt, Li Li, Mark E. Tuckerman, Kieron Burke, and Klaus-Robert Müller. Bypassing the Kohn-Sham equations with machine learning. Nature Communications, 8(1):872, October 2017. ISSN 2041-1723. doi: 117-00839 3. Publisher: Nature Publishing Group.
URL https: :/t::///wwww /www.
nature. wwt.///aaateeee/eeeeeee.eee
Chaoran Cheng and Jian Peng. Equivariant Neural Operator Learning with Graphon Convolu- tion. November 2023. URL https: ctttps://operrevieet wwwpr/veee.e.eeee/e/pe/....ewwwm www./wwweeeeeeee feeet/forumiidee d=E A3UWpncs noteTd==frwxxxtddg
Bruno Zuevis-Zuviría and Luis F. Pacios. Machine Learning of Analytical Electron Density in Large Molecules Through Message-Passing Journal of Chemical Information and Modeling, 61
(6):2658-266) June 2021. ISSN 1549-9596. doi: 10.102//acs, jeim. lc00227. URL https: Publisher: American Chemical Society.
//doi. org. 1021 /acs. jaim. cc00227.

---

Published as a conference paper at ICLR 2026
Beatriz G. del Rio, Brandon Phan, and Rampi Ramprasad. A deep learning framework to emulate density functional theory. upj Computational Materials, 9(1):158, August 2023. ISSN 2057- 3960. doi: 10. 1033///41524-023-011155555 URL https: /www nature www. wwooma/ricles//////ooc s41 02 01111-1011---1
Publisher: Nature Publishing Group.
Jonas Elsborg, Luca Thiede, Alán Aspuru-Guzik, Tejs Vegge, and Arghya Bhowmik. ELECTRA: Cartesian Network for 3D Charge Density Prediction with Floating Orbitals, May 2025. URL
http://arxis org. www./bsgg/b/b.oe.o.oooo......s.a /2503 08305. arXiv:2503.08305 [es].
Pol Febrer, Peter Bjørn Jørgensen, Miguel Pruneda, Alberto García, Pablo Ordejón, and Arghya Bhowmik. Graph2Mat: universal graph to matrix conversion for electron density prediction. Machine Learning: Science and Technology, 6(2):025013, April 2025, ISSN 2632-2153. doi: 10. /ddc871.
115522632-2133/ddeh565 URL https::tp//.... h/tp://xxxxooooo org Publisher: IOP Publishing.
Xiang Fu, Andrew Scott Rosen, Kyle Bystrom, Rui Wang, Albert Musaelian, Boris Kozinsky, Tess Smidt, and Tommi Jaakkola. A Recipe for Charge Density Prediction. November 2024. URL https: http://pppeeeee...eeeee net wttpee/wwww.ee/ooouiiieeeee
Mario Geiger and Tess Smidt. e3nn: Euclidean Neural Networks, July 2022. URL http://
arxiv. org/ abs 2207 09453. arXiv:2207.09453 [cs].
Sheng Gong, Tian Xie. Taishan Zhu, Shuo Wang, Eric R. Fadel, Yawei Li, and Jeffrey C. Grossman. Predicting charge density distribution of materials using local-enviromment- based graph convolutional network. Physical Review B, 100(18): 100(18):18103, November 2019. doi: 10. 1.1.1//yyssee...... 100.184103. URL https: / //////155 liik... a.p..app........ wwt.soorg/oo///// PhysRevB. 100. 184103. Publisher: American Physical Society.
Hohenberg and W. Kohn. Inhomogeneous Electron Gas. Physical Review, 136(3B):B864-877, November 1964. doi: 10.1103//hysevvv33.. 10. 10.1103// PPy03///hyeeevvcmm 136. 8864. Publisher: American Physical Society.
URL https: //link///nnkk aps..pp...o ors.org/oi//
Clemens Isert, Kenneth Atz, José Jiménez-Luna, and Gisbert Schneider. QMugs, quantum me- chanical properties of drug-like molecules. Scientific Data, 9(1):233, June 2022. ISSN 2052- 4463. doi: 10. 00000///1157-022000-111011 s41597- 020 01 390-390000 7. Publisher: Nature Publishing Group.
URL https: /www.... nature. com/ ammmrrriclss////
Peter Björn Jørgensen and Arghya Bhowmik. Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids. npj Computational Materials, 8(1):183, August 2022. ISSN 2057-3960.
URL https: // www. nature cww arti http:////rr..eeee/11144222222220 020////11222222088883333 Publisher: Nature Publishing Group.
doi: (0.10086441544402-000666666
Seongsu Kim and Sungsoo Ahn. Gaussian Plane-Wart Neural Operator for Electron Density Esti- mation. June 2024. URL
forum?
Tsz Wai Ko, Jonas A. Finkler, Stefan Goedecker, and Jörg Behler. Accurate Fourth-Generation Machine Learning Potentials by Electrostatic Embedding. Journal of Chemical Theory and Com- putation, 19(12):3567-3577, June 2023. ISSN 1549-9618, doi: 10. jctc. 2c01146. Publisher: American Chemical So-
146. URL
ciety.
Teddy Koker, Keegan Quigley, Eric Taw, Kevin Tibbetts, and Lin Li. Higher-order equivariant neural networks for charge density prediction in materials. npj Computational Materials, 10(1): 161, July 2024. ISSN 2057-3960. doi: 10. 00244//41554---2--133333333 nature. www.t//ettteeeee.oooooooooeeeeee
URL https: /etps://vae //wp:://www.
024 00014-1111111 1. Publisher: Nature Publishing Group.
Boris Kozinsky, Albert Musaelian, Anders Johansson, and Simon Batzner. Scaling the Lead- ing Accuracy of Deep Equivariant Models to Biomolecular Simulations of Realistic Size. In Proceedings of the International Conference for High Performance Computing. Networking, Storage and Analysis, SC '23, pp. 1-12, New York, NY, USA, November 2023. Association for Computing Machinery. ISBN 979-8-4007-0109-22 doi: 10. 00.114553881784.62770411 URL 3667041.
https: //dl. acm. .wt..rg//o...ooo.oo.
12

---

Published as a conference paper at ICLR 2026
Chenghan Li, Or Sharir, Shunyue Yuan, and Gamet Kin-Lic Chan. Image super-resolution in- spired electron density prediction. Nature Communications, 16(1):4811, May 2025. ISSN 2041- 1723. doi: 10. s414 4.67 67 - 60095 95- 8. Publisher: Nature Publishing Group.
URL https:/ /www natura.ecc..... wwt.aatiaaae///////.ttct...ceatt
Yi-Lun Liao, Brandon M. Wood, Abhishek Das, and Tess Smidt. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. October 2023, URL net wtt.rrmmmmmmcoocmmemaa...m
Peter Lippmann, Gerrit Gerhartz, Roman Remme, and Fred A Hamprecht. yond Canonicalization: How Tensorial Messages Improve Equivariant Message Pass- Peng, F. Sha, and R. (eds.), International Con-
Be-
ing. Y. Yue, A. Garg, ference on Representation Learning, volume 2025, pp. 88067-88087, 2025.
URL
wwwpoeeedings.iir.....oooo https: https::/ httpss///ppccceengg.s.. iclri www...c.....ppp...eeeeeeeeeeeeee htcpppppprrrllllpaaeeelreleeeeee wwwes//pppeeee/23/////1 ffiee//////////////////////////e htr553aalaccefffeecccceceeeeeee oot.aaacccccccaaaaaaaaaoooooooee pal..
Erpai Luo, Xinran Wei, Lin Huang, Yunyang Li, Han Yang. Zaishuo Xia, Zun Wang, Chang Liu, Bin Shao, and Jia Zhang. Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity. June 2025. URL https: http///ppnneeeweeeeeeeeeeeeeeeep net
htt.:rrrewwweeeeeeootmmmm.aa forum?id-
Ilan Mitnikov and Joseph Jacobson. E3STO: Orbital Inspired Squivvainl Molecular Repre- sentation for Electron Density Prediction, October 2024. URL http::/atxtxe org/abs/
2410. 0600661111 arXiv:2410.06119 [physics].
Saro Passaro and Lawrence Zimick. Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs. In Proceedings of the 40th International Conference on Machine Learning,
pp. 17420-274333 PMLR, July 2023. URL https: ppwoceeooeeeeooeoooooooocoooee.e ngs. saigs.ml..... mll wall.ppres//v000
passaro23a. html. ISSN: 2640-3498.
Chen Qian, Valdas Vitartas, James Kermode, and Reinhard J. Maurer. Equivariant Electronic Hamil- tonian Prediction with Many-Body Message Passing, August 2025. URL http:://rre
org/ htrg/abs/200 220...55.... 15108. arXiv:2508.15100 [cond-mat].
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller, III. OrbNet: Deep learning for quantum chemistry using symmetry-adapted atomic-orbital features. The Journal of Chemical Physics, 153(12): Sepeeeet doi:
September 2020. ISSN 0021-9606.
URL http::/: http::///ooooogg///./// 1063/5 0021955.
Zhuoran Qiao, Anders S. Christensen, Matthew Welborn, Frederick R. Manby, Anima Anandkumar, and Thomas Miller. Informing geometric deep learning with electronic interactions to acceler- quantum chemistry. Proceedings of the National Academy Sciences, 119631):e2002221119 2022. doi: 10. 00.1073/pas..2205222 1073 1073/pnas. 220
URL https: www. pnas http:////ccc//////////////////// epprin:: http:://www..naaaa..rraaap/tppth
Joshua A Rackers, Lucas Tecot, Mario Geiger, and Tess E Smidt. A recipe for cracking the quantum scaling limit with machine learned electron densities. Machine Learning: Science and Technol- ogy, 4(1):015027, February 2023. ISSN 2632-2153. doi: 10. 10.1088/2632-215 21532b33,,,, URL
htt.g///..00///333200 https:://de http:////.....o.ooooooo///////// 221122233222131332111-33 aab311
Publisher: IOP Publishing.
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quan- tum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022, URL htttt:: http://wt.........
August 2014. ISSN 2052-4463. doi: 10. 10.103/sdatt.2000 Publisher: Nature Publishing Group.
http//rtcceesssddaa
Roman Remme, Tobias Kaczun, Maximilian Scheurer, Andreas Dreuw, and Fred A. Hamprecht. KineticNet: Deep learning a transferable kinetic energy functional for orbital-free density func- tional theory. The Journal of Chemical Physics, 159(1)):14411,, October 2023. ISSN 0021- 9606, 7690, doi: 10.1063/5.0158275 URL http: arriv. org httg//bs/2300 1133111.1336 13316. arXiv:2305. 13316 [physics].
13

---

Published as a conference paper at ICLR 2026
Roman Remme, Tobias Kaczun, Tim Ebert, Christof A. Gehrig, Dominik Geng, Gerrit Gerhartz, Mare K. lckler, Manuel V. Klockow, Peter Lippmann, Johannes Schmidt, Simon Wagner, Andreas Dreuw, and Fred A. Hamprecht. Stable and Accurate Orbital-Free Density Func- tional Theory Powered by Machine Learning. Journal of the American Chemical Society, 147(32):28851-2888) August 2025. ISSN 0002-7863. doi: 2.1021/jacs.5c06211 URL https://dae http:////dd....11111111//11 jacs. 5c0 621 9. Publisher: American Chemical Society.
Lars Ruddigkeit, Ruud Deursen, Lorenz C. Blum, and Jean-Louis Reymond. Enumeration of 166 Billion Organic Small Molecules in the Chemical Universe Database GDB-17. Journal of Chemical Information and Modeling, 52(11):284--285, November 2012. ISSN 1549-9596. doi: 10. 00.1021/ci304415d. URL https: /tw.iiiiioo/// oogg org/ooogg//oooo 1tt.///111111111111 Publisher: Ameri- can Chemical Society.
Guillem Simeon and Gianni De Fabritiis. TensonNet: Cartesian Tensor Representations for Ef- ficient Learning of Molecular Potentials. November 2023. URL httpps:/
net feet/fozo forum? 1d=3EH1F
Feitong Song and Feng. NeuralSCF: Neural network self-consistent fields for density func- tional theory, June 2024. URL http://wwww..orr/a/////066/ 15873. Ttt1:2400.1887.. [physics].
Qiming Sun, Xing Zhang, Samragni Banerjee, Peng Bao, Mare Barbry, Nick S. Blunt, Nikolay A. Bogdanov, George H. Booth, Jia Chen, Zhi-Hao Cui, Janus Eriksen, Yang Gao, Sheng Guo, Jan Hermann, Matthew R. Hermes, Kevin Koh, Peter Koval, Susi Lehtola, Zhendong Li, Junzi Liu, Narbe Mardirossian, James McClain, Mario Motta, Bastien Mussard, Hung Q. Pham, Artem Pulkin, Wirawan Purwanto, Paul J. Robinson, Enrico Ronca, Elvira R. Sayfutyarova, Maximil- ian Scheurer, Henry Schurkus, James E. T. Smith, Chong Sun, Shi-Ning Sun, Shiv Upadhyay, Lucas K. Wagner, Xiao Wang, Alec White, James Daniel Whitfield, Mark J. Williamson, Sebas- tian Wouters, Jun Yang, Jason M. Yu, Tianyu Zhu, Timothy C. Berkelbach, Sandeep Sharma, Alexander Yu. Sokolov, and Garnet Kin-Lic Chan. Recent developments the PySCF program package. The Journal of Chemical Physics, 153(2):024109, July 2020. ISSN 0021-9606. doi: URL https: doi org 10. 1063/5.6..001 0006074.
Ethan M. Sunshine, Muhammed Shuaibi, Zachary W. Ulissi, and John R. Kitchin. Chemical Proper- ties from Graph Neural Network-Predicter Electron Densities. The Journal of Physical Chemistry December 2023. ISSN 1932-7447. doi: 10. (021/acs.jpcc.3e0666
C, (27(48):23459-2444 URL https: /tw.i wwtp:/.orgg11... 1021/ ipcc. 3c06157. Publisher: American Chemi- cal Society.
Richard Van Noorden. These are the most-cited research papers all time, Nature, 640(8059):591, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL attps::/paapee www.eerssnnpppppmmme.ee www.. cwtppcer_papppppeeeeeeeeee httes//pppe//21///aaeee $x5se//3B/dde2111b0ccccceel cbstaassss..cca.eaaaaaaaaaaa.... html..
Abstract
Zun Wang. Chang Liu, Nianlong Zou, He Zhang, Xinran Wei, Lin Huang, Lijun Wu, and Bin Shao. Infusing Self-Consistency into Density Functional Theory Hamiltonian Prediction via Deep Equilibrium Models, October 2024. URL http:// http://axivv... wwwp.x.og//as///0000000 03794. arXiv:2406.03794 [cs].
Florian Weigend and Reinhart Ahlrichs. Balanced basis sets of split valence, triple zeta valence and quadruple zeta valence quality for to Rn: Design and assessment of accuracy. Physical Chemistry Chemical Physics, 7(18):3297-3305 August 2005. ISSN 1463-9084, doi: 10.1039/ B508541 A. URL https: /pbbs.bba. rsc. org/orggee en content/////////oo aart particlell. alw.eelannnnn////00// cp/b508541e Publisher: The Royal Society of Chemistry.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Co- hen.
3D Steerable CNNs: Leamiing Rotationally Equivariant Features in Volumetric
14

---

Published as a conference paper at ICLR 2026
Data. In Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL htttps://ppoceedigg...oet neuripss wtt.app..ccapaaeee httpp/pappr//////aaaae .tmll.
Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso- Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John R. Kitchin, Daniel S. Levine, Kyle Michel, Anuroop Sriram, Taco Cohen, Abhishek Das, Ammar Rizvi, Sushree Jagriti Sahoo, Zachary W. Ulissi, and C. Lawrence Zitnick. UMA: A Family of Universal Models for Atoms, June 2025. URL http::t/tttiiiieeeeee// http://xxxvvv....aaa/ 2506. 23971. arXiv:2506.2397 [es].
Zilong Yuan, Zechen Tang, Honggeng Tao, Xiaoxun Gong, Zezhou Chen, Yuxiang Wang, He Li, Yang Li, Zhiming Xu, Minghui Sun, Boheng Zhao, Chong Wang, Wenhui Duan, and Yong Xu. Deep learning density functional theory Hamiltonian real space, July 2024. URL http:
//arxiv.org//////.. abs/ 2200 14379. arXiv:2407.1453 14379 [physics].
He Zhang, Chang Liu, Zun Wang, Xinran Wei, Siyuan Liu, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Self- Training for Censity-Functtonnal Theory Hamiltonian Prediction. June fmW.
2024a. URL https:
/tttpps://oeee ptattp:://opnneviieee
net htte///ooommmmmmmmeeeeee
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 molec- ular systems using deep learning. Nature Computational Science. 4(3):210-22) 10-22,,22222 March 2024b. ISSN 2662-8457. doi: 10. 10.1108//44388-0000
8. URL https: /www naturat com/ a articles/ s 02 00605-8 8. Publisher: Nature Publishing Group.
Xing Zhang and Garnet Kin-Lic Chan. Differentiable quantum chemistry with PySCF for molecules and materials the mean-field level and beyond. The Journal of Chemical Physics, 157(20): 204801, November 2022. ISSN 0021-9606. htg/10.10622 0118200..
doi: 10.1063/5.01 18200. URL https: / doi.
15

---

Published as a conference paper at ICLR 2026
A RADIAL CORRECTION FACTOR
While internally all features are represented in basis consisting of Gaussian-yye basis functions, for the final density representation an additional correction factor to the radial part of the basis func- tions is learned. For each atom type Z separately small MLP is used to predict correction factor ON based on the radius r = |||| . As shown in Figure 3C the radius is passed through a Gaussian radial embedding (Passaro & Zitnick, 2023) before being fed into the MLP. The corrected basis functions 2e (r) are then given by 2 (r) = w (r) (1 + (r(r), where the factor Cp is predicted per tensor, i.e. is the same for each basis function that corresponds to the same tensor keep equivariance. The impact of this radial correction factor is evaluated in Appendix C.I.
B SMOOTH ABSOLUTE VALUE ONE OF THE EDGE FUNCTIONS
Taking the smooth absolute value
12 if | < |xl
|x0,
(22)
2X otherwise
of Jabo (r) resulted in slightly improved performance compared to keeping the function unmodified,
(r)
see Appendix C.I. The factor steers the scale of the smooth absolute value function and is here
chosen as x 1000. The density is then given by p(r) Ei,,)E Iobo (r) 9abo (r) +
Eo No (1) (r)
EeEN gl. g (
C ABLATION STUDIES
C.I RADIAL CORRECTION AND SMOOTHI ABSOLUTE VALUE
ablation study of the radial correction factor and taking the smooth absolute value of one of the basis functions in each pair in the expansion of the density is performed. The results are shown in Table 3. For each configuration, three training runs are performed and the mean and standard reported. Both the radial correction and taking the absolute value significantly improve the performance, with the radial correction having the larger effect. The best performance is achieved when both techniques used together.
Table 3: Ablation study of taking the absolute value of one of the basis functions in the pair and applying the radial correction. The small BOA version is trained on the QM9 VASP dataset. Errors are reported as NMAE [%).
Absolute Value
Radial Correction
NMAE [%]
0.204 + 0.002
X
0.167 + 0.004
0.1423 + 0.0006
V
0.1381 + 0.0003
C.2 CHOICE OF BASIS SET
The impact the choice basis sets is studied on two additional smaller basis sets, de SVP and def2- TZVP (Weigend & Ahlrichs, 2005). All hyperparameters other than the used basis are chosen as in BOA small and the models are trained on the QM9 VASP data. The results are shown in Table 4. Smaller basis sets still perform reasonably well, with the medium size basis def2- TZVP still surpassing the previous state of the art. Still, using larger basis sets improves the NMAE significantly.
16

---

Published as a conference paper at ICLR 2026
Table Basis set ablation. All settings other than the basis set are kept constant to the small BOA version and the models are trained on the QM9 VASP dataset. Three training runs are performed for each configuration and the mean and standard error are reported.
NMAE [%]
Number basis functions
de SVP
0.194 + 0.001
103
def2 TZVP
0.1504 + 0.0003
192
def2 QZVPPD
0.1381 + 0.0003
374
C.3 QUADRATIC EXPANSION
To study the impact of the quadratic expansion train additional models using the same hy- perparameters as BOA small. When evaluating the density on the grid we drop ((7) (r) and Jabo (r) and expand the density linearly as p(r) E'EN Io (r) + E(a,)ee Eo No o (r) (1) 7
EnEN En gojluu Dajaa
ra) L(ab)EE, Ee Dabo W (r (r-rra) ra) where the sum over b can
be carried out before evaluating the basis functions w - ra) on the grid. Three models were trained on the QM9 VASP data using this linear expansion and evaluated on the test set. The re- sults are shown in Table 5. The quadratic expansion significantly outperforms the linear expansion, showing the importance the quadratic expansion for accurate density predictions BOA.
Table 5: Quadratic expansion ablation. All settings other than the expansion type are kept constant the small BOA version and the models are trained on the QM9 VASP dataset. Three training runs are performed for each configuration and the mean and standard error are reported.
NMAE [%]
Linear Expansion Quadratic Expansion
0.2716 + 0.0007
0.1381 + 0.0003
C.4 RADIAL CUTOFF
We conducted an additional experiment with larger cutoffs, specifically message passing cutoff of Tmp = SÁ (instead of Tmp = 6Â) and an edge feature cutoff re = 4Á (instead of r. = 3λ) were chosen. Other than the cutoffs the BOA Small settings the QM9 VASP data. Mean and standard error on the test set are reported in Table 6. Using larger cutoffs results in small but consistent improvement in performance.
used and three models are trained on
Table Radial cutoff ablation. All settings other than the cutoffs are kept constant to the small BOA version and the models are trained on the QM9 VASP dataset. Three training runs are per- formed for each configuration and the mean and standard error are reported.
NMAE [%]
Standard cutoffs
0.1381 + 0.0003
Larger cutoffs
0.1343 + 0.0007
D DFT INITIALIZATION AND COULOMB ENERGY ERROR
To evaluate the quality of the predicted electron densities metries closer to practical applications, we calculate the Coulomb energy eтror and investigate the savings in KS-DFT calculations when us- ing the predicted densities as initialization. Both of these evaluations are performed the full QM9
17

---

Published as a conference paper at ICLR 2026
PySCF test set of 10000 molecules using BOA large model. BOA achieves mean absolute error of 66 meV between the Coulomb energy calculated on the predicted densities and the converged label densities, comparing favorably to the 167 meV reported for the ResNet model (Li al., 2025).
KS-DFT calculations using our densities as initial guesses and using a standard superposition of atomic densities as initial guesses performed. The same settings in the data generation de- scribed by et al. (2025) are used. With these settings the BOA initialization reduces the mean number of SCF iterations needed from 15.7 to 10.2, a reduction by 35%. This a result competitive with the 35% reduction reported by Li et al. (2025) on the same data.
E EFFICIENCY
We evaluate the efficiency of ELECTRA Elsborg et al. (2025), SCDP Fu al. (2024), standard BOA, and BOA with def2-TZVP basis models. We note that there is no difference in inference efficiency between the small and large settings, since only the number of probe points and the batch size during training are changed, and we therefore do not need to differentiate between them here. All models evaluated on 40GB A100 GPU using block size maximizing the VRAM usage. The time per molecule is evaluated over the whole QM9 Vasp test set and the mean and standard deviation over these molecules are reported in Table 7. While BOA is roughly 2 times slower than SCDP, using the smaller de f2-2222/ mitigates this effect to some extent and achieves a smaller error than the previous state the art with efficiency comparable SCDP. The efficiency/accuracy trade-off is easily steerable in BOA by employing basis sets of different size.
Table 7: Time Efficiency. The time per molecule needed during inference is evaluated on the QM9 VASP test set for ELECTRA (Elsborg et al., 2025), SCDP (Fu et al., 2024), standard BOA and BOA with Def2-TZZV basis. The mean and standard deviation over all molecules in the test set are reported.
Model
Time per molecule [s]
ELECTRA
0.14 + 0.03
SCDP
0.58 + 0.16
BOA
1.27 + 0.27
BOA def2-T2VE
0.64 = 0.15
F ERROR DISTRIBUTION
The distribution of the electron density prediction errors on QM9 VASP is analyzed in more detail. Figures 5A and 5B show 2D histograms of the absolute errors of the predicted density for BOA large and SCDP (Fu et al., 2024), binned by the distance to the nearest and second-nearest atom. In each bin the absolute errors and normalized by the total error of the SCDP model over all evaluated grid points. Additionally, Figure 5C shows the label electron density distribution binned in the same way and normalized by the total density over all evaluated grid points. Both models show the largest errors in the regions with high electron density. In these critical regions with the nearest atom distance of less than approximately 1.0Å and the second-nearest atom at a distance of 0.5A to 1.5A, the BOA model, however, shows significantly smaller errors than SCDP
summed. The errors evaluated 100 randomly sampled molecules
G FUNCTION MESSAGE PASSING DERIVATION
We show that the messages Mabma are the coefficients in the basis of node a that represent the feature functions hom (r) of node b best in the least squares sense. The message passing step therefore solves
18

---

Published as a conference paper at ICLR 2026
A
C
Label
Dostaao 2tal Nerere Altea
Tantae Se to 2:0 Narrer Meee
de
AusEs (A)
Figure 5: Error distribution analysis. The distribution of the absolute errors of the predicted electron density on the QM9 VASP dataset is shown as a function of the distance to the nearest and second-nearest atom. The errors are evaluated on 100 randomly sampled molecules and normalized by the total error of the SCDP model (Fu et al., 2024). (A) shows the error distribution of BOA large, (B) shows the error distribution of SCDP, and (C) shows the label electron density distribution, normalized by the total density over all evaluated grid points. Both models show the largest errors in the regions with high electron density, but BOA shows significantly smaller errors in these critical regions.
the optimization problem
min Wabuu
hom (r)
Σ
mabmaaly 7e (r Ia)
(23)
2
min Mabba
Σ
homa W Zo
rb)
Σ mabmalu Zu (r ra)
(24)
Taking the derivative with respect to Mabbm and setting it to zero results in
J dr Σ homp Zo (r Io (r - verrro) rro=
dr
Σ
Mabmy Zo (r ro)w ra)
(25)
Wab homp
Was Mabmu
(26)
7 Mobmp
Σ (Waa)
himox:
(27)
which is the solution the optimization problem in Eq. 24 and exactly what implemented in the message passing step.
H EQUIVARIANCE OF BOA
To understand the equivariance of the different BOA layers, it is necessary to understand how the basis functions w (r) transform under rotations. The angular part basis functions given by spherical harmonics, which are known to transform under irreducible representations of the rotation group SO(3). These irreducible representations are indexed by angular momentum E No with dimension 2/ + 1. The basis functions tensors, i.c. there are groups of 2/ basis functions that transform under the irreducible represen- tation indexed by I. We now relabel the basis functions by i and E (-4,...,44 4.), where basis function wi, (r) is part of the ith tensor which transforms according to the representation indexed by 4. The transformation under rotation R SO(3) is given by
now be grouped according their / value into spherical
w (Rr)
Σ m
D mm (R (R)w (r),
(28)
19

---

Published as a conference paper at ICLR 2026
where D (R) is the Wigner D-matrix. A set of basis functions w (r) can therefore be separated into irreducible representations
T
w² w (r)....,,ut (r)) w (r),......
(29)
21
+
resulting in a block diagonal transformation behavior under rotations
3 (Rr) D (R))@2()) (r),
(30)
where D (R) is block diagonal matrix with blocks given by the Wigner D-matrices corresponding to the values of the basis functions. The Wigner D-matrices are orthogonal matrices, so that D (R)
is also an orthogonal matrix, i.. (D² (D² (R))
BOA employs the overlap and Coulomb matrices between the basis functions of different nodes, which are defined
Wab = f dr wr (ro)oro -rrb):
(31)
Cbb = = f dr I dr' w (r ra Na)we -1b)
(32)
|| r'll
Considering a rotation R of the molecule geometry, without changing the orientation of the basis functions, the overlap matrix transforms as
NM = / dr w (r- - Rr.) Zrr))az(t (r (r- Rro) = f dr' wfr((Rr' - Rr,) (Rr'' -Rr,)
ff
(33)
= =s dr' E up
=
1 Zo (r' rra) Σ D (R) -rb) (R)
(34)
Do (R) / dr' w ra wu) 30 (I' rb)
EE (R)
Σ
(35)
1¹
Σ
DZa pp
(R)
D (R)Wab.. u'n
(36)
Wab 2 (R) (D 2 (R))*,
(37)
where W is the transformed matrix. This notation will be used throughout this section. Completely analogously, the Coulomb matrix transforms
Ca. 2 (R)Cab (D&+(R))
(38)
Similarly, one can consider a function f(r) Ey fawg (r). in the same basis with rotated molecule geometry is given by
The rotated function f(r) represented
f(r) f(R (r) Σ
fuw R (r)r=) Σ
fn E
DZ pa
(R
Ju
Σ FunB (39)
+ Fuy D
Dz par (R 'ppp.
(40)
- = (D² F=
D=D²
(41)
H.1
MESSAGE PASSING
The overlap integrals Oabmy defined in Eq. 9 transform
Oabma
Σ
Wrab hommy ΣΣ ui (R)
DZ
Σ DZ (R) who D
(42)
Σ
DZ (R)oabmw's uy W²
(43)
20

---

Published as a conference paper at ICLR 2026
leading to the transformation behavior
Mabmp
Σ (Wan)
wu
Onhmy
-ΣEED
Dan pi
(R)(Wa))
Ne Deo (R) (R)
Rabamo"
10
14
(44)
= =ED 16.
Das, Mammmp'
(45)
of the messages Mabma. The overlap integrals aabmn are scalars, because
alammnn
Σ
hamp Wab pu thny
(46)
Σ
D uie
(R)
hammy
E
Di (R)
Σ D
(R)
Wab p
Σ Dow (R)
(47)
16"1
Σ
hamp'
Wab wou homu' Otabmnn
(48)
For scalar features any MLP is equivariant, since scalars invariant under rotations. The attention weights Gabmn therefore also scalars. The updated node features hamy transform as
hama
Σ
Σ
Rabmn Mabny
(49)
BEN (a)
Σ
Σ
Gabmmn
Dip (R) Mabmay up
(50)
bENNa (a)
W
Σ
D (R)hamp's my
(51)
showing that the message passing layer is equivariant.
H.2 NONLINEARITY
The nonlinearity is based on calculating scalar quantities, which are then transformed by an MLP, and weighting the original features with these scalars. The scalar features lawn are invariant under rotations because
Jamn Σ
hamp Goe hane
(52)
JV
-Σ>
DZo (R) hamy'
DZa (R) Σ
(R. Σ D&c (R)
Σ D RR)aan 4
(53)
Σ hamy'
Σ Oag
lamn.
(54)
Therefore, the weights Wamn are also scalars and the updated node features hamp transform as
hamp Σ
Wamntrany
Wamn Σ
Σ
Dzg (R (R)haane
(55)
n
DZ (R) (R)homa's
(56)
showing that the nonlinearity is equivariant.
21

---

Published as a conference paper at ICLR 2026
H.3 L2 NORMALIZATION
The L2 normalization of the features is done per node and feature channel, i.e. for each a and m.
The normalization factor is scalar because
hamp Was home
(57)
Diug (R)hamy'
Wi Σ Dunn (R) E De (R) Dovie (R)hamv"
(58)
hamp
Σ Wao ham Nam.
(59)
Therefore, the normalized features hamg are straightforwardly equivariant.
H.4 EDGE UPDATE
Exactly the same arguments for the 1.2 normalization and the nonlinearity can be used to show (n)
that abmma and (e) abmn are scalars. Therefore, also the generated weights sbb ((r) and v (-:-) abmo are scalars,
making the update step in Eq. 21 equivariant.
In summary, all operations BOA transform equivariantly under rotations, so that the whole model is SO(3) -qquivaana..
I HYPERPARAMETERS AND TRAINING DETAILS
The model is trained in two stages, the pre-training of the initial guess followed by the full training. During training exponential moving average (EMA) of the previous weights is used. For testing the epoch that performed best in validation chosen in all experiments. A mean absolute error loss on the electron density values on the grid is used. When evaluating the density function values are set to if both of the edge functions have a distance of more than 3 A from the grid point. The electron density prediction on the grid is multiplied with the standard deviation of the labels calculated over the whole training set. The predicted density is additionally multiplied with 0, as this stabilized training in initial experiments.
the grid, edge
1.1 MODEL HYPERPARAMETERS
The detailed list of hyperparameters of the model is given in Table 8. Settings changed for the pre- training of the initial guess are listed in Table 9. One special layer is the MLP for the edge update, which consists of two linear layers, with first SiLU activation and then LayerNorm between them. The first layer has 2 (Ne)² input and output neurons, and the second has 4 No (N 1) output neurons.
The MLP used in the nonlincarity consists of two linear layers with (No)² input and output neurons, separated by first a SiLU activation and then a LayerNorm.
The Gaussian radial embedding employed in the radial correction uses 50 Gaussians evenly spaced between 0ag and 30p (with ag being the Bohr radius). Therefore, 50 values fed into the following MLP, which consists of first linear layer with dimension 50 x 50 then SiLU activation and linear layer with dimension 50 x Nexp, Here Nexp the number of different exponents used in the radial parts of the basis functions, i.e. the number of unique radial functions. For basis functions belonging the same irreducible representations of SO(3), and therefore having the same radial part, the same correction factor is used.
Differences between BOA sizes Going from the small to the large version of BOA, the number of randomly probed points per molecule is increased from 5000 to 6000. Additionally, the batch size increased from 12 to which also leads to an increase in "epochs", since the maximum number of steps remains 500000.
22

---

Published as a conference paper at ICLR 2026
Table 8: The hyperparameters of BOA the QM9 dataset using the small settings.
Hyperparameter
Value
Batch Size
12
Initial Learning Rate
0.001
Learning Rate Schedule
AonealiAnee
Tmax
500000 steps
Optimizer
Adam
B1. ,2,
(0.9, 0.999, 10-$)
Weight Decay
0.0
Max Training Steps
500000
Gradient Clipping Value
0.5
EMA Decay
0.995
# BOA blocks
4
# hidden channels N
32
# edge functions after partial channel mean
8
MLP activation
SiLU
Initial Guess Channels
Absolute Value Scale
1000
scale factor
0.1
Probed Points
5000
Message passing cutoff Tmp
Å
Edge function cutoff re
3 Á
Basis Set
def2-QZVPPD
Table The pretraining hyperparameters that differ from the default settings in Table 8.
Hyperparameter
Value
Initial Learning Rate
0.001
Learning Rate Scheduler
CosineAAncciingWaaRReettttt Resttss
To
250
T Tuult
# Steps
1000
Differences between datasets On the MD datasets, the number of maximum training steps was reduced to 200000.
1.2 BASIS SET
The number of basis functions per atom type and shell of de f2 22-QZVPD basis set is given in Table 10.
Table 10: Uncontracted (primitive) shells per in def2-QZVPPD and total spherical functions.
Element
2
3
4
Total (sph.)
H
7
2
1
0
36
C
16
4
2
83
INN
16
8
4
2
1
83
O
16
2
86
F
16
9
4
2
86
23

---

Published as a conference paper at ICLR 2026
J LLM USAGE
Large Language Models (LLMs) were used to assist with writing (formulation and wording), coding (AI autocomplete), and gathering resources in the preparation of this article.
24


Reviews:
[{"review_id": "7zbdhRfFfm", "reviewer": "Reviewer_igug", "summary": "This paper introduces Basis Overlap Architecture (BOA), a new equivariant message passing neural network for predicting ground-state electron densities. Unlike prior approaches that expand the density linearly in atom-centered basis functions, BOA represents it as a quadratic expansion, inspired by the density matrix representation in Kohn–Sham DFT. The model leverages basis-function overlaps to define message passing between nodes, effectively incorporating both molecular geometry and basis information into the communication scheme.", "questions": "1. **(Efficiency)** Have you evaluated the computational efficiency of BOA in terms of speed and memory usage compared to ELECTRA or other recent electron-density models?\n\n2. **(Basis dependence)** How sensitive is BOA’s performance to the choice of basis set (e.g., def2-SVP vs. def2-QZVPPD)? Have you observed any notable differences in accuracy or stability?\n\n3. **(Non–atom-centric extensions)** Could BOA be extended to incorporate non–atom-centered or floating basis functions, similar to ELECTRA’s floating orbitals, to further enhance flexibility?\n\n4. **(Relation to KS-DFT coefficients)** Do the predicted coefficients correspond directly to the Kohn–Sham orbital coefficient matrix, or are they only related to the density-matrix blocks derived from it?\n\n5. **(Scalability)** What is the computational scaling behavior of BOA with respect to the number of atoms or basis functions? Have you tested the model’s scalability on larger molecular systems?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Novel and physically grounded formulation** : The idea of treating node features as functions represented in a basis and defining message passing through basis overlaps is elegant and well-motivated.\nThis approach embeds physical inductive bias (basis overlap, geometry dependence) directly into the architecture, rather than as post-hoc constraints.\n\n**Comprehensive evaluation**: Experiments on two major datasets (QM9–VASP and MD–DFT) demonstrate strong and consistent improvements. The comparisons to strong recent baselines (SCDP, ELECTRA, GPWNO) are appropriate and show clear quantitative gains.", "weaknesses": "**Lack of quantitative validation on physical observables**: While the model demonstrates impressive accuracy in predicting electron densities, the paper does not quantify how this improvement translates into physically meaningful quantities — such as total energy, dipole moments, or electrostatic potentials. Without showing how the predicted densities affect DFT-derived properties or the self-consistent field (SCF) convergence, the connection between the proposed method and its stated motivation (accelerating or improving DFT) remains incomplete.\n\n**Ablation and efficiency analysis**: The paper lacks quantitative ablations demonstrating which design choices (e.g., basis choice like def2-SVP, basis-overlap attention, quadratic expansion, learned radial corrections) most contribute to the observed performance. The large model reportedly requires ∼94 GB GPU memory, which raises questions about scalability and computational efficiency.\n\n**Minor**\n- Eq. (13) and others: Consider to change $(a,b) \\in \\mathcal E$ to $b \\in \\mathcal N(a)$ or equivalent so that $a$ remains a free index.", "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": "x6m6kVWZl9", "reviewer": "Reviewer_TU9J", "summary": "The paper introduces Basis Overlap Architecture, an equivariant GNN-based model to predict ground-state electron densities. The direct density prediction aims to replace the computationally expensive KD-DFT calculations. The authors develop an innovative message scheme that casts internal network features as basis functions and computes an overlap matrix of interacting vertices to compute messages. Instead of a linear expansion, the model casts the density as a quadratic expansion of the learned basis functions. Attention is calculated based on the overlap matrix and Coulomb matrix for additional physical inductive bias. BOA shows strong performance on density prediction on the QM9 and MD datasets.", "questions": "1. \"Superscripts l and r denote left and right, respectively, for reasons that are clear from Equation 2\" Line 155-156. It's not clear why the node features also have this directionality in the first term in Eq. 2. \n2. \"BOA uses an initial guess based on the atom types of the nodes.\" Line 199- Initial guess of the coefficients for the atom features? \n3. How is the model trained? What loss function is used?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "1. The internal density representation and the model architecture are well motivated, have the proper inductive bias. The problem of fast density generation is important for molecule discovery and characterization\n2. Representing internuclear regions without virtual nodes is a novel innovation", "weaknesses": "1.  \"In contrast to previous work, we however, do not expand the density or its square root directly as\" Lin3 165-166.  citations here would be very helpful to situate the work \n2. Unclear presentation of Figure 1 and Figure 2. What is the MN block in Figure 2. The notation is difficult to follow. In general, the paper was difficult to follow and missing key information", "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": "qEGNltVoBd", "reviewer": "Reviewer_1d9S", "summary": "The authors build an equivariant architecture to perform electron density prediction using discretized bases, ultimately achieving significantly lower density prediction errors than previous work on small molecule datasets. The key innovation is in the message passing scheme used. The paper is generally well written, although there are some parts that I am still trying to understand (see the questions below).", "questions": "Is there a benefit to predicting the electron density in this basis over learning the elements of the density matrix, similarly to what's done in Hamiltonian prediction models? Wouldn't this avoid going through the intermediate grid representation?\n\nWhat is the complexity of the forward pass? For the second class of 'volumetric' prediction methods, I imagine the scalability to larger structures is quite poor. How much does use of these discrete methods improve on that?\n\nThe basis sets used are quite large, and if I remember correctly, QM9 uses HCONF elements. Do the authors expect that this could eventually scale well to heavy elements where the number of basis functions would become even larger?\n\nWhat are the green circles in Fig. 1c? I looked at this subfigure for a while, since it seemed to be the 'visual explanation' that I was looking for, but am still confused.\n\nHow sensitive is the model to radial cutoffs? From what I've seen, MLIPs typically use smaller cutoffs (4-6 A) and Hamiltonian prediction models tend to use larger ones. Since the electron density is a more 'fundamental' quantity directly computed from Hamiltonian and density matrices, I'm wondering if electron density prediction models similarly benefit from larger cutoffs.\n\nCan the authors discuss the error metric a bit more in detail? They use NMAE, which weights all grid points equally. How is the error usually distributed among grid points? Is it concentrated in areas where the electron density has a higher spatial variation? How does the distribution of the error vary between models?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The background work seems thorough (although I am personally only familiar with ELECTRA among them), and there is a clear improvement in electron density prediction.", "weaknesses": "The main body of the paper is methods-heavy and light on results. Certain components (eg, gated nonlinearities) are more well-known and could be shifted to the appendix. More explanatory figures/schematics could be very useful.", "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 (Poster)