Learning local equivariant representations for quantum operators
Authors: Zhanghao Zhouyin
Zixi Gan
Shishir Kumar Pandey
Linfeng Zhang
Qiangqiang Gu
Venue: ICLR 2025 Spotlight (2025)
Keywords: Density Functional Theory
Local Graph Neural Network
Equivariant Neural Network
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: A Strictly Local Equivariant Neural Network for Density Functional Theory Hamiltonian/density matrix/overlap matrix learning, accelerated via invariant parameterization of overlap matrix and SO(2) convolution.

Abstract: Predicting quantum operator matrices such as Hamiltonian, overlap, and density matrices in the density functional theory (DFT) framework is crucial for material science. Current methods often focus on individual operators and struggle with efficiency and scalability for large systems. Here we introduce a novel deep learning model, SLEM (strictly localized equivariant message-passing), for predicting multiple quantum operators that achieves state-of-the-art accuracy while dramatically improving computational efficiency. SLEM's key innovation is its strict locality-based design for equivariant representations of quantum tensors while preserving physical symmetries. This enables complex many-body dependency without expanding the effective receptive field, leading to superior data efficiency and transferability. Using an innovative SO(2) convolution and invariant overlap parameterization, SLEM reduces the computational complexity of high-order tensor products and is, therefore, capable of handling systems requiring the $f$ and $g$ orbitals in their basis sets. We demonstrate SLEM's capabilities across diverse 2D and 3D materials, achieving high accuracy even with limited training data. SLEM's design facilitates efficient parallelization, potentially extending DFT simulations to systems with device-level sizes, opening new possibilities for large-scale quantum simulations and high-throughput materials discovery.

Published as a conference paper at ICLR 2025
LEARNING LOCAL EQUIVARIANT REPRESENTATIONS FOR QUANTUM OPERATORS
Zhanghao Zhouyin Zixi Gan Shishir K. Pandey Linfeng Zhang Qiangqiang Gu
School of Artificial Intelligence and Data Science, University of Science and Technology of China,
Hefei, China
2 AT for Science Institute, Beijing, China
3 DDepartment of Physics, McGill University, Montreal, Canada
*Department of Chemistry, Zhejiang University, Hangzhou, China 6 Birla Institute of Technology & Science, Pilani-Dubai Campus, Dubai, UAE
DP Technology, Beijing, China
7 Suzhou Institute for Advanced Research, University of Science and Technology of China, Suzhou, China
ABSTRACT
Predicting quantum operator matrices such as Hamiltonian, overlap, and density matrices in the density functional theory (DFT) framework is crucial for material science. Current methods often focus on individual operators and struggle with efficiency and scalability for large systems. Here we introduce novel deep learn- ing model, SLEM (strictly localized equivariant pessaae-passingg for predicting multiple quantum operators that achieves state-of-the-art accuracy while dramat- ically improving computational efficiency. SLEM's key innovation is its strict locality-based design for equivariant representations of quantum tensors while preserving physical symmetries. This enables complex many-bodt dependency without expanding the effective receptive field, leading to superior data efficiency and transferability. Using an innovative SO(2) convolution and invariant over- lap parameterization, SLEM reduces the computational complexity of high-order tensor products and is, therefore, capable of handling systems requiring the s and orbitals in their basis sets. We demonstrate SLEM's capabilities across diverse 2D and 3D materials, achieving high accuracy even with limited train- ing data. SLEM'ss design facilitates efficient parallelization, potentially extending DFT simulations to systems with device-level sizes, opening new possibilities for large-scale quantum simulations and high-throughput materials discovery.
INTRODUCTION
Quantum operators, representing observables and the evolution of quantum systems, are the cor- nerstone of describing the microscopic world. In modern quantum science, the advent of density functional theory (DFT) (Hohenberg & Kohn (1964); Kohn & Sham (1965)) has elevated single- particle quantum operators, such as the Kohn-Sham Hamiltonian, density matrix, and overlap ma- trix, to paramount importance in solving complex problems (Jones (2015)). These operators play crucial role in unravelling electronic structures, predicting material properties, and advancing quan- tum technologies. However, as tackle larger and more complex systems, these fundamental operators' efficient and accurate representation has emerged as a pressing challenge computation, demanding new avenues for innovative methodologies.
Recent advances have incorporated machine learning (ML) techniques to accelerate DFT calcula- tions by directly predicting DFT's s output of quantum operators, including charge density (Unke et al. (2021)), overlapping matrix (Yu et al. (2023); Unke et al. (2021)), self-energy (Dong et al. (2024)), wave function (Unke al. (2021)), and Hamiltonian matrix (Yin al. (2024); Yu al. (2023); Gong et al. (2023); Nigam et al. (2022); Unke et al. (2021); Zhong et al. (2023)). By cireumvent- ing self-consistent DFT calculations, such methods have the potential to scale up the electronic
ggugg@ustc.edu....

---

Published as a conference paper at ICLR 2025
(a)
L=0
L=1
L=2
(b)
L=0
L=1
L=2
Fcul
Tou
Teft
MPNN
SLEM
Figure 1: Local design of SLEM MPNN.(a) MPNN aggregation. (b) SLEM aggregation. Balls: nodes, sticks: edges, arrows: aggregation. Teut fixed cutoff;; Teff' effective cutoff. L: layer index.
structure calculations. Some of the approaches utilize Gaussian regressions (Nigam et al. (2022)), kernel-based models (Nigam et al. (2022)). and neural networks (Li et al. (2022)) to predict invari- ant Hamiltonian matrix blocks on localized frames. Notably, powerful equivariant message- mnessage-passing neural networks (E-MPNNs) have demonstrated remarkable accuracy (Han et al. (2024); Musaclian al. (2023); Batatia al. (2022): Joshi et al. (2023); Liao & Smidt; Liao al.; Simeon & De Fabri- tiis (2024); Passaro & Zitnick (2023); Zitnick et al. (2022)). These networks ensure the output tensor blocks' equivariance, respecting atomic systems' physical priors. Typically, they use iterative up- dates to build many-body interactions, achieving high accuracy while enlarging the receptive field. This limits parallelization and, consequently, the model's scalability. The storage-intens intensive quantum tensor prediction tasks exacerbate these limitations, posing considerable challenges for training such models on large datasets or predicting quantum operators for extensive atomic structures Fortu- nately, the electrostatic screening counteracts the long-term dependency in lot of material systems (Huckel & Debye (1923); Resta (1977); Ninno et (2006)). Therefore, quantum operators can be decomposed into elements dependent locally atomic structures, which natrually prefers a strictly local model that avoids expanding the receptive field. The Allegro model (Musaclian et al. (2023)) applied this concept build MI. interatomic potentials (MLIPs), achieving high accuracy and par- allelizability. While MLIPs only concern scalars (energy) and vectors (forces) (angular momentum I = 0 and 1) on each node, predicting quantum operators necessitates targeting both node and edge features on high-order spherical tensors (even up to 6 or 8). This requires locality and repre- sentability for both node and edge. Another significant challenge is computational complexity. To equivariantly mix the features of different angular momentum I, tensor products that scale as O(1)) are required (Passaro Zitnick (2023)). This makes model training, especially on heavy atoms, ex- tremely slow. This limitation hinders the development of unified MI. DFT model that generalizes across the periodic table.
This work presents novel method, the strictly localized equivariant message-passie passing model (SLEM), for efficient representations quantum operators. SLEM employs fully localized scheme to construct high-order node and edge equivariant features for general representation of quantum operators, including the Hamiltonian and density matrix. As illustrated in Fig. 1, the model embeds localized edge hidden states and utilize them construct localized node and edge features without including distant neighbours beyond a fixed cutoff range rem- This design enables the SLEMM model to generalize better, parallelize easier, and scale larger systems. Additionally, fast and efficient SO(2) convolution (Passaro & Zitnick (2023)) is implemented to reduce the O(2)) complexity, with edge-specific training weights, thereby further enhancing the model's accuracy. As for the overlap matrix, it is typically required for property calculations. Previous works have used another E-MPNN that doubles the network sizes (Zhong et al. (2023)) or extracts it from DFT calculations (Gong al. (2023); Yu et al. (2023)), incurring additional costs or incorporating out-of-loop computation steps that complicate inference. contrast, we utilized the two-centre integrals and parameterized the overlap matrix with spherical independent scalars (i.e., Slater- (SK) parameters (Slater & Koster (1954))). This method, inspired by the work of DeePTB (Gu et al. (2024)), fits the overlap operator representation with minimal additional cost.
RELATED WORKS
Message-passing Neural Networks dasssage-asssigg neural networks (MPNNs) (Gilmer et al. (2017)) have been widely applied in the modelling of atomic systems due their exceptional ac-
2

---

