Learning local equivariant representations for quantum operators

## Abstract

Abstract Predicting quantum operator matrices such as Hamiltonian, overlap, and density matrices in the density functional theory (DFT) framework is crucial for understanding material properties. 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, constructing local, equivariant representations for quantum tensors while preserving physical symmetries. This enables complex many-body dependence without expanding the effective receptive field, leading to superior data efficiency and transferability. Using an innovative SO(2) convolution technique, SLEM reduces the computational complexity of high-order tensor products and is therefore capable of handling systems requiring the f 𝑓 f italic_f and g 𝑔 g italic_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.

## 1.   Introduction

Quantum operators, representing observables and the evolution of quantum systems, are the cornerstone of describing the microscopic world. In modern quantum science, the advent of density functional theory (DFT) [, ] has elevated single-particle quantum operators, such as the Kohn-Sham Hamiltonian, density matrix, and overlap matrix, to paramount importance in solving complex problems []. These operators play a crucial role in unraveling electronic structures, predicting material properties, and advancing quantum technologies. However, as we tackle larger and more complex systems, the efficient and accurate representation of these fundamental operators has emerged as a pressing challenge in computation, demanding new avenues for innovative methodologies.

Recent advances have incorporated machine learning techniques to accelerate DFT calculations by directly predicting DFT’s output of quantum operators, including charge density [], overlapping matrix [, ], self-energy [], wave function [], and Hamiltonian matrix [, , , , , ].

Figure: Figure 1: Illustration of the local design of SLEM in comparison with conventional MPNN on a 1D graph. (a) MPNN information aggregation. (b) SLEM information aggregation. Balls and sticks represent nodes and edges; arrows show aggregation direction. $r_{\text{cut}}$: pre-defined cut-off; $r_{\text{eff}}$: effective cut-off after two hidden layer updates. $L$ is the layer index.
Refer to caption: x1.png

By circumventing self-consistent DFT calculations, such methods have the potential to scale up the electronic structure calculations. Some of the approaches utilize Gaussian regressions [], kernel-based models [], and neural networks [] to predict invariant Hamiltonian matrix blocks on localized frames.
Notably, powerful equivariant message-passing neural networks (E-MPNNs) have demonstrated remarkable accuracy [, , , , , , , , ]. These networks ensure the equivariance of output tensor blocks, respecting the physical priors of atomic systems. Typically, they initialize the node and edge features with atomic species and the weighted projection of the spherical harmonics coefficient, then use iterative two-body updates to build many-body interactions. Despite their high accuracy, these iterative updates invariably enlarge the receptive field of node and edge features, incorporating numerous effective neighbours. This expansion limits parallelization and, consequently, the model’s scalability. The storage-intensive nature of quantum tensor prediction tasks exacerbates these limitations, as high-dimensional features in edges and nodes necessitate substantial GPU memory during training and prediction. Consequently, training such models on large datasets or predicting quantum operators for extensive atomic structures poses considerable challenges.

Fortunately, the electrostatic screening counteracts the long-term dependency in a lot of material systems [, , ]. For example, for bulk Silicon, Ge, a typical screening length is around 5 a.u. []. Therefore, the predicted quantum operators can be decomposed into elements with localized dependency on atomic structures. Naturally, a strictly local model that avoids expanding the receptive field, while maintaining sufficient accuracy is at most preferred. In the field of machine-learning interatomic potentials (MLIPs), this concept has been applied in the Allegro model [], where the strictly local representation of interatomic potentials achieves high accuracy comparable to E-MPNNs and is highly parallelizable, particularly for large systems. However, unlike modelling interatomic potential, which only concerns scalars (energy) and vectors (forces) (corresponding to angular momentum $l=0$ and $1$) from node features, the prediction of the general representation of quantum operators necessitates targeting both node and edge features on high-order spherical tensors (even up to $l=6$ or $8$). This adds complexity to constructing strictly localized models for node and edge updates.
Another significant challenge in high-order tensor prediction tasks is computational complexity. To equivariantly mix the features of different angular momentum $l$, tensor products that scale as $\mathcal{O}$($l^{6}$) are required []. Achieving high accuracy often necessitates up to hundreds of thousands of training steps, which significantly limits the applicability of current models to heavy elements. This limitation hinders the development of a unified machine-learning DFT model that generalizes across the periodic table. Therefore, a crucial question, whether one can employ the strictly localized scheme to develop powerful equivariant atom and bond features that can efficiently and accurately reproduce the corresponding quantum tensor block elements with reduced computational complexity, naturally arises.

