Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems

Thank you for your constructive comments and suggestions, they are exceedingly helpful for us to improve our paper. Our point-to-point responses to your comments are given below:

The asymptotic scaling in Collary 1 for the lowest eigenvalue seems unfit to describe Gaussian-type orbitals (MO). Due to their fast decay, the overlap is spatially limited. After a certain system size, more atoms are unlikely to affect the lowest eigenvalue.

Our scaling analysis is grounded in both previous studies and empirical observations from various datasets. Previous studies [1], as illustrated in Figure 13, have established foundations for this scaling behavior. Our experimental results from carbon atom systems (Figure 12), as well as observations from QH9 and PubChemQH datasets (Figure 6, and anonymous link: https://ibb.co/9tyNgZm), corroborate this scaling behavior. It's important to note, however, that we have not explored extremely large systems or highly specialized cases, which may exhibit different scaling behaviors.

[1]: Ida-Marie Høyvik. The spectrum of the atomic orbital overlap matrix and the locality of the virtual electronic density matrix. Molecular Physics, 2020.

l.62/63, "complex transformation" sounds overly complicated given the symmetric nature of the Hamiltonian.

We acknowledge that the term "complex transformation" may imply unnecessary intricacy, considering the symmetric nature of the Hamiltonian. To more accurately reflect the physical operation while maintaining technical precision, we have updated this term to "unitary transformation.”

While the comparison to a regression model is laudable, the comparison to equiformer has many confounding variables. Could the authors compare to their WANet but with a regression head?

Following your suggestion, we implemented a regression head variant of WANet for direct comparison. The results (WANet w/regression: 6.892 MAE, WANet w/WALoss: 0.7122 MAE, Equiformer V2: 6.955 MAE) clearly demonstrate that our improvements stem from the WALoss formulation and architectural design rather than confounding variables. This controlled comparison provides strong evidence for the effectiveness of our approach.

Unfortunately, the real-world DFT speed-ups remain limited.

We thank the reviewer for raising this important point. We report an 18% reduction in SCF cycles, which is a notable improvement, especially considering that our work targets significantly larger and more complex molecular systems compared to previous studies such as QH9 and QHNet, where SCF reductions of 12–33% were achieved compared to minao inital guess [1]. Given the increased complexity and size of our systems, this 18% reduction represents a considerable improvement, reflecting the effectiveness of our model in accelerating convergence for more challenging molecular simulations. This result highlights the scalability of our approach in handling large-scale quantum chemical calculations efficiently. Moreover, Hamiltonian prediction enables multiple downstream applications beyond SCF acceleration. WANet demonstrates its versatility through state-of-the-art performance in predicting various electronic properties such as HOMO-LUMO gaps and dipole moments.

[1]: Yu H, Liu M, Luo Y, et al. Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules[J]. Advances in Neural Information Processing Systems, 2024, 36.

Dear Reviewer,

Thank you for your insightful comments regarding the asymptotic behavior of the smallest eigenvalue in systems described by an overlap matrix. We agree that "only the scaling for very large systems matters for asymptotics," and appreciate your prompt to discuss the scaling behavior of our analysis.

You are correct that for infinitely separated atoms or molecules, the eigenvalues of the total system's overlap matrix match those of the individual, non-interacting subsystems. In this case, the smallest eigenvalue reaches a finite value due to the lack of overlap between distant basis functions. However, we respectfully contend that this example lies outside the relevant asymptotic regime. Asymptotic analysis, by definition, examines the scaling behavior of a system as its size grows while preserving the structure and interactions that define it. Infinitely separated atoms reduce the system to isolated subsystems, effectively forming a different physical system. Consequently, this scenario offers limited insights into the asymptotic scaling of the original system, where finite overlaps persist. For a meaningful analysis, we focus on a more relevant system:

$k$-Nearest-Neighbor Overlaps

Here, we analyzed a system where each basis function overlaps with its $k$-nearest neighbors. The overlap matrix $S \succ 0$ in this case is a symmetric banded matrix with bandwidth $2k + 1$:

$$S_{ij} = \begin{cases} 1, & \text{if } i = j, \\\\ s_{|i -j|} & \text{if } 0 < |i - j| \leq k, \\\\ 0, & \text{if } |i - j| > k, \end{cases}$$

To derive a scaling equation for the smallest eigenvalue, consider the case where $k=1$, which corresponds to the nearest-neighbor model. Other cases can be derived similarly using plane waves and discrete Fourier Transform. The nearest-neighbor model holds physical significance in the context of tight-binding models. For example, in chain-like polyenes, the Huckel model is based on a nearest-neighbor analysis, providing a meaningful framework for discussing this case. The overlap matrix $S$ is given as:

$$S = \begin{bmatrix} 1 & s & 0 & \cdots & 0 \\\\ s & 1 & s & \ddots & \vdots \\\\ 0 & s & 1 & \ddots & 0 \\\\ \vdots & \ddots & \ddots & \ddots & s \\\\ 0 & \cdots & 0 & s & 1 \\\\ \end{bmatrix}$$

The eigenvalues $\lambda_n$ of this tridiagonal matrix can be derived using the known solutions for such matrices:

$$\lambda_n = 1 + 2s \cos\left(\frac{n\pi}{B+1}\right), \quad n = 1, 2, \dots, B$$

The smallest eigenvalue corresponds to $n = B$:

$$\lambda_{\text{min}} = 1 + 2s \cos\left(\frac{B\pi}{B+1}\right)$$

Using the trigonometric identity $\cos(\pi - \theta) = -\cos(\theta)$ and approximating for large $B$:

$$\cos\left(\frac{B\pi}{B+1}\right) = -\cos\left(\frac{\pi}{B+1}\right) \approx -\left(1 - \frac{1}{2}\left(\frac{\pi}{B+1}\right)^2\right)$$

Substituting back:

$$\lambda_{\text{min}} \approx 1 - 2s + s\left(\frac{\pi}{B}\right)^2$$

Since $S \succ 0$, this yields $1- 2s \geq 0$. Here, we investigated the smallest eigenvalue $\lambda_{\text{min}}$ as $B \to \infty$, identifying two distinct cases:

1. Saturating Case: When $s < 0.5$ the smallest eigenvalue saturates at a finite value $1 - 2s$ as $B \to \infty$.

2. Non-Saturating Case: For $s = 0.5$ , $\lambda_{\text{min}}$ decreases indefinitely with increasing $B$.

We would also like to emphasize that, in both cases, there exists a finite range of $B$ where the quadratic scaling holds true, as you correctly identified as the intermediate regime.

Notably, your observation aligns closely with our updated analysis. In the case of Gaussian-type orbitals, the saturating scenario is indeed more likely to occur due to inherent spatial limitations.

The results in Table 1 do not provide enough evidence that WANet architecture is superior to baselines. A comparison to other Hamiltonian predicting models, such as [2] and [3], is necessary to validate the claims about the proposed architecture.

We have extended our evaluation to include PhisNet. While PhisNet achieves competitive accuracy (Hamiltonian MAE: 0.5166), it demonstrates significantly slower inference speed compared to both QHNet and WANet, processing only 0.106 iterations per second—about 4 to 10 times slower than QHNet's 0.45 it/s and WANet's 1.09 it/s and takes much more GPU memory. We believe this might be attributed to PhisNet's inefficient handling of higher order spherical tensors. In contrast, our approach achieves better accuracy (Hamiltonian MAE: 0.4744) while substantially improving computational efficiency. These results demonstrate WANet's ability to effectively balance accuracy and scalability.

One of the paper's main claims is that improved Hamiltonians can be used for downstream regression tasks and to accelerate DFT computations. Claiming that WANet is superior to regression-based models in the HOMO-LUMO gap prediction task seems a little ambitious, given that only one baseline (Equiformer v2) was used. To give a better perspective, it might be worth comparing WANet with SOTA models, such as UniMol+ [4] or other models from [5].

We appreciate the reviewer's insightful comments. In response to the request for more comprehensive baseline comparisons, we have expanded our evaluation to include UniMol+ and UniMol2 [4]. Our results demonstrate that WANet achieves superior performance in HOMO-LUMO gap prediction compared to these state-of-the-art models.

The System Energy MAE is also ~50 kcal/mol, almost 50 times larger than metrics published in [1] and [6].

We thank the reviewer for raising this important point regarding system energy accuracy. While the absolute MAE of 47 kcal/mol appears high, it is important to note that in quantum chemistry applications, relative energy differences between conformational states are typically more relevant than absolute energies. To better evaluate our method's practical utility, we conducted an additional conformational energy analysis:

For each molecule in our test set, we:

1. Sampled 100 conformations.

2. Computed ΔE between each conformation and a reference structure using both DFT and WANet.

3. Calculated the MAE of these energy differences.

This analysis yielded a conformational energy MAE of 1.1 kcal/mol, which is competitive with state-of-the-art end-to-end energy prediction models. These results demonstrate that WANet effectively captures the physically meaningful energy variations required for applications

Importantly, WANet achieves this while solving the more challenging problem of Hamiltonian prediction, which requires learning the full quantum state rather than just scalar energies. To the best of our knowledge, we are the first to evaluate Hamiltonian prediction models using system energy metrics derived directly from the predicted Hamiltonian. This approach provides a more stringent test of physical accuracy compared to traditional property-specific prediction tasks.

Moreover, the acceleration rate seems relatively small (18%), and most SCF iterations are still required when starting from the predicted Hamiltonian.

We thank the reviewer for raising this important point. We report an 18% reduction in SCF cycles, which is a notable improvement, especially considering that our work targets significantly larger and more complex molecular systems compared to previous studies such as QH9, where SCF reductions of 12–33% were achieved compared to the minao inital guess [5]. Given the increased complexity and size of our systems, this 18% reduction represents a considerable improvement, reflecting the effectiveness of our model in accelerating convergence for more challenging molecular simulations. This result highlights the scalability of our approach in handling large-scale quantum chemical calculations efficiently. Moreover, Hamiltonian prediction enables multiple downstream applications beyond SCF acceleration. WANet demonstrates its versatility through SOTA performance in predicting various electronic properties such as HOMO-LUMO gaps and dipole moments.

Thank you for your constructive comments and suggestions, they are exceedingly helpful for us to improve our paper. Our point-to-point responses to your comments are given below:

The WANet model's architecture is hard to understand and poorly motivated. There are no ablations for various building blocks of the model, and no architecture hyperparameter values are provided.

We appreciate the reviewer's insightful question. While the focus of our paper is on the novel loss function and dataset, we acknowledge that a clearer explanation of the model architecture would be beneficial. The WANet architecture builds upon well-established equivariant neural network designs, which have been extensively studied and validated in the field [1,2,3]. We have added references to seminal works on equivariant architectures.

To further address this concern, we have added a supplementary figure (Figure 10, and anonymous link: https://ibb.co/HTrpZSM) illustrating WANet's architecture and its key components. This visual aid will help readers better understand how WANet enhances the scalability of Hamiltonian prediction for large molecular systems. Each architectural element has been carefully selected to address specific computational and physical challenges at scale.

First, WANet's architecture addresses a critical limitation in existing Hamiltonian prediction methods: their inability to scale to larger basis sets and molecular systems. When using larger basis sets, such as Def2-TZVP, higher-order irreducible representations (PubChemQH requires order $L = 6$) are necessary to accurately capture angular dependencies in molecular orbitals. This poses a severe computational challenge for traditional SE(3)-equivariant methods, including QHNet, which become extremely expensive due to their computational complexity. WANet overcomes this bottleneck by introducing SO(2) convolutions, reducing computational complexity from $O(L^6)$ to $O(L^3)$. This improvement enables WANet to efficiently process high tensor degrees. Moreover, by reducing calculations to sparse matrix products, WANet achieves greater computational efficiency and scalability to larger systems. Without SO(2) convolutions, handling the scale of our PubChemQH dataset and larger systems would be significantly more challenging and computationally prohibitive.

Second, large molecular systems exhibit fundamentally different physics at various distance scales. As molecular size increases, long-range interactions become more prominent [3]. WANet's Mixture-of-Experts architecture is designed to efficiently model this complex physics. It employs specialized experts for different interaction ranges, capturing both short-range effects (such as covalent bonding) and long-range phenomena (such as electrostatics). By sparsifying these experts, WANet achieves a rich representation of molecular interactions while maintaining computational efficiency, making it particularly well-suited for large-scale systems.

Third, the Hamiltonian matrix, which fully characterizes the quantum state and electron distribution, presents unique challenges in prediction due to complex electron correlation effects that intensify as the system size grows [2]. To address this, WANet incorporates the MACE architecture's density trick, enabling efficient computation of many-body interactions without explicit calculation of all terms. This approach is crucial for maintaining accuracy as molecular size increases and electron correlation effects become more pronounced, ensuring WANet's scalability and accuracy in Hamiltonian prediction for large systems.

We agree that more detailed ablation studies would strengthen the paper. In the revision, we have included the following:

1. Ablation studies on the building blocks of the model, which is included in Table 9 and below.

2. A more detailed description of architecture hyperparameters, including hidden channel dimensions, the number of layers, and other relevant parameters, which is included in Table 12.

We hope these additions address your concerns and provide a more comprehensive understanding of our model's architecture and performance.

Thank you for your thoughtful questions and remarks. We appreciate the opportunity to provide further clarifications. Below, we address each of your points in detail.

Training Setup for Main Experiment:

We have included the detailed training setup for the main experiments in the revised Appendix K Table 12. For the PubChemQH dataset, we used an 80/10/10 train/validation/test split, resulting in 40,257 training molecules, 5,032 validation molecules, and 5,032 test molecules. We trained all models of maximum 300,000 steps with a batch size of 8, using early stopping with a patience of 1,000 steps. WANet converges at 278,391 steps, QHNet at 258,267 steps, and PhisNet at 123,170 steps. All models used the Adam optimizer with a learning rate of 1e-3 for PubChemQH, along with a polynomial learning rate scheduler with 1,000 warmup steps. We used gradient clipping of 1, and a radius cutoff of 5 Angstrom. For QHNet and PhisNet, we referred to the official implementation for these two models. For detailed model hyperparameters, please refer to Table 12.

Training Setup for HOMO, LUMO, and GAP Predictions:

We have added the training details for the HOMO, LUMO, and GAP predictions in Appendix K Table 13. For a fair comparison, all regression models used identical dataset splits (40,257 training molecules, 5,032 validation molecules, and 5,032 test molecules). We used batch size of 32 and consistent optimizer settings across all models: Adam optimizer with a learning rate of 1e-3 and a polynomial learning rate scheduler with 1,000 warmup steps. WANet followed the same training setup described in Table 12. The complete regression model hyperparameters are detailed in Table 13. For Equiformer V2, Uni-Mol+, and Uni-Mol2 models, we used their original implementations.

Energy Prediction Experiment Details:

For the energy prediction experiment, we employed Gaussian sampling to introduce noise in the atomic positions. Starting from each molecule's initial conformation from the test sets, we generated 100 perturbed geometries by adding random Gaussian noise ($\sigma$ = 0.1 Å) to the atomic coordinates. The potential energy values for these perturbed conformations were computed using the B3LYP/def2-TZVP level of theory. The resulting energy standard deviation between sampled conformations ranged from 10 to around 100 kcal/mol (with average std =49.84kcal/mol), depending on the molecule's flexibility. For baseline models (included in the Table below), the PhisNet model with WALoss achieved 9.90 kcal/mol MAE, the QHNet model with WALoss achieved 10.98 kcal/mol MAE, and the WANet model without WALoss achieved 48.92 kcal/mol MAE. QHNet without WALoss yielded MAEs of 50.56 kcal/mol. Interestingly, we observed that the model without WALoss produce meaningless prediction, demonstrating the effectiveness of WALoss. It worth noting that while more complex sampling strategy could be used, those methods are more computationally expensive as they require DFT calculations.

Additional Remarks:

• We apologize for the missing reference to Figure 10 in the main text. We have now added a reference to it.
• We have included the PhisNet w/WALoss model in the revised Table 1 for completeness.

We thank you again for your insightful questions and hope that our responses have addressed your concerns satisfactorily. We believe that these clarifications will strengthen the paper and make it more accessible to the readers.

Thank you for your constructive comments and suggestions, they are exceedingly helpful for us to improve our paper. Our point-to-point responses to your comments are given below:

The paper's main weakness lies in its insufficient differentiation from prior work, particularly regarding the use of eigenvalues in Hamiltonian prediction loss functions.

We appreciate the reviewer’s feedback and the opportunity to clarify our contributions. To the best of our knowledge, WALoss is the first loss function specifically designed for Hamiltonian prediction that addresses the scalability challenges of large molecular systems. Through a comprehensive review of venues such as ICLR, ICML, NeurIPS, JACS, and other prominent chemistry proceedings, we found no prior work directly tackling this topic. However, we identified a study by Gu et al. (2022)[1] that uses eigenvalues as regression targets. While that study explores eigenvalue-based methods, WALoss introduces key innovations that enable scalability to much larger molecular systems, representing a significant advancement over existing approaches.

The approach by Gu et al. (2022)[1] that directly optimizes eigenvalues necessitates costly eigendecomposition at each training step and encounters difficulties with backpropagation through eigensolvers. In contrast, WALoss precomputes the ground-truth eigenvectors C* once at the start of training and performs an efficient basis transformation on the predicted Hamiltonian. This eliminates the need for repeated eigendecompositions during training, allowing for much faster backpropagation compared to naive methods. Our ablation studies (Table 4) demonstrate that naive eigenvalue-based losses suffer from fundamental limitations for large systems. Beyond numerical instability and convergence difficulties, these approaches fail to produce physically meaningful results, likely due to the challenges in backpropagating through eigensolvers and inconsistencies in eigenspace alignment across training iterations. The WALoss formulation addresses these issues by aligning predicted Hamiltonians with ground-truth eigenspaces through a fixed basis transformation, resulting in more stable training dynamics.

WALoss is specifically designed to handle the unique challenges of large molecular systems (40–100 atoms) with basis sets of up to 2,000 basis—a scale previously intractable with existing eigenvalue-based approaches. The precomputed basis transformation strategy enables efficient training on systems of this size. As shown in our experiments, WALoss significantly improves both the accuracy of predicted energies and the convergence of SCF iterations compared to previous methods, particularly for large molecules in the PubChemQH dataset.

We have expanded the discussion of these technical contributions and included additional relevant references in the revised manuscript. We also welcome any suggestions for other related work that we may have overlooked.

[1]: Gu Q, Zhang L, Feng J. Neural network representation of electronic structure from ab initio molecular dynamics[J]. Science Bulletin, 2022, 67(1): 29-37.

Could you provide details about the wall-clock training time comparison between models with and without WALoss, the GPU memory requirements (particularly for the largest molecules), and how the computational overhead scales with molecule size and atomic number?

Thank you for raising this important consideration regarding computational efficiency. The computational overhead of WALoss is minimal, as we do not perform eigendecomposition at every training step. Instead, we compute eigenvectors only once during initialization and cache them for subsequent use. Our benchmarks confirm this efficiency: WALoss results in only a 3.6% increase in training time (39.22 vs. 37.85 hours of wall-clock time) and negligible additional GPU memory usage (15.86 vs. 15.12 GB) on an NVIDIA A6000 GPU compared to the baseline model. We have added detailed timing and memory profiling results to the appendix. The additional computation primarily arises from the basis transformation step (two matrix multiplications), which scales as $O(B^3)$, where $B$ is the number of basis functions. For future work, we believe the computational cost could be further reduced through fragmentation methods and techniques from linear-scaling DFT calculations [1]. Specifically, leveraging the inherently sparse structure of molecular Hamiltonians—where values are concentrated around the diagonal—could enable distributed matrix operations. This approach could potentially reduce computational cost for large molecular systems.

[1]: Mohr S, Ratcliff L E, Genovese L, et al. Accurate and efficient linear scaling DFT calculations with universal applicability[J]. Physical Chemistry Chemical Physics, 2015, 17(47): 31360-31370.

How were the loss term weights ?The balance between these terms seems crucial for the method's performance, yet the paper doesn't discuss the process of selecting these hyperparameters or their sensitivity analysis.

We determined the loss weights through a principled ablation study on a validation set. While keeping $\lambda_1$ and $\lambda_2$ fixed at 1, we varied $\lambda_3$ from 0.5 to 3 to evaluate the impact of WALoss weighting. The method demonstrated robustness across a reasonable range of values, with $\lambda_1=\lambda_2=1$ and $\lambda_3 = 2.5$ providing the best performance. We have included this analysis (Table 11) in the revised manuscript for clarity, along with detailed performance metrics for different $\lambda_3$ values.

There appears to be a discrepancy in Table 4's Hamiltonian MAE results, where Naive Loss shows a notably lower error (0.0412) compared to other methods. Could you explain this unexpected result and its implications?

Thank you for identifying this error. We sincerely apologize for the typographical mistake in Table 4. The correct Hamiltonian MAE for Naive Loss is 0.4912, not 0.0412. This corrected value aligns with both theoretical expectations and our experimental observations. We have carefully proofread the entire manuscript to ensure there are no other numerical inconsistencies. We appreciate your attention to detail, which helps maintain the rigor and clarity of our presentation.

Thank you for your constructive comments and suggestions, they are exceedingly helpful for us to improve our paper. Our point-to-point responses to your comments are given below:

It is observed that if using WALoss on WANet, while Hamiltonian prediction performance improves, the prediction performance on the two quantum properties degrades. Could authors give explanations why this phenomenon happens on small-scale QH9 dataset but not on large-scale PubChemQH dataset (Table 1)?

We appreciate your observation and would like to clarify our findings. Our results in QH9 show that while the Hamiltonian MAE performance degrades, the prediction accuracy for the two quantum properties improves. This varying impact of WALoss across datasets underscores a key challenge in Hamiltonian prediction: developing an appropriate metric in the matrix space that accurately correlates with the metrics of derived quantum properties, such as system energy, wave function.

In larger systems (i.e., PubChemQH), WALoss enhances optimization stability by providing physically meaningful guidance, leading to improvements across multiple properties. However, in smaller systems (i.e., QH9), the interaction between WALoss and traditional loss terms may result in local optima that prioritize quantum properties over Hamiltonian MAE. This scale-dependent behavior illustrates the complexity of multi-objective optimization in Hamiltonian prediction.

Notably, the improved performance in larger systems is particularly significant, as these are more representative of real-world applications and have historically been more challenging to handle. This finding suggests that the benefits of WALoss for practical applications outweigh the trade-offs observed in smaller systems.

Authors are encouraged to report the results of QHNet with WALoss on QH9 datasets.

We appreciate your suggestion and have implemented it. The results for QHNet with WALoss on the QH9 dataset are now included in Table 2 of the revised manuscript. Our analysis indicates that incorporating WALoss leads to improved quantum property prediction performance on the QH9 dataset, demonstrating the effectiveness of our approach across various model architectures and datasets.

For practical application targets, both prediction accuracy and speed are important for an effective method. Authors are encouraged to give analysis about computational complexity or report inference speed of the proposed method, and compare them with QHNet [1].

We have expanded our efficiency analysis in Section 5.4 to provide a more comprehensive overview of WANet's performance. Our model demonstrates significant improvements in computational efficiency compared to QHNet. Specifically, WANet achieves an inference speed of 1.09 iterations per second—more than double QHNet's 0.45 iterations per second. The training time has been reduced by over 50%, from 90.43 hours for QHNet to 39.22 hours for WANet. Additionally, WANet requires significantly less memory, utilizing 15.86 GB of peak GPU memory compared to QHNet's 26.49 GB.

These efficiency gains can be attributed to one key architectural innovation: our efficient SO(2) convolution design. This enhancement allows WANet to process large molecular systems more effectively while maintaining high accuracy. For a detailed breakdown of these performance metrics, please refer to Figure 3 in the paper.

Hi Haiyang,

Thank you so much your interests in our work! The code and the data will be released soon. Please stay tuned!