Published as a conference paper at ICLR 2025
euracy in capturing the intricate relationships between atomic environments and physical properties (Schütt et al. (2020).. Previous works have predominantly utilized this scheme, achieving remark- ably high precision in reported systems (Schütt et al. (2018); Satorras et al. (2021); Han et al. (2024)). However, as the MPNNs updates, the effective cutoff radius (Feir = N X (um) for each atom's 's features grows linearly with the number of update steps (N), shown in Fig. 1. Consequently, the effective neighbour list scales cubically with Teff. making parallclization intractable. Allegro (Musaelian et al. (2023)) achieved locality by modifying the updating rules by incorporating hidden pair state that depends partially on the center atom. While this framework works successfully for scalars and 1=1 vectors in potential energy prediction, it requires further improvement for fitting more general edge and node features. Other local methods avoid iterative updates that enlarge the receptive field, in- stead creating manually designed descriptors
the local environment (Behler Parrinello (2007); Bartók et al. (2010); Thompson et al. (2015); Zhang et al. (201 Wang et al. (2018); Zhang et al. (2018b)). These methods are generally local, but often balance locality and representation capacity,
Equivariant Message Passing Physical quantities, under the law of nature, should be invariant or equivariant under the spatial and temporal symmetry operations. To model such quantities, a set of neural network models has been developed utilizing equivariant operations. (Thomas et al. (2018); Weiler et al. (2018); Kondor et al. (2018); Kondor (2018)) These neural networks possess physical priors to ensure outputs transform in syne with inputs, making them more generalizable, accurate, and data-efficient in predicting physical quantities. Formally, an equivariant operation from vector space X Y is defined such that:
f(Dx = (x(xf((x) Vg G. G. X
where Dg E GL(X) is the representation of group element on vector space X. Here we consider (3) group, then y y can be composed by irreducible representation (irreps for short) that are the spherical tensors with angular and magnetic momentum index 1, m. and parity p such that [m] S I. Irreps with the same support addition/subrraccion,, while a generalized multiplication is defined as tensor product (8):
x 2 y
ls m3
Σ
ml)) (Uam))))mm
L11 y ma
TAL 7722
Where m))
my (4.ma))aamm (2))))) are Clebsch- (CG) coefficients. Conventional tensor product has the time and memory scales of ( (Passaro & Zitnick (2023)) where mmx is the maximum angular momentum in x and y. Such complexity poses great challenges for quantum tensor prediction. For example, constructing blocks of f-f and 9-g orbital pairs require irreps of maximum order = and = 8. Such high costs make training for large-size systems nearly impossible.
MODEL ARCHITECURE
3.1 PARAMETERIZE EQUIVARIANT QUANTUM OPERATORS
The equivariant parameterization of quantum operators Ô, such as the Hamiltonian and density matrix in the LCAO-based DFT framework, is illustrated in Fig. 2. The matrix element of operator Ô can be expressed as:
Old mallammim = (i, 4, m1 ف j. 12, m²)
(I)
Here and j denote atomic sites, while the angular and magnetic momentum index l, m label the atomic orbitals of the site. We apply the Wigner-Eekart theorem to decomposes the operator indexed by 4, la into a single index ly that satisfies 142 l2l < Ls < (11 + lz):
dd Og
Σ
mll).aa) (12.ma) O1.l2.mm..m
da,
i-j Here, the edge (i + j) and node (i = j) features Is.ma are grouped by the index m into vectors of in with accounting for multiple tensors for the same I. These features can be computed for hopping (i * j) and onsite = j) elements of the quantum operators. Further, by leveraging the Hermitian nature of quantum operators, the parameterized elements be reduced to upper
3

---