This work presents a novel method, the strictly localized equivariant message-passing model (SLEM), implemented in the framework of [DeePTB](https://github.com/deepmodeling/DeePTB) package [] for efficient representations of quantum operators. SLEM employs a fully localized scheme to construct high-order node and edge equivariant features for a general representation of quantum operators, including the Hamiltonian and density matrix. As illustrated in Fig.1, the model embeds localized edge hidden states and utilizes them to construct localized node and edge features without including distant neighbours beyond a prefixed cutoff range $r_{\text{cut}}$. This design enables the SLEM model to generalize more effectively, parallelize more easily, and scale to larger systems. Additionally, a fast and efficient SO(2) convolution [] is implemented to reduce the $\mathcal{O}$($l^{6}$) complexity, with edge-specific training weights, thereby further enhancing the model’s accuracy. As for the overlap matrix, it is typically required along with other operators for property calculations. Previous works have used another E-MPNN that doubles the size of the networks [] or extracts it from DFT calculations [, ], incurring additional costs or incorporating out-of-loop computation steps that complicate inference. In contrast, we utilized the two-center integrals nature and parameterized the overlap matrix with spherical independent scalars (i.e., Slater-Koster-like parameters []). This method, inspired by the early work of DeePTB [], fits the overlap operator representation alongside other operators with minimal additional cost.

This paper is organized as follows. In section 2, we review the related works, focusing on message-passing neural networks and equivariant message passing. In section 3, we detail the architecture of the SLEM model. Section 4 presents our experimental results, comparing the accuracy and efficiency of our model with existing models in reported data sets, and conducting generalization experiments. Finally, we summarize our findings and discuss future directions.

## 2.   Related works

##### Message-passing Neural Networks

Message-passing neural networks (MPNNs) [] have been widely applied in the modelling of atomic systems due to their exceptional accuracy in capturing the intricate relationships between atomic environments and physical properties [].
In the message-passing scheme, atoms are treated as nodes in graphs, with bonds to neighbouring atoms represented as connected edges within a specified cutoff radius.
The embedded atomic features are processed by trainable functions, generating messages from each edge to update the embeddings of central atoms. Formally, the MPNN framework can be summarized as follows:

$$ $\displaystyle\mathbf{e}^{ij,L}$ $\displaystyle=\mathcal{N}_{L}\left(\mathbf{n}^{i,L-1},\mathbf{n}^{j,L-1}, \mathbf{e}^{ij,L-1}\right)$ $\displaystyle\mathbf{m}^{ij,L}$ $\displaystyle=M_{L}\left(\mathbf{n}^{i,L-1},\mathbf{n}^{j,L-1},\mathbf{e}^{ij, L}\right)$ $\displaystyle\mathbf{n}^{i,L}$ $\displaystyle=U_{L}\left(\mathbf{n}^{i,L-1},\sum_{j\in\mathcal{N}(i)}\mathbf{m }^{ij,L}\right)$ $$

Here, $\mathbf{e}^{ij,L}$ represents the edge features, $\mathbf{m}^{ij,L}$ denotes the messages, and $\mathbf{n}^{i,L}$ indicates the node features at layer $L$. $\mathcal{N}_{L}$, $M_{L}$, and $U_{L}$ are the trainable functions for the edge, message, and node updates. $\mathcal{N}(i)=\{j|r_{ij}<r_{\text{cut}}\}$ indicates all the neighbour atoms for atom $i$, with $r_{\text{cut}}$ being the predefined cutoff.
This updating framework allows for the construction of many-body interactions and long-term dependencies, leading to strong performance across various applications.
Previous works have predominantly utilized this updating scheme, achieving remarkably high precision in reported systems [, , ]. However, as updates progress, the effective cutoff radius ($r_{\text{eff}}=N\times r_{\text{cut}}$) for each atom’s features grows linearly with the number of update steps ($N$), as shown in Fig.1. Consequently, the effective neighbour list scales cubically with $r_{\text{eff}}$, making parallelization intractable. This issue is particularly severe when fitting representations of quantum operators, as the $r_{\text{cut}}$ defined by the cutoff radius for the linear combination of atomic orbitals (LCAO) basis is often larger than that defined in other tasks.
A localized approach to constructing many-body interactions without including distant atoms has been proposed. For example, the Allegro model [] modifies the updating rules by incorporating a hidden atom-pair state that depends explicitly on the direction vector and the center atom embedding. While this framework works successfully for 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, instead creating manually designed descriptors of the local environment [, , , , , ]. These methods are generally local but are criticized for insufficiency in describing high-order interactions.

Figure: Figure 2: The design of the SLEM model. (a) the hierarchy structure of the model. Starts with the structure information such as atomic number $Z_{i}$, and the radial and spherical part of the shift vector $\bm{r}_{ij}$ and $\mathbf{Y}^{ij}_{c,l}$, the initialized hidden $\mathbf{x}^{ij}$, $\mathbf{V}^{ij}$, edge and node features $\mathbf{e}^{ij}$, $\mathbf{n}^{i}$ would be generated. The initialized two-body hidden features would be used to predict the Slater-Koster-like Overlap parameters directly and transformed to construct (off-)diagonal blocks for the overlap operator. Others then, will update iteratively using the design’s strict localized updating scheme. (b)-(d) shows the details of the hidden update (b), edge update (c), and node update (d). After the updates, the node and edge features would be used to construct the (off-)diagonal blocks using the Wigner-Eckart theorem for quantum operators such as density matrix and Hamiltonian.
Refer to caption: x2.png

##### Equivariant Message Passing

There are certain natures of the physical quantities that should be invariant or equivariant under e.g. the spatial and temporal operations. For example, under translation, rotation, and reflection, the scalars such as the energy of a system are invariant, while the tensors, such as forces and Hamiltonian matrices should transform accordingly. Such three kinds of operations constitute the group operations of E(3), known as the Euclidian Group (O(3) if without translation). To model such equivariant quantities, a set of neural network models has been developed that incorporate operations that preserve this equivariance. [, , , ]
These neural networks possess physical priors, as their operations ensure the output transforms in sync with the input, making them more generalizable, accurate, and data-efficient in predicting physical quantities. Formally, an equivariant operation from vector space X to vector space Y is defined such that:

$$ $\displaystyle f(D_{X}[g]x)=D_{Y}[g]f(x)\quad\forall g\in G,\forall x\in X$ $$

where $D_{X}[g]\in GL(X)$ is the representation of group element $g$ on vector space $X$. Since the representation takes the form of the direct sum of irreducible representations of the group $G$, it suggests that the vectors $x$ and $y$ themselves are composed by the irreducible representation (irreps for short) of group $G$. Usually, we consider here for operations of group O(3)(since translational symmetry is naturally satisfied by using relative coordinates), the irreps are the spherical tensors with angular momentum index $l$ and parity $p$. So features $x,y$ should also be indexed by $l,p$. For each irrep with angular momentum $l$, an internal degree of freedom, denoted by $m$ exists such that $|m|\leq l$. Therefore, each irreps of order $l$ belongs to a subspace of dimension $2l+1$. Once the $l,p$ index is defined, one can construct the group operations, which are represented by the $(2l+1)\times(2l+1)$ Wigner-D matrices.

Performing numerical operations on such irreps features follows certain rules. For the tensors with the same $l$, addition and subtraction can be conducted. To mix up the tensors of different $l$ indices, an important operation called tensor product ($\otimes$) is applied, which is considered the generalized operation of multiplication on spherical tensors. More formally:

$$ $\displaystyle(\mathbf{x}\otimes\mathbf{y})^{l_{3}}_{m_{3}}=\sum_{m_{1},m_{2}}C _{(l_{1},m_{1})(l_{2},m_{2})}^{(l_{3},m_{3})}\mathbf{x}^{l_{1}}_{m_{1}}\mathbf {y}^{l_{2}}_{m_{2}}$ $$

Where $C_{(l_{1},m_{2})(l_{2},m_{2})}^{(l_{3},m_{3})}$ are Clebsch-Gordan (CG) coefficients. This operation is found extremely useful in constructing edge and node updating functions, as seen in []. However, the conventional tensor product has the time and memory scales of $O(l_{max}^{6})$ [] where $l_{max}$ is the maximum angular momentum of $\mathbf{x}$ and $\mathbf{y}$. Such complexity limits the application like the equivariant forces fields where high inference speed is vital. It poses an even greater challenge for quantum tensor prediction tasks, which require tensor products on very high-order spherical tensors. For example, constructing blocks of $d$-$d$, $f$-$f$ and $g$-$g$ orbital pairs require irreps of maximum order $l=4$, $l=6$ and $l=8$, respectively. Such high costs make training for large-size systems nearly impossible.

## 3.   Model Architecure

### III.1 Parameterize Equivariant Quantum Operators

The quantum operators $\hat{O}$, such as the Hamiltonian and density matrix in the LCAO-based DFT framework, follow the equivariance of the $E(3)$ group and can thus be expressed in tensorized features using group theory. This process is illustrated in Fig. 2. Specifically, the $E(3)$ group consists of spatial rotation, translation, and reflection. Translational symmetry is automatically satisfied by performing operations only on relative coordinates, while rotation and reflection require special treatments. Under the LCAO basis, the quantum tensor block can be expressed as:

$$ $\displaystyle O_{l_{1},l_{2},m_{1},m_{2}}^{i,j}=\langle i,l_{1},m_{1}|\hat{O}| j,l_{2},m_{2}\rangle$ $$

Here $i$ and $j$ denote atomic sites, while $l_{1},m_{1}$, and $l_{2},m_{2}$ refer to the atomic orbitals of site $i$ and site $j$ respectively. The terms $m_{1}$ and $m_{2}$ are the magnetic quantum numbers associated with the angular momentum $l_{1}$ and $l_{2}$. $O_{l_{1},l_{2},m_{1},m_{2}}^{i,j}$ are the matrix element of operator $\hat{O}$.
When the coordinates of the system are transformed by a rotation operation $\mathbf{R}$, equivariance requires the quantum tensor to transform accordingly as:

$$ $\displaystyle O^{\prime i,j}_{l_{1},l_{2},m_{1}^{\prime},m_{2}^{\prime}}=\sum_ {m_{1},m_{2}}D(\mathbf{R})^{l_{1}}_{m_{1},m_{1}^{\prime}}D^{*}(\mathbf{R})^{l_ {2}}_{m_{2},m_{2}^{\prime}}O_{l_{1},l_{2},m_{1},m_{2}}^{i,j}$ (1) $$

Where $O^{\prime}$ is the equivariant block after transformation, and $D(\mathbf{R})^{l}_{m,m^{\prime}}$ is the Wigner-D matrix of order $l$. To construct such equivariance blocks from irreps, the Wigner-Eckart theorem can be applied, which decomposes the operator indexed by $l_{1},l_{2}$ into a single index $l_{3}$ that satisfies $|l_{1}-l_{2}|\leq l_{3}\leq(l_{1}+l_{2})$, which states:

$$ $\displaystyle o_{l_{3},m_{3}}^{i,j}=\sum_{l_{1},m_{1},l_{2},m_{2}}C_{(l_{1},m_ {1})(l_{2},m_{2})}^{(l_{3},m_{3})}O^{i,j}_{l_{1},l_{2},m_{1},m_{2}}$ (2) $$

Here, the edge ($i\neq j$) and node ($i=j$) features $o_{l_{3},m_{3}}^{i,j}$ are grouped by the index $m$ into vectors of $\mathbf{o}_{c,l}^{i,j}$ with $c$ accounting for multiple tensors for the same $l$. These features can be computed for hopping ($i\neq j$) and onsite ($i=j$) elements of the quantum operators. Further by leveraging the Hermitian nature of quantum operators, the parameterized elements can be further reduced to upper diagonal blocks.

The decomposed features $\mathbf{o}_{c,l}^{i,j}$ are aligned with the network prediction through a rescaling procedure. The target node/edge equivariant features $\mathbf{o}^{i,i/j}_{c,l}$ can be standarized and decomposed into constant norms and biases, and the normalized features that have balanced variance. Formally:

$$ $\displaystyle\mathbf{o}^{i,i/j}_{c,l}$ $\displaystyle=\sigma^{Z_{i},Z_{i/j}}_{c,l}\hat{\mathbf{o}}^{i,i/j}_{c,l}+\mu^{ Z_{i},Z_{i/j}}_{c,l}\delta_{l,0}$ $$

Here, $Z_{i}$, $Z_{j}$ are the atom species of atom $i$ and $j$, $\sigma^{Z_{i},Z_{i/j}}_{c,l}$ and $\mu^{Z_{i},Z_{i/j}}_{c,l}$ are norm and bias value for each atom and atom-pair. These values are derived from the dataset statistics and subsequently applied as weights and shift biases in an atom/bond type-specific scaling layer. The improvement of the standardization procedure on enhancing model convergence is evident. For the tensor operators in LCAO basis, the diagonal blocks (or node features) often exhibit highly unbalanced norms across irreps of different $l$, hindering the model from learning a smooth function. Additionally, different irreps in the off-diagonal blocks (or edge features) demonstrate diverse decaying behaviors, making it challenging to learn these decays implicitly. By normalizing in this manner, the decay function for each irrep is constrained to a value ranging from 1 to 0. This facilitates the use of hidden states and a few layers of a Multilayer Perceptron (MLP) activated by a ReLU-like function, to explicitly project the decaying value from the preserved hidden scalar features for each edge. Detailed analysis of the decaying behavior of edge irreps and the norm of node irreps can be found in the supplementary materials (SM).

At this stage, the objective is to generate the node and edge features $\mathbf{n}^{i}_{c,l}$, $\mathbf{e}^{ij}_{c,l}$, and to supervise them towards achieving decomposed target features $\hat{\mathbf{o}}^{i,i}_{c,l}$ and $\hat{\mathbf{o}}^{i,j}_{c,l}$. Once a satisfactory accuracy is attained, the inversion of the Wigner-Eckart theorem facilitates the reconstruction of predicted operator representations such as Hamiltonian and density matrix blocks:

$$ $\displaystyle O^{i,j}_{l_{1},l_{2},m_{1},m_{2}}=\sum_{l_{3},m_{3}}C_{(l_{1},m_ {1})(l_{2},m_{2})}^{(l_{3},m_{3})}o_{l_{3},m_{3}}^{i,j}$ $$

The whole procedure satisfies the rotational, transformational and reflectional symmetry.

### III.2 Parameterize of Invariant Overlap Operators

The overlap matrix, analogous to other quantum operators, is defined as:

$$ $\displaystyle S^{i,j}_{l_{1},l_{2},m_{1},m_{2}}$ $\displaystyle=\langle i,l_{1},m_{1}|j,l_{2},m_{2}\rangle$ $\displaystyle=\int{\phi^{i,l_{1}}_{m_{1}}(\mathbf{r})\phi^{j,l_{2}}_{m_{2}}}( \mathbf{r}-\mathbf{r}_{ij})d\mathbf{r}$ $$

where $\phi^{i,l}_{m}$ is the LCAO basis. The overlap matrix also satisfies the equivariance relation defined in Eq.1. Since the equivariant dependency comes fully from the bases, it is possible to rotate them to align with the axis of $\mathbf{r}_{ij}$ to reduce the angular dependence, simplifying the matrix elements further into scalars via the relation []:

$$ $\displaystyle\int{\phi^{i,l_{1}}_{m_{1}}(\mathbf{r})\phi^{j,l_{2}}_{m_{2}}}( \mathbf{r}-\mathbf{r}_{ij})d\mathbf{r}=s^{Z_{i},Z_{j}}_{l_{1},l_{2},|m_{1}|}({ r}_{ij})\delta_{m_{1},m_{2}}$ $$

Here, the dependence changes to $Z_{i}$, $Z_{j}$, and $r_{ij}$ due to the two-center nature of the overlap integrals. Parameterizing these distance-dependent radial functions involves encoding atomic species with radial information and employing a simple MLP to learn the mapping to target scalars.

$$ $\displaystyle f_{l_{1},l_{2},|m|}(Z_{i},Z_{j},r_{ij})=s^{Z_{i},Z_{j}}_{l_{1},l _{2},|m_{1}|}(r_{ij})\delta_{m_{1},m_{2}}$ (3) $$

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.

### III.3 The SLEM model

The SLEM model architecture is illustrated in Fig.2. The model maintains a set of features, including hidden features $\mathbf{x}^{ij,L},\mathbf{V}^{ij,L}_{c,l}$, node features $\mathbf{n}^{i,L}_{c,l}$ and edge features $\mathbf{e}^{ij,L}_{c,l}$ of hidden layer ($L$). Specifically, the hidden features consist of a scalar channel $\mathbf{x}^{ij,L}$ and a tensor channel $\mathbf{V}^{ij,L}_{c,l}$. The scalar channel is initialized as 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 processed through an MLP to the invariant Slater-Koster-like parameters of overlap. For equivariant quantum operators, such as the Hamiltonian and density matrix, the scalar and tensor channels interact to generate an ordered atom-pair representation, which is 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.

#### III.3.1 Feature initialization

Firstly, the initial scalar hidden feature is computed from the two-body embeddings of atomic species $Z_{i}$ and $Z_{j}$, and the radial distance, as follows:

$$ $\displaystyle\mathbf{x}^{ij,L=0}=\textbf{MLP}_{\text{2-body}}\left(\mathbf{1}( Z_{i})||\mathbf{1}(Z_{j})||\textbf{B}(r_{ij})\right)\cdot u(r_{ij})$ $$

Here, atom species are embedded with one-hot vectors denoted as $\mathbf{1}(Z)$, and a set of trainable Bessel bases $\textbf{B}(r_{ij})$ is utilized to encode the distance $r_{ij}$ between atoms $i$ and $j$. $u(r_{ij})$ are envelope functions [] to add explicit radial dependence. Subsequently, the edge and hidden features are initialized as weighted spherical harmonics of relative edge vectors:

$$ $\displaystyle\mathbf{V}^{ij,L=0}_{c,l}$ $\displaystyle=w_{c,l}\left(L_{N}\left(\mathbf{x}^{ij,L=0}\right)\right)\mathbf {Y}^{ij}_{l}$ $\displaystyle\mathbf{e}^{ij,L=0}_{c,l}$ $\displaystyle=w_{c,l}\left(\mathbf{x}^{ij,L=0}\right)\mathbf{Y}^{ij}_{l}$ $$

Here, the weights are learned from the initialized scalar hidden features $\mathbf{x}^{ij,L=0}$. Layer normalization $L_{N}$ ensures that the hidden tensor features have a balanced amplitude of each edge.
The initial node features are then computed as linear transformations of the aggregated edge features:

$$ $\displaystyle\mathbf{n}_{c,l}^{i,L=0}=\textbf{Linear}\left(\frac{1}{\sqrt{N_{ avg}}}\sum_{j\in\mathcal{N}(i)}\mathbf{e}^{ij,L=0}_{c,l}\right)$ $$

Here $\mathcal{N}(i)$ and $N_{avg}$ are the neighbouring atoms and the average number of neighbours of atom$-i$.

#### III.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 the concatenated equivariant features $\tilde{\mathbf{f}}^{ij}_{c,l}$ and the weighted projection of the edge shift vector ${\bm{r}}_{ij}={\bm{r}}_{i}-{\bm{r}}_{j}$ on the spherical harmonics function $\mathbf{Y}_{l}^{ij}$. Formally:

$$ $\begin{split}\mathbf{f}^{ij}_{c_{3},l_{3}}&=\tilde{\mathbf{f}}^{ij}_{c_{1},l_{ 1}}\otimes w^{ij}_{c_{2},l_{2}}\mathbf{Y}_{l_{2}}^{ij}\\ &=\sum_{c_{1},l_{1},l_{2}}\tilde{w}_{c_{1},l_{1},l_{2}}^{ij}\sum_{m_{1},m_{2}} C_{(l_{1},m_{1})(l_{2},m_{2})}^{(l_{3},m_{3})}\\ &\cdot f^{ij}_{c_{1},l_{1},m_{1}}Y_{l_{2},m_{2}}^{ij}\\ \end{split}$ (4) $$

Here, $\tilde{w}_{c_{1},l_{1},l_{2}}^{ij}=\sum_{c_{2}}w_{c_{1},c_{2},l_{1},l_{2}}w^{
ij}_{c_{2},l_{2}}$ are edge-specific parameters for each tensor product operation. Performing such tensor products on high-order features is computationally intensive. Therefore, we applied the recently developed SO(2) [] convolution to simplify, reducing the computation and storage complexity from $O(l_{max}^{6})$ to $O(l_{max}^{3})$.
The simplification idea is intuitive. $\mathbf{Y}_{l_{2},m_{2}}^{ij}$ are sparse tensors if rotated to align with the edge $ij$, which is nonzero only for $m_{2}=0$. Therefore, it is easier to compute the tensor production in the direction of edge $ij$, and rotate inversely the output afterwards. This step removes the $m_{2}$ index from the summation in Eq. LABEL:tp. Additionally, considering Clebsch-Gordan coefficients with $m_{2}=0$, we find that $C_{(l_{1},m_{1})(l_{2},0)}^{(l_{3},m_{3})}=0$ except for $m_{3}=\pm m_{1}$. This allows further reduction of the summation in Eq. LABEL:tp by replacing $\pm m_{1}$ with a single index $m$. Then the operations can be reformulated formally as:

$$ $\left(\begin{array}[]{l}f^{ij}_{c,l,m}\\ f^{ij}_{c,l,-m}\end{array}\right)=\sum_{c^{\prime},l^{\prime}}\left(\begin{ array}[]{l}w_{c,c^{\prime},l^{\prime},m}-w_{c,c^{\prime},l^{\prime},-m}\\ w_{c,c^{\prime},l^{\prime},-m}\quad w_{c,c^{\prime},l^{\prime},m}\end{array} \right)\cdot\left(\begin{array}[]{l}{\tilde{f}^{ij}_{c^{\prime},l^{\prime},m}} \\ {\tilde{f}^{ij}_{c^{\prime},l^{\prime},-m}}\end{array}\right)$ $$

This represents a linear operation on $\tilde{f}^{ij}_{c^{\prime},l^{\prime},m^{\prime}}$. By employing this method, high-order tensor products for $l=8$, $9$ and even $10$ can be 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 $w^{ij}_{c^{\prime},l^{\prime}}$ are mapped by an MLP from hidden scalar features $\mathbf{x}^{ij,L}$ as $\tilde{w}^{ij}_{c,c^{\prime},l^{\prime},m}=w_{c,c^{\prime},l^{\prime},m}w^{ij}
_{c^{\prime},l^{\prime}}$. This powerful and efficient tensor product layer facilitates the construction of local interactive updates of the features.

#### III.3.3 Hidden updates

To construct many-body interactions, as in Fig. 2(b), the node features $\mathbf{n}^{i}_{c,l}$ and hidden tensor features $\mathbf{V}^{ij}_{c,l}$ would be concatenated and doing tensor product with the projection coefficients of edge shift vector $\bm{r}_{ij}$ on the spherical harmonics functions. The operation is written formally as:

$$ $\tilde{\mathbf{V}}^{ij,L}_{c_{3},l_{3}}=\left(\mathbf{n}^{i,L-1}||\mathbf{V}^{ ij,L-1}\right)_{c_{1},l_{1}}\otimes w_{c_{2},l_{2}}^{ij,L}\mathbf{Y}_{l_{2}}^{ij}$ (5) $$

Unlike most MPNN, the hidden states $\mathbf{x}^{ij}$ and $\mathbf{V}^{ij}_{c,l}$ in SLEM is only a function of the local environment of centre atom $i$, the shift vector of edge $ij$, and the information of atom $j$. Such a design excludes neighbours of $j$ into hidden states.

After the tensor production, the output features $\tilde{\mathbf{V}}^{ij,L}_{c,l}$ will be passed through the gated non-linearity [], and transformed by an E3linear [] 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 include the radial information explicitly and return as the updated feature $\mathbf{V}^{ij,L}_{c,l}$.

The scalar hidden features are updated by mixing the $0\text{th}$ order information from $\mathbf{V}^{ij,L}_{c,l}$ with a latent MLP, which is:

$$ $\mathbf{x}^{ij,L}=\textbf{MLP}\left(\mathbf{x}^{ij,L-1}||\mathbf{V}^{ij,L}_{c, l=0}\right)\cdot u(r_{ij})$ (6) $$

Containing the explicit decaying envelope function $u(r_{ij})$, and the many-body interactions of scalar and tensor features, the sclar hiddien states $\mathbf{x}^{ij,L}$ can be very useful to describing the explicit decaying behaviour of each edge irreps feature along the radial distance.

#### III.3.4 Node updates

The strict local node representation $\mathbf{n}^{ij,L}_{c,l}$ can be constructed naturally from the many-body interactive tensor features $\mathbf{V}^{ij,L}_{c,l}$. We follows the MPNN style to create the message from node-$j$ to node-$i$. Formally:

$$ $\mathbf{m}^{ij,L}_{c_{3},l_{3}}=\left(\mathbf{n}^{i,L-1}||\mathbf{V}^{ij,L} \right)_{c_{1},l_{1}}\otimes w_{c_{2},l_{2}}^{ij,L}\mathbf{Y}_{l_{2}}^{ij}$ (7) $$

Here again, we exclude the neighbouring information of atom-$j$ in $\mathbf{m}^{ij,L}_{c_{3},l_{3}}$ via this partial updates of $\mathbf{V}^{ij,L}_{c,l}$. However, all necessary interactions are contained in this format, including the interactions of atom-$j$, bond $i-j$, atom-$i$, and the environment of $i$.

Each message then is passed through a gated activation and E3Linear layer, and weighted separately by weights learnt from the hidden scalar features, by:

$$ $\mathbf{n}^{ij,L}_{c_{3},l_{3}}=\frac{1}{N_{avg}}\sum_{j\in\mathcal{N}(i)}w^{ ij}_{c_{3},l_{3}}\mathbf{m}^{ij,L}_{c_{3},l_{3}}$ (8) $$

The weights here are different from those in hidden updates since it is learnt from $\mathbf{x}^{ij,L}$ directly without passing through a normalization. Therefore, the absolute radial decay is enforced in the weights, giving strong priors that the message from a closer atom is often more important while maintaining flexibility. Meanwhile,
combining the updates of hidden scalar $\mathbf{x}^{ij,L}$, we can see the weights $w^{ij}_{c_{3},l_{3}}$ are dependent on the features $\mathbf{V}^{ij,L}$ and $\mathbf{n}^{i,L-1}$, as shown in Eq.6. Such style aligns surprisingly with the equivariant graph attention mechanism [, ] that shows powerful expressibility in various tasks. Here ${w}^{ij}_{c_{3},l_{3}}$ corresponding to an attention score computed from $\mathbf{V}^{ij,L}$ and $\mathbf{n}^{i,L-1}$. Therefore, through this update, the node features’ dependencies are kept strictly local.

#### III.3.5 Edge updates

The locality of edge features $\mathbf{e}^{ij,L}_{c,l}$ can be naturally preserved as long as the node features are localized strictly. By mixing the information of node features on both sides, localized updates can be formulated as:

$$ $\displaystyle\tilde{\mathbf{e}}^{ij,L}_{c_{3},l_{3}}=\left(\mathbf{n}^{i,L-1}| |\mathbf{V}^{ij,L}||\mathbf{n}^{j,L-1}\right)_{c_{1},l_{1}}\otimes w_{c_{2},l_ {2}}^{ij,L}\mathbf{Y}_{l_{2}}^{ij}$ $$

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:

$$ $\displaystyle\mathbf{e}^{ij,L}_{c_{3},l_{3}}=w^{ij}_{c_{3},l_{3}}\tilde{ \mathbf{e}}^{ij,L}_{c_{3},l_{3}}$ $$

Here we reframe the updating rules of SLEM with the MPNN framework, which looks like this:

$$ $\displaystyle\mathbf{V}^{ij,L}$ $\displaystyle=\mathcal{V}_{L}\left(\mathbf{n}^{i,L-1},\mathbf{V}^{ij,L-1}\right)$ $\displaystyle\mathbf{e}^{ij,L}$ $\displaystyle=\mathcal{N}_{L}\left(\mathbf{n}^{i,L-1},\mathbf{V}^{ij,L}, \mathbf{n}^{j,L-1}\right)$ $\displaystyle\mathbf{m}^{ij,L}$ $\displaystyle=M_{L}\left(\mathbf{n}^{i,L-1},\mathbf{V}^{ij,L}\right)$ $\displaystyle\mathbf{n}^{i,L}$ $\displaystyle=U_{L}\left(\mathbf{n}^{i,L-1},\sum_{j\in\mathcal{N}(i)}\mathbf{m }^{ij,L}\right)$ $$

Where $\mathcal{V}_{L}$ 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 a scheme, in principle, would have much better transferability, data-efficiency, as well as scalability. Most importantly, it is possible to parallelize such update computations, allowing for training and inference on multiple devices. This is also vital for inferencing on very large atomic structures. The following section will focus on validating the effectiveness of this framework via learning equivariant DFT Hamiltonians and density matrices and overlap, which have very high-order tensors and need extremely high prediction accuracy. We demonstrate our method is much better than the previously reported ones in training and inference speed, transferability and the potential power to parralize the inference on multi-GPU devices.

**Table 1: MAE (in meV) for Hamiltonian matrix predictions using SLEM and other methods on materials with LCAO basis up to $d$ orbitals. Numbers in parentheses indicate parameter count.**
| Systems with LCAO-basis up to $d$-orbitals |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Material | SLEM | DeepH-E3 | HamGNN |  |  |
| (0.7M) | (4.5M) | ( 1.0 M) | ( 4.5M) | (2.8M) |  |
| MoS2 | 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 meV) for Hamiltonian matrix predictions using SLEM and DeepH-E3 models [] on materials with LCAO basis up to $f$ and $g$ orbitals. Numbers in parentheses indicate parameter count. The MAE of HfO2 in DeepH-E3 is absent due to out-of-memory errors.**
| Systems with LCAO-basis up to $f$ and $g$-orbitals |  |  |
| --- | --- | --- |
| Material | SLEM | DeepH-E3 |
| (1.7M) | (1.9M) |  |
| GaN | 0.21 | 0.87 |
| HfO2 | 0.28 | - |

