Learning local equivariant representations for quantum operators

We appreciate the reviewer's kind comments about our work. We especially, thank the recognition for the experiment that is well validated. Your kind applause encourages us. We are also very happy to see the construction suggestions including the data openness, the transferability test and the multi-GPU inference. These suggestions help improve our manuscript and provide guidance for further research. We have uploaded all the data used in this work in the “AIS Square” with the link of https://www.aissquare.com/datasets/detail?pageType=datasets&name=Quantum_Operator_Dataset&id=286 . We hope this can increase the reproducibility. In the following discussion, we will address the questions and suggestions listed one by one.

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

The theoretical framework of DFT fundamentally supports the strict locality assumption in our method. 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 and electronic screening effects in condensed matter systems. The only difference lies in how large the local spread is. For periodic systems, particularly metals and semiconductors, electronic screening effectively reduces long-range Coulomb interactions. The screening length varies by material type around 4 Å in semiconductors like silicon, and Germanium, and smaller in insulators such as 2.76 Å in Diamond Therefore, we are safe to choose a cutoff radius to match the sum of basis orbital radii, adequately captures these screening effects. Regarding strongly correlated systems, it's important to distinguish between method limitations and framework limitations. While these systems present challenges, the limiting factor is the accuracy of the underlying DFT calculation rather than our locality assumption. Our method can still effectively learn the DFT mapping, inheriting both its capabilities and limitations. For molecular systems where screening effects are weaker, we acknowledge the need for larger cutoff radii. This motivated our development of the LEM variant, which introduces trainable but exponentially decaying environmental dependencies to handle the weaker screening effects in molecular systems. Our benchmark results on the QH9 dataset demonstrate the effectiveness of this adaptation:

|dataset: QH9-Stable-size-ood| unit (1e-6 Ha) |QH-net(reported) | QH-net(reproduced) | LEM(Ours) | |75.0 | 83.2 | 57.7 |

More formal results will be added to the appendix of the manuscript after the conclusive test shortly. Therefore, rather than "breaking down," our method maintains validity within the DFT framework across different physical regimes, with adjustable parameters to accommodate varying degrees of locality.

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

We apologize for not clearly specifying the cutoff selection. We will clarify this here, and add a description in the experiment section, in section 4.1 of the revised manuscript.

Typically, we set the cutoff radius being the sum of two atomic orbital cutoffs used in DFT calculations for training labels. Specifically, for numerical atomic orbitals (used in SIESTA, ABACUS, OpenMX), The cutoff is explicitly defined in orbital files or generation inputs, typically 5-10 a.u. For Gaussian-type orbitals, The cutoff is determined during data processing as the distance beyond which matrix elements become negligible.

The reason for this setting is that in a quantum operator matrix learning task, the cutoff radius needs to cover all significant matrix elements on the LCAO basis. For matrix elements between two atomic orbitals, their magnitude becomes negligible when the interatomic distance exceeds the sum of the two orbital cutoff radii. Therefore, our graph construction naturally aligns with this physical constraint by connecting atoms within this cutoff distance. A key advantage of our strict localized scheme is that the effective cutoff remains constant at this physical value, unlike MPNN-based methods (e.g., DeePH-E3, HamGNN) where it increases by a factor of L (number of network layers) due to message passing.

We thank the reviewer for the suggestions and questions. Some suggestions improve reliability and provide direct guidance for improvement. For example, we fully understand your concerns about the soundness of the experiment, data generation processes, and the comparison of open benchmarks. To address these concerns, we added detailed data generation processes and uploaded all our datasets for reproducibility.

On the other hand, we want to highlight that quantum operator learning for periodic material systems is still an emerging field without standardized open benchmarks. We've done our best to increase the soundness of our work by using data from multiple sources and comparing it with the available methods. Below, we will carefully address the comments point by point and respectfully request the reviewer reconsider the marks, considering the novelty and contribution of this work.

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

Response: For the model selection, we use DeepH-E3 and HamGNN since they are the most representative (if not the only available) equivariant models in this area, that support Hamiltonian prediction task on periodic materials. While there exist other models that also predict the Hamiltonians, they focus on the molecular system and do not support periodic materials that our SLEM model specialises for. Therefore, our comparison with these two models represents the most comprehensive evaluation possible within the current landscape of the field.

For the dataset selection, we share the reviewer's interest in evaluating recognized open benchmark datasets. However, unlike in molecules, widely recognised benchmark datasets such as QH-9 and nabla-DFT for Hamiltonian/Density/Overlap operators on periodic materials are unavailable until today. Therefore, to ensure comprehensive evaluation, we used a combination of public and in-house datasets:

• For 2D materials: We utilized the public graphene and MoS2 datasets from DeepH-E3 model [X. Gong, et. al, Nat Commun 14, 2848 (2023)]
• For 3D materials: Due to the lack of public datasets, we generated our data using ABACUS DFT software [P. Li, et. al, Comput. Mater. Sci. 112, 503 (2016)], which is openly available and uploaded on https://www.aissquare.com/datasets/detail?pageType=datasets&name=Quantum_Operator_Dataset&id=286.

Our model demonstrated a consistent advantage in accuracy in all the above-mentioned datasets. We believe this proves the advance of our method. While SLEM primarily targets periodic systems, as suggested by the reviewer, we also evaluated its performance on molecular datasets like QH-9. For these smaller systems where all atom pairs fall within typical cutoffs (resulting in fully connected Hamiltonians), we implemented a semi-local variant called LEM (Localized Equivariant Message-passing). Our comparative analysis with QH-net on the QH-9 benchmark (will be updated to the appendix shortly) demonstrates excellent accuracy, further validating our method's versatility across both periodic and molecular systems.

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

Response: We would like to clarify that E3NN's performance metrics (MAE) are not included in our accuracy comparison tables because E3NN is a general-purpose equivariant neural network library rather than a specific model for quantum operator prediction. E3NN does not provide ready-to-use implementations for quantum operator learning tasks, making a direct accuracy comparison impossible.

The E3NN comparison in our computational scalability analysis serves a different purpose - comparing tensor-product implementations. The standard tensor product implementation in the E3NN package, while widely used, becomes computationally intensive when handling materials with higher orbitals. DeepH-E3 improved upon this by using a low-rank approximation for the weight matrix, and our SLEM model further advances this direction through an SO(2)-based tensor product approach. This comparison specifically examines the computational efficiency of these different tensor product operations, rather than model accuracy. We sincerely hope our clarification satisfies your question.

I would like to thank the authors for their thoughtful responses and the progress highlighted in their rebuttal. I have updated my ratings for the paper.

Dear reviewer:

We sincerely thank you for your careful review and constructive feedback on our paper. We particularly appreciate your recognition of our two key technical contributions - the strictly local SLEM architecture and the parameterization of invariant overlap operators. Your positive assessment of these innovations as being "novel" and "technically sound" is very encouraging.

In your review, you raised two main concerns: one regarding our evaluation testing along the trajectories of molecular dynamics, and another about the theoretical soundness about the strict local design of our approach, especially for the Hartree term in DFT. We would like to address these concerns point by point:

1. For the first concerns about the generalization test on MD trajectory, and suggestion of test on transferring to different materials with same chemical species.

We appreciate the reviewer's suggestion about cross-system generalization tests. While such extensive transferability (training on AB, BC, AC systems and testing on ABC) would be valuable for general-purpose models, we want to point out that specialized models trained for specific systems are both sufficient and preferable in many practical applications. Many research scenarios focus on studying the dynamic evolution and properties of specific materials or molecular systems, where the primary concern is the model's ability to generalize across different configurations of the same system, rather than across entirely different chemical compositions. Nevertheless, to demonstrate our model's capability for cross-system generalization, we have conducted evaluations on the well-established QH9 dataset, where the training and testing sets consist of molecules with different chemical compositions and varying numbers of atoms. Our model has shown excellent transferability in these tests, achieving competitive performance compared to the reported results of QH-net. We list the current results here, and more extensive results will be updated in the paper shortly.

|dataset: QH9-Stable-size-ood| unit (1e-6 Ha)

|QH-net(reported) | QH-net(reproduced) | LEM(Ours) |

|75.0 | 83.2 | 57.7 |

Additionally, we would like to emphasize that the generalization capability along MD trajectories is of significant practical importance. In real applications, our SLEM model only requires training data from a small portion of the MD trajectory. Once trained, it can effectively predict quantum operators over much longer time scales. This capability is particularly valuable when combined with Machine Learning Force Field (MLFF) methods. In such scenarios, DFT calculations are only needed for a short time period to generate the energy, force, Hamiltonian/overlap, and density matrix data for training both MLFF and SLEM models. Subsequently, by running AIMD with MLFF and predicting the Hamiltonian along long-time simulations, we can efficiently compute many important electronic properties that are time/temperature dependent or only emerge after extended MD simulations. This approach significantly reduces the computational cost of DFT calculations while maintaining accuracy.