Published as a conference paper at ICLR 2025
(Z,, (())
x
Two-body MLP
Embedding
Riddin
50123 Layer
V
Cate
Lagar Norm
Hidden Update
Laten MLP
Linger
Edge Update
VILL
Node Update
(d)
X.
Edge
SOX2) Layer
Layer
Output MLP
Output Layer
tu 2
inverr
Lincerr
Paruuettars
Liveducible Marak Elemeuus
XP
Embed MLP
(0410g)
(dh (Olá))
Figure 2: Design of the SLEM model. (a) Hierarchical structure of the model. Starts with the atomic number Z, the radial and spherical part of the shift vector Túj and Y the initialized hidden fea- tures x5, víj, along with edge and node features e n'. are generated. The two-body hidden fea- tures predict the SK parameters constructing (off- )diagonal blocks for the overlap operator. Others features are then iteratively updated using the designed strictly localized updating scheme. (b)-(d) shows the hidden update (b), edge update (c), and node update (d). Node and edge features are used to construct the diagonal blocks for quantum operators.
triangular blocks. This almost half the demanding storage. Then, we standardised o balance the variance. Formally:
jJ 0, o. Z. g + Me.l 2,,2; 01,0
Here, Z; Zj are the atom species of atom i and j. eel 2, and He,l 2.,2, are norm and bias value for each atom and atom-pair. These values are derived from the dataset statistics and applied as weights and biases in an atom/bond type-specific scaling layer. This step helps to resolve the unbalanced norms across diagonal and off-diagonal elements of quantum operators while facilitating using ReLU-activated network learn the normalized radial dependent decaying function from
After supervising on the normalized features Oz inverse transform from Eq. 2 applied to recon- struct the predicted operator representations such Hamiltonian and density matrix blocks:
Σ
clla.ms) (12,mm) Ig:mg
Lg.ma
The whole procedure satisfies the rotational, translational and reflectional symmetry.
3.2 PARAMETERIZE OF INVARIANT OVERLAP OPERATORS
The overlap matrix, analogous to other quantum operators, is defined as:
S mam
(i, 1, m. 12, m2) = Seh (r) 2 r.y)dr
(5)
where of! is the obital from LCAO bases. The overlap matrix also satisfies the equivariance relation. Since the equivariant dependency comes fully from the bases, it is possible to rotate them to align further simplify the matrix
with the axis of raj, reducing the angular dependence. Therefore, elements into scalars via the relation (Podolskiy & Vogl (2004)):
2, / Othl (r) Tíj drr 11,22, (Fiz (rsg)mm,,n 0m,,m
(6)
4

---

Published as a conference paper at ICLR 2025
Here, the dependency changes Z, Z,, and raj, indicating the two-center nature of the overlap in- tegrals. Parameterizing these distance -eeancee-deedeentt radial functions involves encoding atomic species with radial information and simple MLP to learn the mapping to target scalars.
I |m/ (Z. Zz Tiy) S /2/////m (Tiz) (mi))mm,,1
2,
After getting these parameters, we use them to construct the overlap matrix aligned with the bond axis, and then rotate it back to the original orientation.
3.3 THE SLEM MODEL
The SLEM model architecture is illustrated in Fig. 2. The model maintains set of features, includ- 1/.L of layer (L). Specifically, ing hidden features xLi. V node features nz,7 and edge features c,t
4.2
the hidden features consist of scalar channel x and a tensor channel V The scalar channel is initialized an embedded vector containing two-body information, including the atomic species and the radial distance of the atom pair. These initialized two-body radial features are then mapped by an MLP to the invariant SK parameters for overlap. For equivariant quantum operators, such as the Hamiltonian and density matrix, scalar and tensor channels interact to generate an ordered atom-pair representations, which are used to construct local node and edge representations. After iterative updates, the representation is scaled by the statistical norm and biases, thereby achieving the final prediction.
3.3.1 FEATURE INITIALIZATION
Firstly, the initial scalar hidden feature is computed from the two-body embeddings of atomic species Zg and Zy. and the radial distance Piy. as follows:
x -= MLP 2-bedy (1(Z,) [1 ((ry)) u(rr))
Here, is vector concatenation. Atom species are embedded as one-hot vectors 1(2), and set of trainable Bessel bases ((r.)) is utilized to encode the distance Fúj between atoms and j. u(ri)) are envelope functions (Batzner al. (2022)) to add explicit radial dependency. Subsequently, the edge and hidden features are initialized as weighted spherical harmonics of relative edge vectors:
L-j = Wwe.! (LN (x)))==)) Y
L we.1 (xb.L=1) Yİ
Here, the weights are learned from the initialized scalar hidden features x Layer normaliza- tion LN ensures that the hidden tensor features have balanced amplitude of each edge. The initial node features are then computed as linear transformations of the aggregated edge features:
d. 1
Linear
L=O cdl
(10)
VNNw
JEV(()
Here N/))) and Nagg are the neighbouring atoms and the average number of neighbours of atom i.
3.3.2 SPEED UP TENSOR PRODUCT
To integrate the information from the equivariant features, the tensor product is employed in all updating blocks of the SLEM model. Generally, the tensor product in SLEM is performed with concatenated equivariant features (t) and the weighted projection of the edge shift vector ry
r; -r, ry the spherical harmonics function Y Formally:
Σ 191
(a.ma)
(m) Per di Y
(11)
11. En, Woreaahh da w 12 edge-specific parameters for each tensor product oper- ation. Performing such tensor products high-order features is computationally intensive. There- fore, we applied the recently developed SO(2) (Passaro Zitnick (2023)) convolution to reduce the computation and storage complexity from 0 (16 to O which refer the Appendix 18.
Here, 5/
5

---

Published as a conference paper at ICLR 2025
3.3.3 HIDDEN UPDATES
To construct many-body interactions, as in Fig. 2(b), the node features ni, and hidden tensor fea- tures V would be concatenated and doing tensor product with the projection coefficients of edge shift vector Piy on the spherical harmonics functions. The operation is written formally as:
VinL ca.fa = (n - [ c
w
f 02.22
Y
(12)
Unlike most MPNN, the hidden states x and V in SLEM depend only on the local environment of centre atom i, the shift vector Pij. and the atomic type and coordinate informations of atom j. Such design excludes neighbours of into hidden states.
After the tensor production, the output features V will be passed through the gated non-linearity (Batzner et al. (2022)), and transformed by "E3linear" (Geiger & Smidt (2022)) layer to mix up the information across different channels. The new hidden feature will be multiplied by the weights learned from normalized scalar features to explicitly include the radial information and return the updated feature ViJ. The scalar hidden features updated by mixing the Oth order information from V with a latent MLP, which is:
x ULL MLP 1t IIV
u((is)
(13)
The scalar hidden states I incorporate an explicit decaying envelope function u(ryy) and many- body interactions of scalar and tensor features. This formulation effectively captures the decay behaviour of each edge irrep feature as a function of radial distance.
3.3.4 NODE UPDATES
The strictly local node representation na can be constructed naturally from the many-body inter- active tensor features V We follow the MPNN style to create the message from node j to node i. Formally:
m
c
HiL Y c2.l
(14)
U.L Here again, exclude the neighbouring information of atom j in m ogyls via the partial updates of v while maintaining necessary interactions. Each message then is passed through gated activation and L3Linear layer, weighted separately by weights learnt from the hidden scalar features, and aggregated to update the node feature by:
VI
oZ
U.L cgg
8.--1 sadda
Σ JEN()
m i Ca.fa
(15)
Naeg
Here ranged from 0 to 1. The weights here differ from those in hidden updates they are directly learnt from x without normalization. Therefore, the absolute radial decay is enforced in the weights, providing strong prior that the messages from atoms shorter distances generally more significant. Meanwhile, the update of x and consequently w depends on the features Vis. and ni.--1, as shown in Eq.13. This structure aligns with the equivariant graph attention mechanism (Liao & Smidt; Liao et al.) which has demonstrated powerful expressibility in various tasks. Here w ds corresponding an attention score computed from Vij. and i Therefore, through this update, the dependencies of node features strictly local.
3.3.5 EDGE UPDATES
The locality of edge features e can naturally preserved as long as the node features are strictly local. By mixing the information of node features on both sides, localized edge updates can be formulated as follows:
J.L 2
(n ||
L Weyd 1
(16)

---

Published as a conference paper at ICLR 2025
Table 1: MAE (in meV) for Hamiltonian matrix predictions using SLEM and other methods materials with LCAO basis up to orbitals, Numbers in parentheses indicate parameter count.
Systems with LCAO-basis up to d-orbitals
SLEM
DeepH-E3
HamGNN
Material
(0.7M)
(4.5M)
(1.0 M)
(4.5M)
(2.8M)
MoS₂
0.34
0.14
0.46
0,55
1.20
Graphene
0.26
0.14
0.40
0.28
0.35
Si(300K)
0.10
0.07
0.16
0.10
0.19
Table 2: MAE (in m) for Hamiltonian matrix predictions using SLEM and DeepH-E3 models (Gong et al. (2023)) materials with LCAO basis up to f and g orbitals. Numbers in parentheses indicate parameter count. The MAE of HfO, DeepH-E3 is absent due to out-of out-of-memory errors.
Systems with ICAO-basis up to f and g-orbitals
SLEM
DeepH-E3 (1.9M)
Material
(1.7M)
GaN
0.21
0.87
HfO₂
0.28
Similarly, the updated edge features are processed via a gated activation and an E3Linear layer, and then multiplied with weights learnt from the hidden scalar features (without normalization) as:
Tنh e eggla a
ij.-11 eggls
v a²
w Li ca 0g.la
(17)
The following section focuses on validating the effectiveness of this framework via learning equiv- ariant DFT Hamiltonians, density matrices and overlap.
RESULTS
4.1 BENCHMARK THE ACCURACY AND DATA-EFFICIENCY
Using diverse datasets includes up g orbitals, evaluate model's performance in fitting Hamiltonian, density matrix, and overlap matrix. For Hamiltonian, we use systems including 2D systems of monolayer MoS2 and graphene from existing datasets (Li et al. (2022)), well as 3DD bulk silicon generated for this study. To test SLEM's capability with high-order tensors, we also train the models on generated datasets of bulk GaN and HIO, systems, which include f and orbitals. Structures in the Si, GaN, and HfO, systems are sampled via molecular dynamics using neural network potentials (Wang et al. (2018)). For density matrix and overlap matrix evaluations, we focus exclusively on the datasets of Si, GaN, and HfO, datasets, as the reported datasets for MoS, and graphene lack density and overlap matrix data. Our model supports an atom coon-scccific reat setting, each value corresponding to the radial basis cutoff of the numerical atomic orbitals, which aligns with the setting in the compared method.
Table presents a comparison of mean absolute error (MAE) values in Hamiltonian prediction for graphene, MoS,, and Si systems, whose LCAO basis extends up to d orbitals, among the SLEM, DeepH-E3 (Gong al. (2023)), and HamGNN (Zhong et al. (2023)) methods. Our SLEM model achieves state-of-the. accuracy is achieved with relatively small model size of only 0.7 million (M) trainable parameters. Furthermore, we extended our comparison between the SLEM and DeepH-E3 methods to include GaN and HfO, systems, where the LCAO basis extends up ʃ and g orbitals, respectively. As
accuracy, exhibiting the lowest MAE across all systems. Notably, this high

---

Published as a conference paper at ICLR 2025
Table MAE for predicting density matrix using SLEM on materials with LCAO basis up to f. and orbitals. Model settings align with those in Table 1.
SLEM density matrix model
Materials MAE
Silicon 8.9e-5
GaN
HfO₂ 3.9e-5
2.3e-5
1.5
1.0
0.5
E- E-FF(e
0.0
SLEM
DFT
0.5
-10
-15
-20
R2
G
02
G
V2
Figure 3: Comparison of band structures for Si MD trajectory snapshot: SLEM prediction vs. DFT calculation. Predicted band structures are obtained from either diagonalization of the predicted Hamiltonian or NSCF DFT calculation using predicted charge density, yielding indistinguishable results. Inset: Visualization of charge density distribution for the same structure.
shown in Table 2, the SLEM model consistently presents the lowest MAE, further demonstrating its high accuracy and versatility across different orbital complexities. Fig. illustrates the band structure of silicon structures computed directly from the SLEM-predicter Hamiltonian, where the eigenvalues are indistinguishable from those obtained using DFT. For the density matrix, fitting re- sults are presented in Table 3. The results demonstrate very high accuracy (order of 1e-5). approach- ing the machine precision limit float32 numbers. In Fig. 3, we use the trained model to predict the density matrix and visualize its real-space distribution. This capability is particularly important applications such as charge distribution analysis or tracking electron transfer. Furthermore, predicted density can be directly used for non-self-consisten field (NSCF) DFT calculations. The resulting band structure for silicon, as an example, is highly accurate and matches the DFT output, with MAE of only 1.09 meV in eigenvalues compared self-consistent DFT results. The overlap matrix, represented by invariant SK parameters in our SL.MM model, achieves exceptionally high accuracy as demonstrated Table 4, approaching the machine precision limit for float32 numbers. Notably, our simplified parameterization enables this high accuracy with only a minimal increase in model complexity. For instance, in a silicon model designed fit only the Hamiltonian, our typical parameter count is 0.7 M. The inclusion of overlap matrix prediction adds merely 0.01 M parameters, which is about 1.4% total model size.
Additionally, the strict localization scheme of our SLEM model confers superior data efficiency, requiring fewer DFT-calculated data points for training. To quantify this efficiency, we conducted an experiment using randomly split subsets of the original data, comprising 20%, 40%, 60%, and 80% of the full training set. We trained both the SLEM model and DeepH-E3 method on these subsets and evaluated their performance
consistent validation sets. Results in Table 5 demonstrate SLEM's high accuracy across all subsets, thus highlighting its remarkable data efficiency. This data efficiency implies that SLEM users generate smaller, more cost-effective training sets. Moreover, SLEM's excellent data efficiency and transferability make it particularly well-suited as backbone for developing universal DFT models, especially for systems involving heavy elements.
8

---

Published as a conference paper at ICLR 2025
Table 4: MAE for predicting overlap matrix using SLEM's parameterization materials with LCAO basis up to d. f. and orbitals. Model settings align with those in Table 1.
SLEM overlap prediction
Materials MAE
Silicon 5.6e-5
GaN
HfO₂
4.7e-5
6.3e-5
Table 5: Comparison of validation MAE (in mcV) for SLEM model and DeepH-E3 (Gong et al. (2023)) method trained on randomly split datasets (Li et al. (2022)) with varying training ratios. Model settings align with those in Table 1.
MoS2
Partition SLEM
100%
80%
60%
40%
20%
0.34
0.37
0.39
0.37
0.37
DeePH-E3
0.46
0.72
0.84
1.03
1.46
Graphene
Partition
100%
80%
60%
40%
20%
SLEM
0.26
0.26
0.27
0.21
0.26
DecPH-E3
0.40
0.30
0.33
0.36
0.60
4.2 EFFICIENCY AND SCALABILITY
In materials science, chemistry, and biology, many signíficant properties emerge in systems con- taining heavy atoms. These heavy atoms introduce high-order spherical tensors when representing quantum operators. Scaling to such systems challenging due to the computational complexity of tensor products used to construct complex spherical tensors, which scales O(1)). Conven- tional tensor production methods struggle with training and inference on systems containing heavy atoms, making it difficult model these important phenomena efficiently. Moreover, inferring large material systems while training with small structures is particularly valuable, which requires paral- lelising the model inference by assigning partitions of the large atomic structure to multiple GPU workers. However, most current models struggle with this task. As the receptive fields expand through iterative graph updates, the minimum size of each partitioned subgraph increases, reducing the effectiveness of such partitioning. The SLEM model addresses these challenges by efficiently constructing high-order tensor products and assisting parallelization through its strict locality.
For efficiency, the implementation of SO(2) convolution reduces the tensor product computational complexity from ((1) to O(13), which is then further reduced by the parallelization matrix operations to nearly O()) benefiting from PyTorch. Figure compares the wall time and GPU memory consumed by tensor product operations using SO(2) convolution versus the conventional method employed in DeepH-E3 (Gong al. (2023)) and E3NN (Geiger & Smidt (2022)). The SO(2) convolution approach demonstrates remarkable superior efficiency, enabling our method to handle all possible basis choices in LCAO DFT. We also evaluated memory usage and training time for typical systems, comparing our model with DeepH-E33. The results, displayed in Fig. 5, show that our model consistently outperforms in both metrics. Notably, the advantage becomes more pronounced as the basis set size increases, highlighting our method's scalability. Additionally, the SLEM model's strict locality design significantly enhances parallelization. This localized approach allows for dividing the atomic graph into sub-gras mub-graphs, enabling independent computation of node and edge features on separate devices. This extremely important when expanding DFT simulation to large systems. In practice, for HfO, with 4s2p2d2 lgg basis, typical model of 1.7M parameters can predict the quantum operators for up to 10a atoms on devices with 32GB memory. Despite system with 104 107 would require over 300 GB of memory, necessitating parallelization across multiple GPUs. Therefore, a strictly localized model such as SLEM holds significant potential for expanding simulation system sizes.
linear scaling, the inference

---

Published as a conference paper at ICLR 2025
(b)
S0(2)+r
SO(2)
10"
3-7
DeepH-E3
E3NN
10"
GPU Memory (GB)
wall time(s)
1-12
10
2-
10
Angular Momentum (f)
Angular Momentum
Figure 4: Comparison of time and memory consumption for different tensor-product implementa-
tions. (a) Time consumption vs. angular momentum (!) for different models, including the SO(2)- based SLEM model (triangles) with and without radial part (r). DecpH-E3 (cross) (Gong et al. (2023)), and E3NN (square) (Geiger & Smidt (2022)) models. Inset: Log-scale fit with slopes of for the SLEM model and 3.7 for the other two models. (b) Memory consumption vs. 1. The SLEM model demonstrates over two orders of magnitude improvement in both time and memory efficiency compared to DeepH-E3 and E3NN.
(a)
3
SLEM
DecPH-E3
27
GPU Memory (GB)
2.16
Wall Time (8)
2
10
1.36
0.73 0.88
0.97
3.6
4.2
0.71
1.9
1.6
10
Graphene
MoS2
GaN
Graphene
MoS2
GaN
Figure 5: Comparison of training time per iteration and memory consumption with SLEM and DeepH-E3 (Gong et al. (2023)) models.
CONCLUSION
This work presents SLEM (Strictly Local Equivariant Message-passe passing) model, novel approach for predicting quantum operator representations in materials science. By employing a strict locality design and integrating SO(2) convolution, SLEM achieves state-of-the-art performance in predict- ing Hamiltonians, density matrices, and overlap matrices for diverse materials, including systems with heavy atoms. The model's efficiency and scalability open new possibilities for large-scale quantum simulations and high-throughput materials discovery. Notably, SLEM's locality enables efficient parallelization through atomic graph partitioning, potentially extending its applicability to extremely large systems at device-level scales. SLEM's intrinsic support for multiple quantum op- erators, coupled with its novel overlap matrix parameterization, significantly reduces computational costs and dependence ciency and transferability, making it particularly well-suited for developing universal DFT models, especially for systems involving heavy elements. These advancements position SLEM powerful tool for simulating complex systems in materials science and computational chemistry. Future work will focus developing robust sampling methodologies for active learning and integrating SLEM with existing software ecosystems to fully leverage its capabilities real-world applications, further advancing the field of computational materials science.
post-training DFT software. The model demonstrates superior data effi-
10

---

Published as a conference paper at ICLR 2025
CODE AVAILABILITY
The implementation of the SLEM model and its semi-local variant (named LEM) are open-source and accessible via the DeePTB GitHub repository. To facilitate the integration of DFT outputs with machine learning models, we have developed and released supplementary tool, dftio. This tool enables parallel parsing of DFT outputs into machine-readable data format, enhancing the efficiency of data preparation for model training and analysis.
ACKNOWLEDGMENTS
We appreciate Mingkang Liu for his contribution in validating the SLEM/LEM model on QH9 datasets during the review process. The results presented in the Appendix A. We sincerely acknowledge his efforts and regret that we were unable to include him as a co-author due to ICLR authorship policies. We also appreciate insightful discussions with Siyuan Liu, Feitong Song, Jijie Zou, and Duo Zhang. We acknowledges the funding support from AI for Science Institute, Beijing (AISI). The computational resource is provided by Bohrium Cloud Platform from DP technology.
REFERENCES
Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi. Gaussian approximation poten- tials: The accuracy of quantum mechanics, without the electrons. Phys. Rev. Lett., 104:136403, Apr 2010,
doi: Reee 104. 136403. URL https://linkk N0.1103/PhysRevLee 1tt.eet......114410011. 136403.
org/doi/
104.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher order equivariant message passing neural networks for fast and accu- rate force fields. Advances in Neural Information Processing Systems, 35:11423-11436, neurips. 2022. URL https://proceddd/ horpaeeeeeeebbbeeaeaeeeeeeeeeeee file/
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korn- bluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, 2022. 1038 41477 022-29239- 5.
URL https: / /oii. org. ortg/...000../////
Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-energy surfaces. Phys. Rev. Lett, 98: 164001, Apr 2007. doi: oo.100//hyseeeeeet..... 146401. URL htwps: /tt.///innkk./ aps org /taprg//oooo /tlpg///////// http:......//////eeeeeetttt 98 146401.
Xinyang Dong, Emanuel Gull, and Lei Wang. Equivariant neural network for green's functions of molecules and materials. Phys. Rev. B, 109:075112. Feb 2024. doi: 10.1103/PhysRev..Ik
075112. URL https::p//.. //t///in..//////////////////l/.. aps.appss.....
www.kk......g./......ioo.oo. org///i/1d///// 1ttp::.......//h/eeeeeee 109.. 075.11
Matthias Erzzerhof and Gustavo E Scuseria. Assessment the perdew-burke-ernzern exchange- correlation functional. The Journal of chemical physics, 110(11):5029-500) 1999.
Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks. arXi::2207.09453 2022. URL
https: /atp://aaavv///aa www.wiwwww....a.ooooo
094553.
Justin Gilmer, Samuel S Schocnholz, Patrick F Riley, Oriol Vinyals, and George Dahl. Neural mes- sage passing for quantum chemistry. In International conference machine learning, pp. 1263- 1272. PMLR, 2017. URL https: / d1. acm. org/ /ttg////111..55/3055533333
11.11.555533303 55555/3335581... 55355535311111335555222 3305512.
Xiaoxun Gong, He Li, Nianlong Zou, Runzhang Xu, Wenhui Duan, and Yong Xu. General frame- work for E())-qquivariant neural network representation of density functional theory Hamilto- nian. Nature Communications. 14(1):2848, 2023. URL https:/ /oii org/ 1038/ 941467-023- -000223333665000 8.
Qianggiang Gu, Zhanghao Zhouyin, Shishir Kumar Pandey, Peng Zhang, Linfeng Zhang, and Weinan E. Deep learning tight-binding approach for large-scale electronic simulations at finite temperatures with ab initio accuracy. Nature Communications, 15(1):6772, 2024.

---

Published as a conference paper at ICLR 2025
DR Hamann.
Optimized norn-conserving vanderbilt pseudopotentials. B Condensed Matter and Materials Physics, 88(8):085117, 2013
Physical Review
Jiaqi Han, Jiacheng Cen, Liming Wu, Zongzhao Li, Xiangzhe Kong, Rui Jiao, Ziyang Yu, Tingyang Xu, Fandi Wu, Zihe Wang, et al. survey of geometric graph neural networks: Data structures, models and applications. arXiv:2403 00085, 2024. URL https: aawii.iiiivv///./ org.iioorgg//ooo org. httgg/bb//4400 00485.
P. Hohenberg and W. Kohn. Inhomogeneous electron gas. Phys. Rev. 136:B864-8871, Nov 1964. doi: 10. Setnhssev..3...86.. URL https:///inkk http:s//iiinaas.....o. PhysRev. 136. 136.B864.
aps. org/ /1o/d//i111110 1103/
Huckel and P Debye. Zur theorie der elektrolyte. i. gefiterpunkteericcdiiuugg und verwandte erscheinungen. Phys. 7, 24:185-206, 1923.
Simão M João, Miša Andelković, Lucian Covaci, Tatiana G Rappoport, João MVP Lopes, and Aires Ferreira. Kite: high-performance accurate modelling of electronic structure and response functions of large molecules, disordered crystals and heterostructures. Royal Society Open Sci- ence, 7(2):191809 809, 2020. URL. https:: royal societypubl shi shing. org/ do' / / 10. 10.1098/rsee rsos....s.... 191809.
R. O. Jones. Density functional theory: Its origins, rise to prominence, and future. Rev. Mod. Phys., 87:897-923, Aug 2015. doi: 10. Mod//RerMoPyyy.8....5 897. URL hhttps:///htt aps.ar.......... 00i//0.1103 RevModPhys. 87..887888.. 87.897.
Chaitanya Joshi, Cristian Bodnar, Simon Mathis, Taco Cohen, and Pictro Lio. On the expres- sive power of geometric graph neural networks. In International conference on machine learn- ing. pp. 15330-15355. PMLR, 2023. URL https: /ttp///ooooeeeeeeeoooooooc ngs www.mmllrr....... mml..press/v200.
Doshi23a/joshi23::
W. Kohn and L. Sham. Self-consistent equations including exchange and correlation effects. Phys. Rev., 140:A1133-A1133 Nov 1965. doi: 10.1 10.1103/PhysRe.. htt.///hysee..1...1133333 140.A1133. URL https:
//link. aps wwtp///g/////eeoeoo http://///.....././..... 1103/PhysRev. 140 A1133.
Risi Kondor. N-body networks: covariant hierarchical neural network architecture for learn- ing atomic potentials. arXiv:/803.01588, 2018. URL httts:://aaxii// orgp///////./.....
01588.
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-Gordi nets: fully fourier space spherical convolutional neural network. Advances in Neural Information Processing Systems, 31, 2018. URL httpp:::pppoceeeeepeeette 4tt.://00004421111--------------
neurips. ec cc/paper__ppeeeeeeeeeeeeeeeeeeee llappe_falle_.____aaaaaaaaaaeaee files/paper/2e/ ppppp//pppp////////ppp//
pdf.
He Li, Zun Wang, Nianlong Zou, Meng Ye, Runzhang Xu, Xiaoxun Gong, Wenhui Duan, and Yong Xu. Deep-learning density functional theory hamiltonian for efficient ab initio electromic-structure calculation. Nature Computational Science, 266) 2022. URL https: / /tw.i//o//////////////.. .or.
022 0026666.
Pengfei Li, Xiaohui Liu, Mohan Chen, Peize Lin, Xinguo Ren, Lin Lin, Chao Yang, and Lixin He. Lcrge-scale ab initio simulations based on systematically improvable atomic basis. Computational Materials Science, 112:503-511 2016.
Yunhai Li, Zhen Zhan, Xueheng Kuang, Yonggang Li, and Shengjun Yuan. TBPLaS: A tight- binding package for large-scale simulation. Computer Physics Communications, 285:108632, April 2023, ISSN 0010-4655, doi: 10. 0010166..pc:2222.1086333
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations.
Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivari- ant transformer for scaling to higher-degree representations. In The Twelfth International Confer- on Learning Representations.
12

---

Published as a conference paper at ICLR 2025
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kombluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atom- istic dynamics. Nature Communications, 14(1):579, February 2023. ISSN 2041-1723. doi:
10.1038/s41 0411088/411467-023--333 467-023-363329-
Jigyasa Nigam, Michael J. Willatt, and Michele Ceriotti. Equivariant representations for molecular Hamiltonians and N-center atomic-scale properties. The Journal of Chemical Physics, 156(1): 0141 01 2022. ISSN 0021-9606. doi: 10.106335.0072784 URL https: //ttts:://di.... org//0. /C. 1063 0072784.
D. Ninno, Trani, G. Cantele, K. J. Hameeuw, G. Iadonisi, E. Degoli, and S. Ossicini. Thomas- Fermi model of electronic screening in semiconductor nanocrystals. Europhysics Letters, 74(3): 519, mar 2006. doi: 10. 10.1209//l//20005- 10544-9. URL https: //tp wew.:o/ii.ooog.... 100 1209/ 201/12005-1054442
Saro Passaro and C Lawrence Zitmick. Reducing SO(3) convolutions to SO(2) for efficient equiv- ariant gnns. In International Conference on Machine Learning, pp. 27420-27433 PMLR, 2023.
URL https::///// /dl. acm. org/ hts///oi/1..00 5555/ 3618408. 111440886666900 9548.
A. V. Podolskiy and P. Vogl. Compact expression for the angular dependence of tight-binding Hamiltonian matrix elements. Phys. Rev. B, 69:233101, Jun 2004. doi: 10.1 10.1103/PhsseevB.55 233101. URL https: http///nnnkke wwa.kaaps.co wrt.oo/g///oooo...oo /1c.
69. 69.233101.
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):177,0014
Raffaele Resta. Thomas-fermi dielectric screening in semiconductors. Phys. Rev. B, 16:2717- 2722, Sep 1977. doi: 10.1 10.1103/hhsReev..1222 16.2217.... URL http::///innkaas..orrrr. / 10.1103/ 1111100/Physeez2 27717.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(1 21:2) 2875, 2012.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural net- works. In International conference on machine learning, pp. 9323-9332. PMLR, 2021. URL https: http:///ooeeeeeninneeeee mlr.press. / vt1//vv1v/v///////// satorras2la wttps///aaaaatttraaaa..a... pdf.
K. T. Schütt, H. F. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet A deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24): 241722, 03 2018. ISSN 0021-9606 doi: 10.1063/1.501979, URL batpss///oll. org 10. 1063/1.501977
Kristof Schütt, Stefan Chmiela, O Anatole Von Lilienfeld, Alexandre Tkatchenko, Koji Tsuda, and Klaus-Robert Müller. Machine learning meets quantum physics. Lecture Notes Physics, 2020.
Guillem Simeon and Gianni De Fabritiis. TensorNet: Cartesian tensor representations for efficient learning molecular potentials. Advances Neural Information Processing Systems, 36, 2024. URL https: oppnrevieww net wwtp:/eeee///iiiiieooeoaoo
forum? Bae.ummiid--BiPPaaa--
J. C. Slater and G. F. Koster. Simplified LCAO method for the periodic potential problem. Phys. Rev., 1494448-15 Jun 1954. doi: 10. 00.1003/Physeev.9..1 1498. URL org/doi// / / 103 1103/Physeev 94.1498
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation eauiiiiionn-uuvviiinnnnaa neural networks for 3d point clouds. arXiv: /80X..8800.00 802.08219, 2018. URL htt.hhttss///aaa ptttttpss//axxii..ooo
wrg/abs/////oo.tooo.oo...o... htt.ttpps./aaaiw.orrr//aassssss 8028 08219.
Aidan Thompson, Metin Aktulga, Richard Berger, Dan S Bolintineanu, Paul S Crozier, Pieter J In't Veld, Axel Kohlmeyer, Stan G Moore, Trung Dac Nguyen, et al. Lammps-a a flexible simulation tool for particle-based based materials modeling at the atomic, meso, and continuum scales. Computer Physics Communications, 271:108171, 2022.
Michael Brown,
13