**Table 3: MAE for predicting density matrix using SLEM on materials with LCAO basis up to $d$, $f$, and $g$ orbitals. Model settings align with those in Table LABEL:hamiltonian.**
| SLEM charge density model |  |  |  |
| --- | --- | --- | --- |
| Materials | Silicon | GaN | HfO2 |
| MAE | 8.9e-5 | 2.3e-5 | 3.9e-5 |

## 4.   Results

### IV.1 Benchmark the accuracy and data-efficiency

Figure: Figure 3: Comparison of band structures for a 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.
Refer to caption: extracted/5735747/band.png

We benchmark our model’s performance in fitting Hamiltonian, density matrix, and overlap matrix using a diverse dataset. For the Hamiltonian, the datasets include systems with up to $d$ orbitals, encompassing 2D systems like single-layer MoS2 and graphene from reported datasets [], as well as 3D bulk silicon generated in this work. Given SLEM’s proficiency with high-order spherical tensors, we also train it on our generated datasets of bulk GaN and HfO2 systems, which include $f$ and $g$ orbitals, respectively. Structures in the Si, GaN, and HfO2 systems are sampled via molecular dynamics using neural network potentials []. For density matrix and overlap matrix evaluations, we test the SLEM model exclusively on the datasets of Si, GaN, and HfO2, as the reported datasets for MoS2 and graphene lack density and overlap matrix data.