2. For the second concern about the theoretical soundness of the strict local design: The theoretical foundation for our local design comes from the behaviour of Hamiltonian matrix elements. For example, the Hartree potential matrix element between two atomic orbitals takes the form: $V_{ij}^H=\int\phi_i^*(r-R_A)V_H(r)\phi_j(r-R_B)dr$ where $R_A$ $R_B$ are the centres of atomic orbitals, this integral decays rapidly with the distance $R_A-R_B$ due to the localized nature of atomic orbitals. The reviewer raised an important point about the four-center integral (ij|kl) which appears in the Hartree-Fock method or the hybrid functionals DFT: $(ij│kl)=\int\int\phi_i^*(r_1-R_A )\phi_j(r_1-R_B )\frac{1}{|r_1-r_2|}\phi_k^*(r_2-R_C)\phi_l(r_2-R_D )dr_1dr_2$ Indeed, this integral can be significant even when the ij pair is distant from the kl pair, because when both overlaps (i|j) and (k|l) are large (meaning $|R_A-R_B|$ and $|R_C-R_D|$ are small), the Coulomb interaction $1/|r1-r2|$ decays slowly with the distance between the pairs. However, our work primarily focuses on standard DFT with local/semi-local exchange-correlation functionals (LDA/GGA/meta-GGA), where such four-center integrals do not appear in the Hamiltonian. As for the hybrid functionals, though not the focus of our work, due to the screening effects, the Coulomb interaction decays more rapidly, and we can still treat it using the local framework by using a larger cutoff, as has been demonstrated in recent work in DeepH [Tang, Z. et al. Nat. Commun. 15, 8815 (2024)].

In our implementation, the cutoff radii (rcut) are chosen based on the inherent localization of atomic orbitals used in DFT calculations. Since typical DFT orbital cutoffs range from 5 to 10 Bohr, our model's cutoff for the distance between orbital centres $(R_A-R_B)$ is naturally set to twice this value (10-20 Bohr). For periodic systems, this cutoff is quite natural - matrix elements for orbital pairs beyond this distance are set to zero, which is consistent with the locality of the basis functions. For small molecules, where all atoms might fall within this cutoff range, all orbital pairs are considered, effectively becoming a full-range calculation. This explains why our model performs well across different types of systems while maintaining its physically-motivated local design. Besides, to further improve the flexibility and applicability of our method, we also developed a semi-localized version called LEM, where the interaction between distant atoms is included but decays exponentially with their distance, with a trainable decay factor. We agree with the reviewer's concern and add a section to discuss the limitations of this method.

Thanks for the updates on the empirical evaluation. For the strict locality, I'm not entirely aligned yet. SLEM not only performs well, but actually outperforms is still quite anti-intuitive to me.

our work primarily focuses on standard DFT with local/semi-local exchange-correlation

I was not talking about the exact exchange in hartree fock, but on the computation of coulomb interaction (hartree term) which is also computed via contracting with the four center integral. As you agreed above, distant shell pairs would still have an non-negligible four center integral, and consequently non-negligible in the hartree term. The Cauchy-Schwartz inequality used in integral screening does help remove some of the entries, but I doubt using a hard cutoff is legit, otherwise it would have already been applied in electron integral.

Since typical DFT orbital cutoffs range from 5 to 10 Bohr

Not sure what you mean by "orbital cutoff range". I didn't find a description on what type of orbitals are used in this study, then I assume it is GTOs being used, which would not have any hard cutoff (thus the above discussion on the four center integrals). If you're using strictly local numerical orbitals then these would make sense.

For periodic systems, this cutoff is quite natural

In periodic systems, the computation often breaks down into a local and a long range term. The long range effects still need to be considered, and should vary with the configuration.

LEM

Good to know this variant, could you comment on the relative performance of LEM and SLEM?

About the DFT orbital cutoff

We are using Numerical Atomic Orbitals (NAO) in this study. Generally, numerical atomic orbital has a cutoff radius, beyond which, the orbital value is ignored since it is too small. This type of basis is widely used in LCAO DFT software for periodic systems such as SIESTA, OpenMX, ABACUS etc. For molecule systems, the GTO are more generally accepted and Indeed does not have a certain cutoff radius. However, since the basis decayed exponentially with radial distance, the cutoff is defined as the value beyond which the quantum operator matrix elements are small to ignore, which can be obtained during the data processing step.

We apologise for the lack of a detailed description of DFT calculations. We employ ABACUS[3] software to generate our data, and use NAO basis set generating follows [4,5] using SG15 pesuodopotential [6]. We have added a section in the appendix to discuss the data generation process, and a section for the cutoff selection criteria with more detail.

About LEM