---

Published as a conference paper at ICLR 2025
A.P. Thompson, L.P. Swiler, C.R. Trott, S.M. Foiles, and G.J. Tucker. Spectral neighbor analysis method for automated generation quantum-accuratt interatomic potentials. Journal of Com- putationna Physics, 285:316-330, 2015. ISSN 0021-9991. doi:
2014.12.018. URL https:/ /ttp::/wwww... /ww. sci
encedi www.sccincediiiiic rect. com/ watm//cienne/arrtiillcn
000219991140083355
Oliver Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus- Robert Müller. SE(3)-equivariant prediction of molecular wavefunctions and elec- tronic densities. Advances in Neural Information Processing Systems, 34:14434 14447, 2021. URL httpsss//pooceeeeng..eeeee wwtppripps.cppppppee.eepee cctppppee//2///11eeeeeee 78f1893678a fbeas 90b1 Fw01b9cfb860----0000000 Paper. pdf.
Jonathan Vandermause, Steven B. Torrisi, Simon Batzner, Yu Xie, Lixin Sun, Alexie M. Kolpak, and Boris Kozinsky. On-the-fly active learning of interpretable Bayesian force fields for atomistic rare events. upj Computational Materials, 6(1):1-11, March 2020. ISSN 2057-3960. doi: 10.
1038//41524-000-02833
Han Wang, Linfeng Zhang, Jiequn Han, and Weinan E. DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics. Computer Physics Commu- nications, 228: 178-184, July 2018. ISSN 0010-4655. doi: ...0166j.cc..211.03.0111
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Co- hen. 3D Steerable CNNs: Learning rotationally equivariant features in volumet- ric data. Advances in Neural Information Processing Systems, 31, 2018. URL https: /tpos://rree http://pccceeeegsseeee neurips. CC/ cc/ppper.p.....//ee/peeeeeee.. files/ Btteee/papee/221/EE111 Pawaaafaaapepppdae
Shi Yin, Xinyang Pan, Fengyan Wang, Feng Wu, and Lixin He. A framework of co(3)-equivarian non-linear representation learning and its application to electronic-structure hamiltonian predic-
tion. arXiv:2405.05722, 2024. URL https: /ttp://iii..../aa/////..... 05722.
Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Efficient and equivariant graph networks for predicting quantum hamiltonian. In International Conference on Machine Learning. pp. 14412-0442, PMLR, 2023. URL
v202/ /y23i. html.
Haiyang Yu, Meng Liu, Youzhi Luo, Alex Strasser, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules. Advances Neural Information Processing Systems, 36, 2024.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep potential molecular dy- namics: A scalable model with the accuracy of quantum mechanics. Phys. Rev. Lett., 120: 143001, Apr 2018a. doi: 10. 10.1103/hyseevLeet.. 120. 120. 143001. URL https: /ttp:///iik.... aps orggooogiooooo doi/ 10. 1tt.1113/Phyyh www.:///hh/.eeeet....m 120 143001.
Linfeng Zhang, Jiequn Han. Han Wang, Wissam Saidi, Roberto Car, and Weinan E. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. Advances in Neural Information Processing Systems, 31, 2018b. URL htlle//22acd62222266666666333
www.iie...p.... cc. wcc/paper__te _fpper_fille wttrffffeeeepe////011/ pdf.
Yuzhi Zhang, Haidi Wang, Weijie Chen, Jinzhe Zeng. Linfeng Zhang, Han Wang, and Weinan E. DP-GEN: concurrent learning platform for the generation of reliable deep learning based po- tential energy models. Computer Physics Communications, 253:107206, 2020. ISSN 0010-4655. doi: htss://doi.org/100 http://0...0011....2000000010000 com/ http://ciinneerrrrllee//////11./
URL https::/ww..scieeeeee /www///ww//// scieecedreedr..r.......
Yang Zhong, Hongyu Yu, Mao Su, Xingao Gong, and Hongjun Xiang. Transferable equivariant graph neural networks for the hamiltonians of molecules and solids. upj Computational Materials, 911):182, 2023. URL ht /doi. orgg 038/ 41 524 020000130---
14