**Table 4: MAE for predicting overlap matrix using SLEM’s parameterization on materials with LCAO basis up to $d$, $f$, and $g$ orbitals. Model settings align with those in Table LABEL:hamiltonian.**
| SLEM overlap prediction |  |  |  |
| --- | --- | --- | --- |
| Materials | Silicon | GaN | HfO2 |
| MAE | 5.6e-5 | 4.7e-5 | 6.3e-5 |

Table LABEL:hamiltonian presents a comparison of mean absolute error (MAE) values in Hamiltonian prediction for graphene, MoS2, and Si systems, whose LCAO basis extends up to $d$ orbitals, among the SLEM, DeepH-E3 [], and HamGNN [] methods. Our SLEM model achieves state-of-the-art accuracy, exhibiting the lowest MAE across all systems. Notably, this high accuracy is achieved with a 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 HfO2 systems, where the LCAO basis extends up to $f$ and $g$ orbitals, respectively. As shown in Table LABEL:hamiltonian_fg, the SLEM model consistently presents the lowest MAE, further demonstrating its high accuracy and versatility across different orbital complexities.
Fig. 3 illustrates the band structure of silicon structures computed directly from the SLEM-predicted Hamiltonian, where the eigenvalues are indistinguishable from those obtained using DFT.
For the density matrix, fitting results are presented in Table LABEL:density. The results demonstrate very high accuracy (order of 1e-5), approaching the machine precision limit of 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 for applications such as charge distribution analysis or tracking electron transfer via methods like differential charge density. Furthermore, the predicted density can be directly used for non-self-consistent field (NSCF) DFT calculations. The resulting band structure for silicon, as an example, is highly accurate and matches the DFT output, with a MAE of only 1.09 meV in eigenvalues compared to self-consistent DFT results.
The overlap matrix, represented by invariant parameters in our SLEM model, achieves exceptionally high accuracy as demonstrated in Table LABEL:overlap, 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 to fit only the Hamiltonian, our typical parameter count is 0.7 M. The inclusion of overlap matrix prediction adds merely 0.01 M parameters, representing an increase of only about 1.4% in total model size.