The accuracy of the LEM model on some periodic systems reported in the paper (such as silicon) is similar to SLEM, while the data efficiency of LEM is not as good. However, in other cases when the screening effect is weak (for example the metal-molecule junction we studied recently), the SLEM would be less accurate than LEM. We need to admit that we only have limited comparison examples about SLEM and LEM since this paper is mostly focused on the prior one. However, we highly agree that a thorough comparison of local, semi-local and conventional MPNN models to discover the specific applied area would be a great contribution both for theoretical understanding and to guide the user communities. We are looking forward to studying this in future work.

[1] Resta, Raffaele. "Thomas-Fermi dielectric screening in semiconductors." Physical Review B 16.6 (1977): 2717.

[2] Ninno, Domenico, et al. "Thomas-Fermi model of electronic screening in semiconductor nanocrystals." Europhysics letters 74.3 (2006): 519.

[3] Li, Pengfei, et al. "Large-scale ab initio simulations based on systematically improvable atomic basis." Computational Materials Science 112 (2016): 503-517.

[4] Peize Lin, Xinguo Ren and Lixin He, Strategy for constructing compact numerical atomic orbital basis sets by incorporating the gradients of reference wavefunctions, Phys. Rev. B 103, 235131 (2021).

[5] Mohan Chen, Guang-Can Guo, and Lixin He, Systematically improvable optimized atomic sets for ab initio calculations, J. Phys.: Condens. Matter 22, 445501 (2010).

[6] Hamann, D. R. "Optimized norm-conserving Vanderbilt pseudopotentials." Physical Review B-Condensed Matter and Materials Physics 88.8 (2013): 085117.

For the anti-intuition of outperforming SLEM

We fully understand your concern. The SLEM model introduces extra constraints and, therefore should not be more accurate than well well-constructed MPNN model. We agree with this consideration. The good accuracy of SLEM, from our understanding, should be contributed mostly by the model design, parameterization of the equivariant features (as in 3.1, we standardized all the irreducible matrix element features using statistical quantities from the dataset) and the usage of SO(2) convolution. However, SLEM's data efficiency in periodic material systems mostly comes from the strict locality, as it helps the model to capture the dependency more physically in localized systems with a strong physical prior.

About the long-range effect of Hartree Term

Indeed, if we split the Hartree term in the above form, there exist four centre integrals that $\langle ik|lj\rangle$ is large even when A and B are far away from each other. However, the integration results are a collective effect that sums over all pairs of $kl$, which is hard to interpret directly. Moreover, we are looking for a locality in the Hamiltonian matrix element. It is a combination of the LCAO basis, hartree term, external potential, and xc functional. Therefore, we need to discuss the localization as a collective effect, which is contributed by the following properties:

$V_{ij}^H=\int\phi_i^*(r-R_A)\left[V_H(r)+V_{ext}(r)+V_{xc}(r)\right]\phi_j(r-R_B)dr$

1. The decaying behaviour of the LCAO basis.

   For all LCAO bases (whether GTO, NAO, or Slater-like), the radial part all decays rapidly with distance. Therefore, the long-range term in $V_H, V_{ext}, V_{xc}$'s contribution in the integration (that decays as 1/r) would be decreased fastly (often exponentially such as in GTO or Slater-like orbital).

2. The system's electronic neutralization condition.

   The electronic neutralization condition requires 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)=\sum_In_I(r-R_I)$ Then the sum of $V_H(r), V_{ext}(r)$ would be:

   $-\sum_I\frac{Z_I}{|r-R_I|}+\sum_I\int\frac{n_I(r'-R_I)}{|r-r'|}dr'$

   When $|r-R_I|$ is large, we can approximate $|r-r'|\approx|r-R_I|$, therefore: $-\frac{Z_I}{r-R_I}+\int\frac{n_I(r'-R_I)}{|r-r'|}dr'\approx-\frac{Z_I}{|r-R_I|}+\frac{Z_I}{|r-R_I|}=0$

3. The screening effect.

   we refer to [1,2] the study of screening radius $R_s$, which describes the system's electrostatic potential's reaction to the vibration of charges.

   In [1], the screening radius of typical insulator and semiconductor systems is reported as:

   Diamond: 2.76 a.u.

   Silicon: 4.28 a.u.

   Germanium: 4.71 a.u.

   In [2], for some nanoparticles, the $R_s$ is reported as:

   Si191H148: 5.36 a.u., Ge191H148: 5.61 a.u.

Combining with the above effects, we can see the $H_{ij}$ Hamiltonian elements under LCAO basis can be approximated locally dependent on the neighbouring area of atom $A,B$.