---

Published as a conference paper at ICLR 2025
Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems, 35:8054-8067,
2022. URL. https: http:////oooedooeeeeeeoooooe ngs neurips cc httpp/ppppeelleepppeeeeeeeeeeeee File/3501bea acc61 htt.ee/eeecccocoooaoffooefeeeaed Conerrereeppe--c
A BENCHMARK ON QH9
We also train our model on the open molecule Hamiltonian and Overlap dataset, QH9 (Yu et al. (2024)), test our method's accuracy and transferability on large dataset. We compare our re- sult with the reported benchmark QHNet, which shows considerable accuracy improvement, as dis- played in Table. We should notice that our model performs well the Hamiltonian learning tasks, by decreasing 1/3 of the enror reported previously. The density matrix is inferred from the learned Hamiltonian since the data is not available in QH9. Our accuracy only 1/3 of the compared one.
QH9 is an extended electronic structure dataset based on QM9 (Ruddigkeit et al. (2012); Ramakr- ishnan et al. (2014)), a well-known benchmark for molecular property prediction. QH9 contains 13,081 static structures and 2998 molecular trajectories. We test our method based on the static (Ootooddiisributiin(OOOO)) tasks defined by QHNet's paper. In this task, the testing set has a very different molecular structure and atom number from the ones in the training set. The dataset is split as 1w..000111..999.... As discussed in Appendix G, we used semi-local variant of SLEM to prevent the possible break of strict locality in the molecular system.
Table 6: MAE for Hamiltonian matrix (H) and density matrix (D) prediction.
QH9 dataset benchmark
Method
H[10- E,
all
non-diag
diag
all
LEM
56.57
44.86 80.26
155.24 135.63
0.0219
QHNet
83.22
0.0643
B ON THE STRICT LOCALITY APPROXIMATION
Here we want to further discuss the strict locality approximation made in our method, In the quan- tum operator learning task, we express the physical operators using an LCAO basis. In this expres- sion, the theoretical framework of DFT fundamentally supports the strict locality approximation. In LCAO-based DFT calculations, the Hamiltonian matrix elements exhibit inherent locality. This locality manifests through two mechanisms: the spatial decay of LCAO basis functions, the electron neuralized conditions of the physics system, and electronic screening effects especially in condensed matter systems. These effects make the dependency of the quantum operator elements approximately local. For periodic systems, particularly metals and semiconductors, electronic screening effectively reduces long-range Coulomb interactions. The screening length varies by material type around Á in semiconductors like silicon, and Germanium, and smaller in insulators such as 2.76 A in Diamond Therefore, we safe to choose cutoff radius match the sum of basis orbital radii, adequately captures the correct physical relations. Taking the Hamiltonian matrix elements as an example:
V = f 6r RA) (Vx (r)+ + (rae(r) + Vic (r)) (r RB
The locality is combination of the LCAO basis, the Hartree term, the external potential, and xe functional (we ignore the kinetical term for its local nature). Therefore, we need to discuss the localization collective effect, which is contributed by the following properties:
15

---

Published as a conference paper at ICLR 2025
The decaying behaviour of the LCAO basis. For all LCAO bases (whether GTO, NAO, Slater-like), the radial part all decays rapidly with distance. Therefore, the long-range term in Vn, Vezt: Vae's contribution in the integration (that decays as 1/r) would be decreased fastly (of- ten exponentially such as in GTO or Slater-like orbital).
The system's electronic neutralization condition. The electronic neutralization condition re- quires that the long-term part of the Hartree term and the external term cancel each other. To derive this, we assign the electron density to the atomic center, by writing as:
n(r) Enr(+ R1)
Then the sum of (r), (ewe(r) would be:
Z1
Σ / ny(r' Ri) dr
R1
When |r- Ry| is large, we can approximate
|r -r-1 | 20 Ir- Ryl. therefore:
Z1 R₁ +
11(rr' R1) dr ~
Z1
0
R₁
Rr
The screening effect. we refer to Huckel Debye (1923); Resta (1977); Ninno et (2006) the study of screening radius Rs, which describes the system's electrostatic potential's reaction to the vibration of charges.
In Resta (1977), the screening radius of typical insulator and semiconductor systems is reported as: Diamond: 2.76 a.u, Silicon: 4.28 a.u. and Germanium: 4.71 a.u. In Ninno et al. (2006), for some nanoparticles, the R, is reported as: Si191H148: 5.36 a.u. and Ge191H148: 5.61 a.u.
These effects altogether motivate us to revise the current method for machine learning electronic structures and design SLEM,, striet localized model suitable for learning the correct mapping from structures to quantum operators with the correct physical priors-a more reliable and data- efficient method. Mcanwhile, we also notice the existing limitation of this method some special systems, we refer to the discussion in Appendix G for detail.
C COMPARISON OF SLEM WITH MPNN
Neural Networks In the message-passing scheme, atoms are treated as nodes in graphs, with bonds to neighbouring atoms represented as connected edges within specified cutoff radius. The embedded atomic features are processed by trainable functions, generating messages from each edge update the embeddings of central atoms. Formally, the MPNN framework can be summarized as follows:
LULL =N, (n 1 .tt.L-1)
m u.L = My (n 5,6-1 e
i. = Ut.
n 1.h-11
Σ
m i.
JEV(
Here, e.ih represents the edge features, miL denotes the messages, and n'-4 indicates the node features at layer L. NL. Mr, and Un are the trainable functions for the edge, message, and node updates. A(i) [jlray < reur} indicates all the neighbour atoms for atom i, with reur being the predefined cutoff. This updating framework allows for the construction many-body interactions and long-term dependencies, leading to strong performance across various applications. Here
16

---

Published as a conference paper at ICLR 2025
reframe the updating rules of SLEM with the MPNN framework, which looks like this:
V V5 (n+ visle1
e =No (n V n ,eti
m UHI -M (n vv..))
n' L = UL
n n-
m
JEN())
Here Vz. is the neural network for hidden feature construction. This update scheme constructs many- body interactions to build equivariant edge and node features while preserving the absolute locality by excluding atoms outside the constant cutoff radius. Such scheme, in principle, would have much better transferability, data efficiency, and scalability when we have a strong prior that the input and output are dependent locally.
D MORE ON TENSOR PRODUCT
To integrate the information from the equivariant features, the tensor product is employed in all updating blocks of the SLEM model. Generally, the tensor product in SLEM performed with the concatenated equivariant features P and the weighted projection of the edge shift vector ra
r; -r r. the spherical harmonics function Y Formally:
p ca.la
P
& 22
Σ c1,11d2
Σ mi..1ng
clla.ma)
Pii Y (1.mm) dimm
(18)
/1.11.1 Eng Wo Way 12.12 are edge-specific parameters for each tensor product oper- ation. Performing such tensor products on high-order features computationally intensive. There- fore, we applied the recently developed SO(2) (Passaro & Zitnick (2023)) convolution to simplify, reducing the computation and storage complexity from ( (1 The simplification idea intuitive. Yúj 12,m2 are sparse tensors rotated to align with the edge ij, which is nonzero only for m2 = 0. Therefore,
Here, Wh
easier to compute the tensor production in the direction edge ij,
and rotate inversely the output afterwards. This step removes the m2 index from the summation in Eq. 18. Furthermore, considering the Clebsch-Gordan coefficients with my = O, we find that (la.ma)
C (f1.m ()) except for my =m This allows further reduction of the summation in Eq. 18 by replacing i- with single index m. Then the operations be reformulated formally as:
Jeim fat.-
Σ C
We,e', m Weeel.V... We.c!....m
P P.P..-m
This represents linear operation on Plj fe By employing this method, high-order tensor products for 8, and even 10 can efficiently calculated, which is essential for heavy element systems where the f, g or even higher orbitals are used in the LCAO basis for DFT calculations. Finally, the weights for the new SO(2) tensor product method are multiplied by edge-specific parameters, where wh are mapped by an MLP from hidden scalar features x as (i ",e', f' = We,e' w p- This powerful and efficient tensor product layer facilitates the construction of local interactive updates of the features.
E DATA GENERATION
In this section, discuss the data generation process used in experiments of this work. The data sampling process includes the Hamiltonian, overlap and density matrix of materials Si, GaN, HfO,-
First, we perform ab initio molecular dynamic simulation based neural network force field using DeePMD (Wang et al. (2018)) and Lammps (Thompson et al. (2022)). typical input file of the lammps sampling looks like:
17

---

Published as a conference paper at ICLR 2025
variable
NSTEPS
equal 500000
variable
THERMO FREQ
equal 100
variable
DUMP _REQ
equal 1000
variable
TEMP
equal 300
variable
PRES
equal 1.00
variable
TAU_T
equal 0.10
variable
TAU_P
equal 0.50
units
metal
boundary
P p p atomic
atom_style
neighbor
1.0 bin
read_data
conf. srnf.lammps-oee lammps-data
mass
1 28.085
pair_style pair_coeff
deepmd graph.pb
* Si
thermo_style
custom step temp pe ke etotal press vol 1x ly 12 xy XZ $ [THERMO_RREQ)
thermo
restart
1000000 dpgen.
velocity
all create S(TEMP) 449414
fix
all nvt temp (TEMP) } $ (TEMP) 0.1
dump
1 all custom [REM] lammpsst id type y vy
timestep
0.001
run
100000
undump dump
2 all custom $ (DUMMP_FRE)) sample. lammpstrj id type x VX vy v2
run
$ (NSTEPS) upto
The samples are taken after 10t steps when the system reaches equilibrium.
After the configuration sampling, we then applied ABACUS (Li et al. (2016)) to compute the quan- tum operators, including the Hamiltonian, density and overlap matrix of each configuration. Among them, 150 frames are randomly split as the training set, and 30 frames are used for testing. We use SG15 ONCV numerical atomic orbitals, pscudopotentials (Hamann (2013)), and PBE functionals (Ernzerhof Scuseria (1999)) for the DFT calculation. A typical DFT calculation input looks like:
INPUT_ PARAMETERS
Created by Atomic Simulation Enviroment
ntype
ecutwfc
100
scf_nmax
100
smearing_method
gaussian
smearing_sigma
0.002
basis_type
1cao
ks_solver
genelpa
mixing_type
pulay
mixing_beta
0.7
scf_thr
1e-07
out_chg
symmetry
calculation
scf
out_band
18

---

Published as a conference paper at ICLR 2025
force thr out out_stru kspacing Ispinorb out_wfc_lcao
0010
0.08
0
dft aft_functional
pbe
out_mat_hs2
True
After the calculation converged for all, processed the data format into machine-leaming read- able datasets and started training. The dataset used in this work uploaded in the opensource platform AlSquare via this link: https: fttt:::/www.aisssuarrs.. com datasets//eeeaeee
Quant un__uuunrrpprrrreeeeee -at
F CUTOFF SELECTION
The cutoff used in the SLEM model is important since it decides the size of the atom and bond environment consider when learning the map from the atomic structure to the quantum operators. There are two cutoff concepts in the SLEM model. One is the cutoff for the bond rb, which is used build the one-to-one correspondence of the graph structure of atomic data to the diagonal and off-diagonal blocks of the quantum operators. Another one is the cutoff for environment
For rb. since the bond whose radial distance beyond r3 would all be considered this value should be fixed according to the LCAO orbital used in DFT data generation. For numerical atomic orbitals (NAOs), the cutoff is defined clearly in orbital files or DFT software, therefore, we just double (since two orbitals constitute bond) and assign the value to each atom there Th. For Gaussian- type orbitals (GTOs) and Slater-type orbitals (STOs), since they decay exponentially along radial distance, often set some precision threshold during data processing and find the longest bond in off-diagonal blocks with a value larger than the preset threshold. The length of the bond will be considered тb-
The environmental cutoff Te is very system dependent. This decides how large atomic/bond envi- ronments relate to the mapping. This should follow the locality behaviour of the physical system. For example, in typical periodic crystals such as semiconductors (Si, GaN, etc.), medals (Graphene), and insulators (HfO2) the r. are quite local, therefore, is safe systems experience strong nonlocality behaviour, should be very large. One should carefully bal- ance the efficiency of choosing large Te that ensures strict locality and the usage of a semi-loca or MPNN-based method.
the r. rb. However, when
G LIMITATION
The SLEM model, benefiting from the Strietly localized design, performs better in accuracy and transferability as demonstrated in various datasets. We acknowledge that such a localized hy- pothesis is most suitable for describing periodic systems, which is commonly adopted in physics, chemistry and material science research. For confined systems such as molecules, the absence of screening effect could lead to long-term dependency that is uncovered within prefixed cutoff. In these cases, the SLEM model wouldn't perform as well as it does in periodic ones. For gener- ality, also proposed semilocal model called LEM (Localized Equivarient Message-Passing where the interaction between distant atoms is included but decays exponentially with their distance, with trainable decay factor. All the designed models are ready /
in GitHub repository:
cow. com dww//eeeeooeeeeeeeoeeeeeeeeeeeee ng DeePTB.
H FUTURE INVESTIGATION
While SLEM demonstrates efficacy across various applications, several areas warrant further inves- tigation:
19

---

Published as a conference paper at ICLR 2025
Sampling Methodology for Active Learning: Scientific computation tasks require high con- fidence in calculated results, which can be challenging for data-driven approaches. Devel- oping robust sampling workflow for active learning is essential for reliability. While techniques like uncertainty-driven sampling with Gaussian regression (Vandermause et al. (2020)) model ensembles (Zhang al. (2020)) have addressed confidence issues in ma- chine learning force fields, how to design sampling workflow specifically for quantum operator models remains an open question.
Software Integration for Post-Processing: Integrating the model with existing software is vital, especially for high-throughput calculations and large atomic systems beyond conven- tional DFT capabilities. An efficient, parallelizable solver for extracting physical quantities from the model's predictions is highly beneficial. While some software based on stochas- tie techniques (Li al. (2023); João al. (2020)) shows promise, open-sourced and highly-optimized solution for non non-orthogonal bases remains unavailable.
Future research addressing these challenges will further enhance the applicability and reliability of SLEM and similar quantum operator models in computational materials science and chemistry. We are excited and look forward to seeing these happen.
20


Reviews:
[{"review_id": "2SLn7XVwPu", "reviewer": "Reviewer_hW5o", "summary": "This paper proposes a data driven method to solve KS-DFT. Instead of solving the KS system to consistency, this method puts the configuration of atoms through a carefully designed SO(2) neural network and directly predicts the quantum operators at self consistency. While there are many other existing works trying to accomplish the same thing, the key technical contributions of this paper can be summarized as the following two points:\n\n- The SLEM architecture, compared to traditional methods, SLEM has a strictly local design, its effective cutoff does not increase as more layers are added. This architecture is more scalable/parallelizable because the dependency is much smaller.\n- The parameterization of the invariant overlap operators, which enables the prediction of overlap operator without requiring lots of parameters.\n\nThe above innovations is verified to be effective in the empirical evaluations on various systems. The advantages reported include \n\n- Better generalization due to the more restricted model assumption.\n- Better scaling behavior w.r.t angular momentum.\n- Faster iteration speed and smaller memory footprint.", "questions": "My questions are stated in the above concerns, I would be happy to raise my score if the authors address them.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The two key contributions are novel, they are clearly explained in the paper and based on my understanding they are technically sound.\n- The strictly local structure is significant and likely to be widely adopted in the future given the nice properties, not only more parallelizable, but also leads to lower errors.", "weaknesses": "I have a major concern on the way the evaluation is carried out\n\n- The training and testing happens on the same system using trajectories of molecular dynamics. Although this type of evaluation may be also used in the baselines that the author compare to, I feel it is not sufficient. A good generalization is not surprising if the MD trajectories have a good coverage of different atomic geometric configurations. We run into the chicken egg dilemma, if DFT is already calculated on a system, why would we need to fit a blackbox model and do it again faster. In my opinion, the evaluation should be performed across different materials in a combinatoric way, i.e. train on material made of `AB` and `BC` and `AC`, evaluate on material made of `ABC`. Or on materials that contain same elements but in different proportions, i.e. train on `1A2B` and evaluate on `2A1B`.\n\nAnother concern is on the theoretical soundness\n\n- Based on my understanding, only atoms pairs the has a distance smaller than rcut is considered, which means that the interaction between atomic orbitals from distant atoms are not considered, however, the operator needs to include all pairs of interactions to form a matrix. How are the noninteracting entries of the matrix set?\n- Although the empirical study favors a strictly local structure, it is unintuitive theoretically. For example, the hartree term in DFT is a slow decaying term, using `(ij|kl)` to represent the four center integral, and `(i|j)` as the overlap; When `(i|j)` and `(k|l)` are both large, the `(ij|kl)` term is not negligible even when `ij` is distant from `kl`. It would added to the soundness of this paper if the authors could provide a theoretical support.\n- With the above said, it is crucial to discuss the limitation of this method, i.e. in which scenario would this method fail due to the strictly local assumption.", "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": "HosyC5wFG0", "reviewer": "Reviewer_VuZs", "summary": "The paper introduces Strictly Local Equivariant Message-passing (SLEM), a deep learning model designed to predict multiple quantum operator matrices, such as Hamiltonians, density matrices, and overlap matrices, within the density functional theory (DFT) framework. SLEM tries to address key challenges in efficiency and scalability that traditional methods face when handling large quantum systems. The authors validate SLEM’s performance across 2D and 3D material systems, demonstrating high accuracy for the experiments.", "questions": "1) What was the reason for selecting only two models for comparison with SLEM, and how were these models chosen? Would you consider expanding your evaluation to include additional, widely recognized benchmark datasets, as previously suggested?\n\n2) In the computational scalability analysis, SLEM was compared with E3NN. However, I noticed that E3NN’s performance metrics, such as MAE, were not included in the accuracy tables. Did E3NN achieve better accuracy compared to SLEM, even if it was more computationally expensive?\n\n3) The paper states that structures for the Si, GaN, and HfO₂ systems were sampled via molecular dynamics using neural network potentials. Since these potentials may introduce approximations relative to traditional quantum mechanical methods like DFT, could you clarify the potential impact on the quality and reliability of your training data? How do you ensure that these approximations do not compromise the accuracy and generalizability of SLEM’s predictions?\n\n4) Could the authors share more details about the in-house simulations used for data generation, such as specific parameters and configurations? Providing this information would greatly aid in the reproducibility of the study as well as the quality of the data.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "By focusing on a strictly localized equivariant message-passing framework, the authors present a creative way to address the challenges of efficiency and scalability in quantum mechanical computations. The use of SO(2) convolutions to manage high-order tensor complexity is particularly novel, as it reduces the computational burden associated with f and g orbitals.\n\nThe methodological rigor of the paper is detailed theoretical justifications for the design choices of SLEM. The authors provide mathematical foundations for the model's ability to preserve physical symmetries while maintaining strict locality utilizing quantum mechanical properties. The paper is well-organized, with a logical flow from the problem motivation to the model formulation and experimental validation. \n\nThe contributions of this work is important, especially for the fields of quantum chemistry and materials science. The ability to efficiently predict quantum operators without expanding the receptive field could accelerate DFT calculations and open new possibilities for simulating large-scale quantum systems.", "weaknesses": "The experimental comparisons presented in the paper are limited to only two other models, and these comparisons are not consistently provided for all experiments. Expanding the range of baseline models, including more well-established methods, would strengthen the validation of SLEM’s computational efficiency, scalability, and accuracy. Incorporating additional well-known benchmark datasets, such as QH9 [1], nablaDFT [2], and potentially QM9 [3, 4] (used in models like HamGNN), could provide a more reliable and widely recognized basis for evaluation. This would help clarify the practical advantages of SLEM more comprehensively. Additionally, while the authors conduct in-house simulations for some datasets, details about these datasets are not provided. More transparency about the data generation process and the choice of neural network potentials for sampling MD simulations would address concerns about the accuracy, quality, and reproducibility of the dataset. Given that datasets are often created using well-established computational methods like DFT, clarifying these choices would be beneficial.\n\nMoreover, the mathematical formulation of SLEM is quite complex and may be difficult for practitioners who are not experts in quantum mechanics or advanced tensor operations. Providing a more accessible explanation or simplified overview could make the approach more approachable. Additionally, combining some of the experimental tables into one could improve readability and streamline the presentation.\n\nIn addition, the mathematical formulation of SLEM is highly complicated and could be difficult to implement for practitioners not deeply familiar with quantum mechanics and advanced tensor operations. The paper could be improved by providing a more accessible explanation. Furthermore, some of the experiment tables can be merged into one which would improve the readibility. \n\n\n[1]: Yu, H., Liu, M., Luo, Y., Strasser, A., Qian, X., Qian, X., & Ji, S. (2024). Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules. Advances in Neural Information Processing Systems, 36.\n\n[2]: Khrabrov, K., Shenbin, I., Ryabov, A., Tsypin, A., Telepov, A., Alekseev, A., ... & Kadurin, A. (2022). nabladft: Large-scale conformational energy and hamiltonian prediction benchmark and dataset. Physical Chemistry Chemical Physics, 24(42), 25853-25863.\n\n[3]: Ruddigkeit, L., Van Deursen, R., Blum, L. C., & Reymond, J. L. (2012). Enumeration of 166 billion organic small molecules in the chemical universe database GDB-17. Journal of chemical information and modeling, 52(11), 2864-2875.\n\n[4]: Ramakrishnan, R., Dral, P. O., Rupp, M., & Von Lilienfeld, O. A. (2014). Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1), 1-7.", "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": "PPSqZXt4bn", "reviewer": "Reviewer_nQRm", "summary": "This paper introduces SLEM, a model for predicting quantum operator matrices that achieves state-of-the-art accuracy and improves computational efficiency. It is strictly local and uses a SO(2) convolution technique, enabling it to handle systems with heavy atoms and high angular momentums while maintaining efficiency. The work shows experimental validation across various materials, demonstrating better performance in both accuracy and computational cost compared to existing methods", "questions": "1) Are there specific cases where you expect the strict locality assumption to break down? The method perhaps could not handle systems with strong electronic correlations.\n2) It is not clear from the paper (at least I could not find it) which is the cutoff used, or if different cutoffs were used for different systems. This could enable comparing the receptive field of the model proposed in this paper and other approaches using message-passing.\n3) Evaluating the model performance in a transferability task would enhance significantly the submission.\n4) Evaluating the scaling of the model in a multigpu setting would also enhance significantly the submission.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main innovation and strength of this paper is combining strict locality, with an architecture heavily inspired in Allegro, with the SO(2) convolution trick, offering a compelling solution to handling heavy atoms and large systems efficiently while maintaining accuracy. The work can be considered well validated through comprehensive benchmarks, and they demonstrate the previous claims of both improved accuracy and reduced computational costs compared to state-of-the-art methods. The authors also create new datasets, and making them available would improve this paper's impact", "weaknesses": "A major limitation of the paper is its insufficient demonstration of transferability. While the authors show good performance on individual systems, they train and evaluate on the same type of material (e.g., training on Si and testing on Si configurations). There's no evaluation of cross material transferability. For instance, training on light elements and testing on heavy elements, or training on one crystal structure and testing on another. Also, the parallelization benefits, while promising (and I believe them, since they should be inherited from Allegro), are mainly theoretical with regards to this paper. The paper lacks concrete scaling studies on large systems or benchmarks on multiple GPU setups that would validate their claims about improved parallelizability.", "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)