**Table 5: Comparison of validation MAE (in meV) for SLEM model and DeepH-E3 [] method trained on randomly split datasets [] with varying training ratios. Model settings align with those in Table LABEL:hamiltonian.**
| MoS2 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Partition | 100% | 80% | 60% | 40% | 20% |
| SLEM | 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 |
| DeePH-E3 | 0.40 | 0.30 | 0.33 | 0.36 | 0.60 |

Figure: Figure 4: Comparison of time and memory consumption for different tensor-product implementations. (a) Time consumption vs. angular momentum ($l$) for different models, including the SO(2)-based SLEM model (triangles) with and without radial part ($r$), DeepH-E3 (cross) [], and E3NN (square) [] models. Inset: Log-scale fit with slopes of 1.2 for the SLEM model and 3.7 for the other two models. (b) Memory consumption vs. $l$. The SLEM model demonstrates over two orders of magnitude improvement in both time and memory efficiency compared to DeepH-E3 and E3NN.
Refer to caption: x3.png

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 training data, comprising 20%, 40%, 60%, and 80% of the full dataset. We trained both the SLEM model and DeepH-E3 method on these subsets and evaluated their performance on consistent validation sets. Table LABEL:data-efficiency presents the results, demonstrating SLEM’s maintained high accuracy across all subsets, thus highlighting its remarkable data efficiency.
This data efficiency implies that SLEM users can generate smaller, more cost-effective training sets. Moreover, SLEM’s excellent data efficiency and transferability make it particularly well-suited as a backbone for developing universal DFT models, especially for systems involving heavy elements.

### IV.2 Efficiency and Scalability

Many interesting and important properties in material science, chemistry and biology appear with systems composed of heavy atoms. The inclusion of heavy atoms would introduce very high-order spherical tensors when representing quantum operators such as the DFT Hamiltonian. Therefore, scaling to such systems would be quite a challenge because the tensor product used to construct complex spherical tensors scales as $O(L^{6})$. It is very hard to train and inference on systems with heavy atoms using conventional tensor production.
Moreover, inferring large material systems while training with small structures is particularly valuable. It therefore requires parallelising the model inference by assigning partitions of the whole atomic structure to multiple GPU workers. However, this is very difficult for most of the currently available models. Since the increasing receptive fields from the iterative updates of the graph would increase the minimal size of each partitioned subgraph, which makes such partition less effective. The SLEM model addresses these challenges by efficiently constructing high-order tensor products and assisting parallelization through its strict locality.

Firstly, for efficiency, the implementation of SO(2) convolution reduces the tensor product computational complexity from $O(L^{6})$ to $O(L^{3})$, which is then further reduced by the parallelization of matrix operations to near $O(L)$ benefiting from PyTorch. In Fig. 4, we compare the wall time and GPU memory consumed by tensor product operations using SO(2) convolution and the normal method employed in DeepH-E3 [] and E3NN []. Clearly, tensor production via SO(2) convolution is far more efficient than conventional methods. This reduction of cost makes our method very efficient, which is capable of handling all possible basis choices in LCAO DFT. We also record the memory and time required when training on some of the typical systems and compare with that of DeepH-E3 []. The results are displayed in Fig. 5. We can see, that in all systems, our model trains faster in time while requiring less memory. Notably, the larger the basis set, the greater the advantage our method has in both memory and time efficiency.

Figure: Figure 5: Comparison of training time per iteration and memory consumption with SLEM and DeepH-E3 [] models.
Refer to caption: x4.png

Additionally, the strict local design assists the parallelization greatly. Since the edge and node features in the SLEM model are constructed locally, each node and edge feature is only dependent on the neighbour subgraph of the whole system, as:

$$ $\displaystyle\mathbf{n}^{i,L}$ $\displaystyle=U_{L}\left(\mathbf{n}^{i,L-1},\sum_{j\in\mathcal{N}(i)}\mathbf{m }^{i,L}\right)$ (9) $\displaystyle=U(Z_{i},\{\mathbf{Y}^{ij},Z_{j}\}_{j\in N(i)})$ $$

Therefore, it is possible to divide the whole atomic system into sub-graphs and compute their node and edge features independently on different devices. This is extremely important when expanding the simulation power of DFT to very large systems. In practice, we find that for HfO2 with $4s2p2d2f1g$ basis, a typical model of 1.7M parameters can only predict the quantum operators of structures up to $10^{3}$ atoms on devices with 32GB memory. Despite the linear dependency, the inference on a system with $10^{4}\sim 10^{7}$ would require a memory with potentially over 300 GB, which is only available when parallelised to multiple GPUs. Therefore, a strictly localized model such as SLEM holds significant potential for expanding simulation system sizes.

## Discussion

This work proposed SLEM, the strictly local equalvariant message-passing model for the prediction of quantum operator representations, as well as a systematic view of the parameterization of equivariant quantum operators. The strict local design of SLEM builds many-body dependence without expanding the effective receptive field, contributing to excellent accuracy, data efficiency and transferability. The integration of SO(2) convolution [] enables rapid and memory-efficient tensor production, which is crucial for modeling heavy atoms. Furthermore, SLEM’s locality potentially allows for extension to extremely large atomic systems through parallelization over partitioned atomic graphs, thus enabling first-principles electronic structure simulations at device-level scales. Beyond its modeling capabilities, SLEM intrinsically supports the representation of multiple quantum operators, including the Hamiltonian, density matrix, and overlap matrix. By parameterizing the overlap matrix using a two-center integration formula, it is transformed into invariant quantities readily obtainable from SLEM’s initialized two-body embedding features. This approach significantly reduces computational costs and minimizes the model’s reliance on DFT software post-training.

While SLEM demonstrates efficacy across various applications, several areas warrant further investigation:

- •
Sampling Methodology for Active Learning: Scientific computation tasks require high confidence in calculated results, which can be challenging for data-driven approaches. Developing a robust sampling workflow for active learning is essential for reliability. While techniques like uncertainty-driven sampling with Gaussian regression [] or model ensembles [] have addressed confidence issues in machine learning force fields, designing a 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 conventional DFT capabilities. An efficient, parallelizable solver for extracting physical quantities from the model’s predictions is highly beneficial. While some software based on stochastic techniques  [, ] shows promise, an open-sourced and highly-optimized solution for 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.

## 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](https://github.com/deepmodeling/DeePTB). To facilitate the integration of DFT outputs with machine learning models, we have developed and released a supplementary tool, [dftio](https://github.com/floatingCatty/dftio). This tool enables parallel parsing of DFT outputs into a machine-readable data format, enhancing the efficiency of data preparation for model training and analysis.

###### Acknowledgements.

###### Acknowledgements.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: My questions are stated in the above concerns, I would be happy to raise my score if the authors address them.
soundness: 4
strengths: - The two key contributions are novel, they are clearly explained in the paper and based on my understanding they are technically sound.
- 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.
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:

- 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.
- The parameterization of the invariant overlap operators, which enables the prediction of overlap operator without requiring lots of parameters.

The above innovations is verified to be effective in the empirical evaluations on various systems. The advantages reported include 

- Better generalization due to the more restricted model assumption.
- Better scaling behavior w.r.t angular momentum.
- Faster iteration speed and smaller memory footprint.
weaknesses: I have a major concern on the way the evaluation is carried out

- 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`.

Another concern is on the theoretical soundness

- 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?
- 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.
- 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.

---

confidence: 5
contribution: 3
initial_rating: 3
presentation: 3
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?

2) 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?

3) 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?

4) 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.
soundness: 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.

The 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. 

The 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.
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.
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.

Moreover, 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.

In 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. 


[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.

[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.

[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.

[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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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.
2) 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.
3) Evaluating the model performance in a transferability task would enhance significantly the submission.
4) Evaluating the scaling of the model in a multigpu setting would also enhance significantly the submission.
soundness: 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
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
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.

Meta-review: The paper proposes an equivariant neural network that predicts Hamiltonian (converged Fock matrix in DFT) for molecular systems. The authors highlighted the benefit of reduced computational and memory cost by imposing a strictly local design to the architecture, together with other tricks e.g. a proper processing of the overlap matrix, and the usage of SO(2) convolution for lighter cost in tensor product.

I agree with reviewers' appreciation that this work contributes to the bottlenecking issue of efficiency in this domain, and even showing better generalization. Some concerns are raised by the reviewers, including the rationale of the strictly local assumption, generalization results over chemical space, choices of baselines and metrics and dataset details. In the rebuttal, the authors provided a reasonable argument for the locality at least for periodic systems, further results on QH9 that can be accounted for chemical space generalization evaluation, and explanations on baselines and metrics and enriches details in curating the dataset. Correspondingly, two reviewers have raised their scores.

In addition, I also asked two questions regarding the SOTA claim (considering e.g. PhiSNet) and the novelty of using SO(2) convolution (considering MLFFs that use this). The reply says that the proposed method mainly targets at periodic systems while PhiSNet at molecular systems, and that the proposed method outperforms QHNet on QH9 which has a similar performance as PhiSNet, and that the contributions of strictly local design and parameterization of the overlap matrix are more central contributions. This could be an acceptable answer, but I hope the authors could provide more details on the experimental setup and protocols to consolidate that the improvement is not from extensive tuning, and reasoning the improvement on QH9 given that the authors emphasized the beneficial cases are mainly periodic ones due to better satisfaction of quantum screening.

In respect to reviewers' preference, I recommend an accept to this submission. But I hope the authors could further improve the paper accordingly for the final version, e.g., the more comprehensive explanations on the locality assumption, their own dataset details, and configurations on the QH9 dataset, apple-to-apple comparison with other baselines, and interpretation of the better generalization in the non-periodic case.

Decision: Accept (Spotlight)