Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems
Authors: Yunyang Li
Zaishuo Xia
Lin Huang
Xinran Wei
Samuel Harshe
Han Yang
Erpai Luo
Zun Wang
Jia Zhang
Chang Liu
Bin Shao
Mark Gerstein
Venue: ICLR 2025 Spotlight (2025)
Keywords: AI for Science
Quantum Chemistry
EGNN
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present a new framework to enhance the scalability and applicability of Kohn-Shan Hamiltonian.

Abstract: Density Functional Theory (DFT) is a pivotal method within quantum chemistry and materials science, with its core involving the construction and solution of the Kohn-Sham Hamiltonian. Despite its importance, the application of DFT is frequently limited by the substantial computational resources required to construct the Kohn-Sham Hamiltonian. In response to these limitations, current research has employed deep-learning models to efficiently predict molecular and solid Hamiltonians, with roto-translational symmetries encoded in their neural networks. However, the scalability of prior models may be problematic when applied to large molecules, resulting in non-physical predictions of ground-state properties. In this study, we generate a substantially larger training set (PubChemQH) than used previously and use it to create a scalable model for DFT calculations with physical accuracy. For our model, we introduce a loss function derived from physical principles, which we call Wavefunction Alignment Loss (WALoss). WALoss involves performing a basis change on the predicted Hamiltonian to align it with the observed one; thus, the resulting differences can serve as a surrogate for orbital energy differences, allowing models to make better predictions for molecular orbitals and total energies than previously possible. WALoss also substantially accelerates self-consistent-field (SCF) DFT calculations. Here, we show it achieves a reduction in total energy prediction error by a factor of 1347 and an SCF calculation speed-up by a factor of 18\%. These substantial improvements set new benchmarks for achieving accurate and applicable predictions in larger molecular systems.

Published as a conference paper at ICLR 2025
ENHANCING THE SCALABILITY AND APPLICABILITY OF KOHN-SHAM HAMILTONIANS FOR MOLECULAR SYSTEMS
Yunyang Li
Zaishuo X
Lin Huang Xinran Wei* Han Yang* Sam Harshe
Zun Wang
Chang Liu
Jia Zhang * Bin Shao
Mark B. Gerstein
(huang. lin,
jia zhang ) (emicrosoff com markegersteinlab... org
Yale University
UC Davis
MSR A14Science
ABSTRACT
Density Functional Theory (DFT) is a pivotal method within quantum chemistry and materials science, with its core involving the construction and solution of the Kohn-Sham Hamiltonian. Despite its importance, the application of DFT is frequently limited by the substantial computational resources required to construct the Kohn-Sham Hamiltonian. In response to these limitations, current research has employed deep-learning models to efficiently predict molecular and solid Hamiltonians, with roio-translationa symmetries encoded in their neural networks. However, the scalability of prior models may be problematic when applied to large molecules, resulting in non non-phyne non-physical predictions of ground-state properties. In this study, we generate a substantially larger training set (PuhChemQH) than used previously and use create scalable model for DFT calculations with physical accuracy. For our model, we introduce a loss function derived from physical princi- ples, which we call Wavefunction Alignment Loss (WALoss). WALoss involves performing basis change on the predicted Hamiltonian to align it with the ob- served one; thus, the resulting differences can serve surrogate for orbital energy differences, allowing models to make better predictions for molecular orbitals and total energies than previously possible. WALoss also substantially accelerates self-consistent-fiet (SCF) DFT calculations. Here, we show it achieves reduc- tion in total energy prediction error by factor of 1347 and an SCF calculation speed-up by a factor of 18%. These substantial improvements set new benchmarks for achieving accurate and applicable predictions in larger molecular systems.
INTRODUCTION
Density functional theory (DFT) (Kohn & Sham 1965 Hohenberg & Kohn 1964 Martin 2020
has been widely used in physics (Argaman & Makov 2000:Jones] 2015: Van Mourik et al. 2014).
chemistry (Levine et al. 2009 Van Mourik et al.) 2014), and materials science (March 1999.
Neugebauer & Hickel, 2013) to study the electronic properties of molecules and solids. This methodology particularly valued for balanced blend of computational efficiency and accuracy, rendering a versatile choice for investigating electronic structure (Kohn et al., 1996; Parr & Yang 1995). spectroscopy (Neese. 2009 Orio et al., 2009), lattice dynamics (Dal Corso et al. 1997 Wang et al. 2021), transport properties (Bhamu al.) 2018), and more. The most critical step in applying DFT to a molecule constructing the Kohn-Sham Hamiltonian, which consists of the kinetic operator, the external potential, the Coulomb potential (also known the Hartree potential), and the exchange-correlation potential (Kohn & Sham 2020).
1965 Hohenberg & Kohn 1964 Martin
The Hamiltonian matrix contains crucial information about molecular systems and their quantum states (Kohn & Sham 1965 Hohenberg & Kohn| 1964; Yu al., [2023b] Zhang al. 2024). This matrix facilitates the extraction of various properties, including the highest occupied molecular orbital (HOMO) and lowest unoccupied molecular orbital (LUMO) energies, the HOMO-LUMO gap, total
"Equal contribution.

---

Published as a conference paper at ICLR 2025
energy, and spectral characteristics (Eisberg & Resnicck 1985 Zhang al. 2024). These properties
are vital for analyzing conformational energies (St.-Amant et al. 1995). reaction pathways (Farberow
let al. 2014), and vibrational frequencies (Watson & Hirst 2002). However, the efficiency of DFT is constrained by the self-consistent field (SCF) iterations required to solve the Kohn-Sham equations and achieve consistent charge density. These iterations scale as ((NTT) ~ O(N'T) (Tirado- & Jorgensen 2008 Yu et al. (023a)), where N represents the number of electrons and 7 denotes the number of SCF cycles, making them particularly resource-intensive for large systems.
Consequently, this computational intensity underscores the critical need to develop more efficient methodologies to determining the Hamiltonian without relying on SCF iterations. This challenge has catalyzed interest in leveraging deep learning to predict Hamiltonians directly from atomic configurations while adhering to the inherent symmetries of molecular systems (Unke et al. 202 la Kochkov et al. 2021 Yu et al. 2023b Li et al. 2022, Schütt et al. Gastegge et al. 2020
Hermann et al. 2020 Yin et a 2024 Zhang et al. 2024 Luo et 2025). Notably, the Hamiltonian
matrix is subject to unitary transformations under molecular rotations due to its spherical harmonics component. To tackle these transformations, researchers have developed -equivauuvari neural networks such as PhisNet (Unke al., 2021b) and QHNet (Yu et al. 2023b). These networks employ high-order spherical tensors and the Sebsch-GGordon tensor product to construct the predicted Hamiltonians.
However, current implementations of SE(3)-equivarian neural networks face considerable scalability challenges when applied to large molecular structures. Notably, prevailing techniques for predicting the Hamiltonian matrix predominantly utilize mean absolute error (MAE) or mean square error (MSE) as the loss functiio (Unke et al. 202II] Yu al. (023b). We contend that these elementwise losses alone are insufficient for accurately learning Hamiltonians in large systems. To underscore this point, we introduce the PubChemQH dataset, which comprises molecular Hamiltonians for structures with atom counts ranging from 40 to 100. In contrast, the previously curated dataset (Yu et al. 2023a) is limited to molecules with no more than 31 atoms. Figure demonstrates a phenomenon enabled by the introduction of the PubChemQH dataset, which we term Straling-Inducee MAE AAp-iipliayyiliiy Divergence (SAD). The applicability of the Hamiltonian is assessed by evaluating the system energy eπror derived from pyscf (Sun et al. 2020 Sun 2015 Sun et al. 2018). While small molecules yield system energies close to the ground truth with relative tomil MAEs up to 200%, larger molecules exhibit profound inaccuracies. Energy errors escalate to much as 1, 1,000,000 kcal/mol at only 0.01% % relative MAE, rendering the Hamiltonians inapplicable. This striking disparity emphasizes the inadequacy of MAE for large systems and highlights the pressing need for new methodologies to enhance the scalability and applicability of predicted Hamiltonians.
To address these limitations, this work aims to enhance Hamiltonian learning for large systems by utilizing wavefunctions and their corresponding energies as surrogates to improve Hamiltonian applicability. The wavefunction is crucial for Hamiltonian prediction as it encapsulates the quantum state of system, enabling the validation of Hamiltonians based on their ability to accurately reflect the system's 's energy and dynamics. However, learning the electron wavefunction is non- trivial; inaccuracies in the machine learning-based Hamiltonian may lead to significant error in both electronic structure and wavefunctions. This challenge motivated us to introduce the Wavefunction Aligument Losy function (WALoss), which aligns the eigenspaces of predicted and ground-truth Hamiltonians without explicit backpropagation through eigensolvers. Additionally, to improve the scalability of Hamiltonian learning, present WANet, modernized architecture for Hamiltonian prediction that leverages eSCN (Passaro & Zitnick 2023) convolution and a sparse mixture of pair experts. Our pipeline and main contributions are summarized in Figure
BACKGROUND
In the framework of Density Functional Theory (DFT), a molecular system is defined by its nuclear configuration M = {z, R.,, where Z represents the atomic numbers of the nuclei and R their positions within the molecule. DFT focuses on determining the ground state of a system consisting of N electrons by minimizing the total electronic energy with respect to the electron density p(r). Here, p(r) is a functional of the set of N one-electron orbitals (1)( I where r E TR33 specifies the spatial coordinates of an electron.
We define the scalability as the model's accuracy under large molecules.
2

---

Published as a conference paper at ICLR 2025
QH9 Molecule
PubChemQH Molecule
102
Gauastan
Gauesaan
120300
Qithlr
Intiaa Gueds
Inital Suess
10000
Energy Eiror (kcalho))
Energy Error (kcaltmol)
100
MAE
MAP
Figure Visualization of the SAD phenomenon. The y-axis (in symmetrical log scale Webber 2012 represents the system energy error derived from the perturbed Hamiltonian, while x-axis shows the relative MAE(HLH')) MAE, defined as Mean(H") where H" represents the ground truth Hamiltonian and H denotes the predicted or perturbed Hamiltonian. The relative MAE induced by model learning errors or Gaussian perturbation. The Gaussian perturbation ensures that the perturbed matrix remains Hermitian. The initial guess line is derived from nao (Almlöf ct al. 1982 Van Lenthe et al. 2006). Current statc-of-the-ane models, such as QHNet, achieve relative MAE of up 10 on PubChemQH molecules. For small molecules, a 10 relative MAE is sufficient for accurate system energy predictions (left panel), but this accuracy does not extend to large systems (right panel).
For computational efficiency, these orbitals are represented using basis set (oa (r)) depending on the molccular geometry, denotes the number of basis and varies for different basis sets. The expansion coefficients of the orbitals organized in a matrix C E BBN, allowing each orbital (r) to be expressed as: (T) Σx= Caioa DFT seeks the minimal electronic energy by solving for the optimal coefficient matrix C via the Kohn-Sham equations, represented as H(C)C SCe, where H(C) E Bx denotes the Hamiltonian matrix. Notice that H is function of C, and can be computed using density fitting with complexity of O(33). ). S E B is the overlap matrix with elements S.g == I [ ())))()))) where ot denotes the complex conjugate of 6. represents diagonal matrix containing the orbital energies. This forms generalized eigenvalue problem, where comprises the eigenvectors and the diagonal elements However, the solution this problem is complicated by the interdependence between H(C) and C. To resolve this, traditional DFT employs the self-consistent field (SCF) method, an iterative process that refines the coefficients ((k) through successive approximations the Hamiltonian matrix (*), Each iteration commences by computing (k) leveraging ((k-1) and solves new eigenvalue problem:
are the eigenvalues.
(k) (c(k-1)) (.(k) (C())
converging stable Hamiltonian H and its corresponding eigenvectors C*, which define the electron density and other molecular properties derived from the Kohn-Sham equations.
Problem Formulation The objective of Hamiltonian prediction is obviate the need for self- consistent field (SCF) iteration by directly estimating the target Hamiltonian H* from a given molecular structure M. To achieve this, one could parameterize a machine learning model (.M)). The learning process is guided by an optimization process defined as:
Σ
dist (M),
arg min
D
(ED
where (D] denotes the cardinality of dataset D, and dist(;;)) is predefined distance metric.
SE(3) Equivariant Networks SE(3)-equivarian neural networks incorporate strong prior knowl- edge through equivariance. These networks utilize equivariant irreducible representation (irreps)
features built from vector spaces of irreducible representations to achieve 3D rotation equivariance. The vector spaces are (26 didmmennion,,, where degree & N represents the angular frequency of the vectors. Higher l values are critical for tasks sensitive to angular information, such as predicting the Hamiltonian matrix. Vectors of degree referred as type-l vectors, are rotated using Wigner-D matrices D( E 1(2(+1)x(2(-1) when rotating coordinate systems. Eigenfunctions of rotation in R can be projected into type- vectors using the real spherical harmonics y()) 52 - R where g = (f e R² [ = 1] is the unit sphere. Equivariant GNNs update irreps fea- tures through message passing of transformed irreps features between nodes, using tensor prod-
3

---

Published as a conference paper at ICLR 2025
PubchemQH
Resource
Broad Applicability
Integration
Cheerlory
Larger Dataset
Scalable Model
Inforence Spoed
Solve SCF
Solve
rm 000 600
Energy
* Mec
Downstream Applications
GHI PubChemOH
QHitrt
WAlett
SCF Acceleration
Dipole Moment
b
Model Architecture
x
x
WANet
HUMO-LUMO
Training
Gap
Energy Levels
Total Energy
Physics-based Loss Function
HC = SCe
Node Convolution Laver
PubChemQH Resource
Basls Change
(C) /CCC'
Mature of Pair Experts
Many Body Layor
molecules Def2TZVP basis
MAE-MSEE
Expansion
SCF CUDA Accelerated DFT
Wlveffunction Aligmment
ourregate
Nom-baagonal
Diagonal
5 [a,
Assemble
Scaling Induced MAE Applicppiyyyiiityy Divergence
Predicted Hamiltonian
Applicability
[., 15:1)
improvements
Figure 2: (a) We introduce PubChemQH, a new resouree for Hamiltonian learning that facilitates the exploration of the scalability challenge known SAD. (b) We present WANet, modern architecture designed for accurate Hamiltonian prediction. WANet incorporates SO(2) convolution, mixture of pair experts, and many-body interaction layer. mixture of pair experts constructs the non-diagonal block, while many-body interaction layer constructs the diagonal block. (c) Our loss module, WALoss, performs a basis transformation of the predicted Hamiltonian using the ground-truth Hamiltonian and the overlap matrix. This enhancement aims to improve the applicability of the predicted Hamiltonian weal-world scenarios. Our final loss function combines MAE-MSE loss with WAL.oss.
mg -2 (.).m)).(a.mm) " " where C denotes the Clebsch- ucts: (w 8 of y = E'r ---=-f -m==-t,
En,
ml))
Gordan coefficient, ls satisfies 121 f2 < E3 l1 + Ez. Note that denotes the m-th element in the irreducible representation with -f S m I and m E N. In equivariant graph neural networks, the message function m/, from source node target node is calculated using SO(3) convolution. The
fo-th degree of mis can be expressed as mis (f.) = Σt,e, We,,t,& 2 (e.) X (((f))((s)) where
Wr,,e, are the weight vectors, I, represents the irreducible representation of the source node and Fta P
WAVEFUNCTION-ALIGNNENN Loss
The applicability of the predicted Hamiltonian H depends significantly on the accuracy of its cigenvalues (orbital energies) and eigenvectors (basis coefficients). To assess the alignment between H and the actual Hamiltonian H' with respect to their eigenspaces, we present the following theorem:
Theorem 1. Let H and H represent Hamiltonian matrices, and S the overlap matrix. Define the perturbation matrix AH = - H. Let (H, and (H, S) be the ith generalized eigenvalues of H and H. respectively. Assume spectral gap 5 separates the generalized eigenvalues of H. H. x(-) denotes the conditional number of given matrix. ll2 represents the spectral norm. |A = eigenspace of and A satisfy:
Then, the difference in eigenvalues and the angle 0 between the
R(S) S
R(S)
AH
(i, (HH
|AH|I
sin 0 <
8
4

---

Published as a conference paper at ICLR 2025
Proof.
The proof is delegated to Appendix J.1
Corollary (Perturbation Sensitivity Scaling). Assume (S) of the overlap matrix S scales as: Amin (S) = +
A
where e > 0, A > 0, and a > 0 describe the saturation behavior. Then,
B
1+
No.
for a perturbation matrix AH, under the first-order perturbation approximation, the eigenvalue perturbation is bounded by:
A (H. - (H, sO(0B'/2 (BBB)++Blu))
Proof.
The proof is delegated to Appendix J.2
Remark The theorem highlights that the difference between the predicted and actual Hamiltonian matrices, when only considering the element-wise norm, lead to unbounded differences in eigen-
x(S) eigenvalue perturbations to the size of the basis B. This phenomenon underscores the catastrophic scaling associated with increasing B, which is manifestation of the aforementioned SAD phe- nomenon. We provide thorough discussion on a in Appendix J.3] To validate our theoretical a(S)
values/eigenvectors due to a significant [[S ratio. The corollary further elucidates the sensitivity of
analysis, empirically evaluated the distribution of TS] across both the QH9 and PubChemQH datasets (Figure 6). The results demonstrate that molecules in PubChemQH exhibit substantially higher ratios compared to QH9, indicating increased perturbation sensitivity in larger molecular systems. This provides strong empirical evidence for the SAD phenomenon and aligns with our theoretical predictions. Consequently, these findings suggest that when designing an effective super- visory signal for learning or optimization tasks involving Hamiltonian matrices, is crucial to take into account the interaction of the overlap matrix and the corresponding Hamiltonian to mitigate the potential instability caused by perturbations.
In light of this, we propose new loss function: the Wavefunction Alignment Loss. It is designed to preserve the integrity of the eigenstructure related to molecular orbitals. Let and e* represent the eigenvalues (orbital energies) of H and H*, respectively, and Ö and C* denote the corresponding eigenvectors (basis coefficients). We define primary form of the loss function by directly applying the Frobenius norm to the eigenvalues Calign = ". Σ ||ا [), [|2, where E,, and tt,, are derived
from solving the generalized eigenvalue problems: HČ = SČè and H* SC e", respectively.
However, this formulation has notable limitations: (1) Generalized eigenvalue problems suscepti- ble to numerical instabilities due to ill-conditioned matrices, leading to erroncous gradients during backpropagation through iterative eigensolvers, complicating optimization. (2) The loss function assigns uniform weights to all orbital energies, which is not practical as some orbital energies hold more significance. To address these issues, we begin by applying the following algorithm (Ghojogh et al. 2019 Golub & Van Loan 2013 to perform a simultaneous reduction of the matrix pair (H",))
Algorithm 1 Simultancous reduction of a matrix pair (H",S)
Require: Groud-truth Hamiltonian matrix H' and overlap matrix S
Ensure: Diagonal matrix e* and matrix C* such that (C*) TSC aad (C*) (C')
1: Compute the Cholesky decomposition S = GGT
2: Define M H H-G.HHGTT
3: Apply the symmetric QR algorithm to find the Schur form (Q*) = Q"
4: Compute C*
When the overlap matrix ill-conditioned, the eigenvalues e* computed by Algorithm [ can suffer from significant roundoff emror.. To mitigate this, we modify the algorithm by replacing the Cholesky decomposition of S with eigen (Schur) decomposition V SV = Σ. where V is the matrix of eigenvectors and Σ is the diagonal matrix of eigenvalues. We then substitute G with VΣ -1/2 This modification effectively reorders the entries of M" placing larger values towards the upper left-hand corner, thereby enhancing the precision in computing smaller eigenvalues (Golub Van Loan 2013 Wilkinson, 1988).
5

---

Published as a conference paper at ICLR 2025
Claim 1. Under optimal convergence condition, the ground truth eigenvector C* should diagonalize the predicted transformed matrix H, thereby satisfying the relation (C*) E".
HC*
Accordingly, we propose a refined loss function:
CwA = - Σ (HO (syyle
(2)
==1
where the subscript ())s denotes the i-th sample. This modified loss function explicitly penalizes deviations from the expected eigenstructure. It important to note that while (C*) HC may not yield diagonal matrix, the diagonal nature of e" implicitly enforces the predicted eigenstructure through the non-diagonal clements of the loss function. Here, provide another derivation of the WALoss by drawing a parallel to first-order perturbation theory. In perturbation theory, for a Hamiltonian H H' + AH, where AH the perturbation. The first-order energy shift is de T AH using the unperturbed eigenvectors vf* of H*, with the assumption that these eigenvectors remain fixed while adjusting the eigenvalues. Similarly, WALoss employs the ground- truth eigenvectors C* from H'C e" to evaluate the predicted Hamiltonian H" + AH. The loss could be rewritten as LWA + -
(t)) (H" + ()) C) E/1) [| assessing how well H* + AH reproduces e* in the C* basis. Here, (C*) (H AH)C* - + (C*) AHC* where (C*) AHC" mirrors the perturbation term wt AHH*
widely known that energies associated with occupied molecular orbitals (and LUMO) make up the most energy of the molecule. Thus. to capture the most relevant part of the eigenspectrum, we calculate the loss function with increased weight on the + lowest eigenvalues. Here, corresponds the number of occupied orbitals and additional eigenvalue accounts for the Lowest Unoccupied Molecular Orbital (LUMO). Let I represent the set of indices corresponding to the k + lowest eigenvalues, and ((),s)) )5 indexes the j-th eigenvalue for the i-th sample. The loss function is defined as follows:
1 n
Σ i=1
1 (Ci)), (c + Σ Σ JET
[C H (C))
(c 1
where denotes the number of samples, and p, are hyperparameters where > This adaptation ensures that the loss function places greater emphasis on the eigenvectors and eigenvalues corre- sponding to the occupied orbitals and LUMO. This improves the model's focus the critical part of the eigenspace, which is crucial for practical applications.
WANET
Here, we present WANet, modernized architecture for Hamiltonian prediction. First, unlike previous approaches, we propose a streamlined design for Hamiltonian prediction that consists of two essential components: the Node Convolution Layer and the Hamiltonian Head. The Node Convolution Layer operates on a localized radius graph, performing graph convolution to capture intricate atomic interactions. This block serves dual purpose: first, it generates an irreducible node representation. providing powerful input for the subsequent Hamiltonian Head. Second, it can be initialized with pretrained EGNN or reprogrammed for other downstream tasks, constituting a unified framework for molecular modeling. The Hamiltonian Head constructs both pairwise and many-body irreducible representations using the Clebsch-Gorddo tensor product. These representations are then utilized to assemble both the non-diagonal and diagonal components of the Hamiltonian matrix. The model architecture and ablation studies are detailed in Figure [0 and Table 9 respectively.
4.1 NODE CONVOLUTION LAYER
For the Node Convolution Layer, we replace the traditional SO(3) convolutions with Equivariant Spherical Channel Network (eSCN) (Passaro & Zitnick 2023 Liao et al. 2023). The eSCN
k
N for paired orbitals.

---

Published as a conference paper at ICLR 2025
framework primarily utilizees SO(2) linear operations, optimizing the computation of tensor products involved in the convolution process. Traditionally, SO(3) convolutions operate on input irreducible
representation (irrep) features v and spherical harmonic projections Ym, (ft=). By applying
rotation matrix Dis to Tis. aligning it with the canonical axis where f and m = 0, the spherical
harmonic projection Yik, (1)) becomes non-7ero exclusively for mf = 0. This condition
(.l...m)) simplifies the tensor product to ( (e..):) which remains non-zero only when m, = -- Subsequently, eSCN has demonstrated that this reformulation could be reparameterized using (O operations on the rotated tensors, and simplify the computation from O to O(L²), L is the degree of the representaion. A detailed mathematical framework of this method is elaborated in the
Appendix D
4.2 HAMILTONIAN HEAD
Sparse Mixture of Long-Short-Ramne Range Experts We introduce variant of the Gated Mixture-of-
Experts (GMoE) (Shazeer al. [017]] Clark et al. 2022 Riquelme al. 2021 Zoph al. 2022 Jiang et al. 2024) model by incorporating a Mixture-of Experts (MoE) layer tailored for pairwise molecular interactions. This enhancement draws inspiration from the Long-Shorr Short-Rangge Range Message Passing framework (Li et al. 2023), which differentiates between handling proximal and distal interactions through specialized layers. Our approach differentiates interaction dynamics based on distance, with closer pairs experiencing distinct interaction profiles compared to more distant pairs. This differentiation is achieved through a novel layer that substitutes the conventional pair interaction layer with sparse assembly of expert modules, each functioning autonomously Layer. We define the Pair Construction layer with the function F".
Pair Construction pair for the n-th expert, and delineate the output of the MoE layer with N experts as: (x ..x) Σn= Pr (x)) Faair (Zi, xx), where Pr. (a,, .x)) are the gating probabilities computed by the gating network, and denotes scalar multiplication. The gating probabilities are obtained by applying the Softmax function over the gating scores of all experts: xx)) Σm (x)] where = rbf( (|2))l2) applies a radial basis function the Euclidean distance | 122 with distance cutoff, and (z) represents the gating score for the n-th expert, computed as: (=) = z W + c,. Here, W are learned gating weights for expert n, and 6n is injected noise to encourage exploration and promote load balancing among experts. Specifically, we use Gumbel noise (Jang et al. 2016); en = log (- log (Uw)), Un ~ Uni 1). This noise enables differentiable approximation of the top-K selection, allowing for sparse expert utilization while maintaining gradient flow during training. To further promote load balancing among the experts, we introduce an auxiliary load balancing loss (Shazeer al., 2017):
(x)(G.,(=)
Cload Gaaadbaaanniing = N E (Eirg) (((e))))))) which encourages the gating network to allocate Ei.,,
("E,,,)) routing probabilities evenly across experts, preventing underutilization of any single expert. The Pair
Construction layer for each expert is defined as: (Fpuie(Iz. 2x))) Ei,d, WTid (2 C )
where æ and u/² are the 1;-th and 15-th irreducible representations of source node s and target node t, respectively; W are the learned weights that couple these representations into the output representation lo: and & denotes the tensor product.
Many-Body Interaction Layer Considering many-body interactions for the diagonal components of the Hamiltonian in molecular systems captures essential electron correlation effects (Szabo & Ostlund 2012: Jensen 2017). These interactions provide more accurate representation of the collective behavior of atoms, beyond pairwise approximations. This leads to a precise description of key quantum phenomena like electron delocalization and exchange interactions (Szabo & Ostlund 2012). For this purpose, we employ the methodologies of the MACE framework (Batatia et al.
2022 Kováes al. 2023 Batatia al., (023). Central to MACE is the adept conversion of first-
order features into higher-order features using the so-called density trick (Duval et al., 2023). This procedure initiates with the formation of generalized Mlebsch-Gordon tensor products from the
first-order features:
BM
Σ 2m
Cm II
Wig P.: where Im (l1m1, (.. b.w),,
7

---

Published as a conference paper at ICLR 2025
where I'm represents the input tensor, and Bir,,, the resultant tensor for the v-body. The coefficients Cmme are the generalized Seebsch-Gordan coefficients, ensuring the L-equivariance of the output tensor B M. Moreover, Omn is notably sparse and can be pre-computed efficiently.
5 EXPERIMENTS
In this section, evaluate WANet with WALoss the QH9 and PubChemQH datasets. Our evaluation metrics include MAE for the Hamiltonian, CHоMO- CLUMO, Ca, Coce, cosine similarity for the eigenvectors C, and relative SCF iterations compared to the initial guess, which are commonly used in previous works (Unke et al. 2021b Yu et al. 2023b a). Additionally, we introduce two new physics-related metrics MAE for Corb and System Energy to provide more comprehensive evaluation. Detailed descriptions of these metrics are provided in Appendix C.1
5.1 RESULTS ON THE PUBCHEMQH DATASET
Dataset Generation Process In our study, we investigated the scalability of Hamiltonian learning by utilizing a CUDA-accelerated SCF implementation (Ju et al. 2024) to perform computational quantum chemistry calculations, thereby generating the TubChemQH dataset. We began with ge- ometries from the PubChemQC dataset by (Nakata & Maeda 2023), selecting only molecules with molecular weight above 400. This filtration process resulted in dataset comprising molecules with 40 to 100 atoms, totaling over 50,000 samples. We chose the B3LYP exchange-correlation
functional (Lee et al. 1988 Beeke 1993 Vosko al.) 1980 Stephens al. 1994) and the Def2TZV
basis set (Weigend & Ahlrichs 2005: Weigend 2006) to approximate electronic wavefunctions. Generating this comprehensive dataset represents a substantial computational effort, requiring ap- proximately one month of continuous processing using 128 NVIDIA-V VVDI GPUs. We provide a comparison between the PubChemQH dataset and the QH9 dataset in Appendix H
Table I: PubChemQH experimental results. The energy units are presented in kcal/mol, N/A indicates that the metric is not applicable to specific model. The best-performing models are highlighted in bold. Detailed
training setups are provided Table 12 and Table 13
Facaaa
(arMAE..
Simitarty
Enery
GIUNe:
8.790
75,230
$3309
57790
2187:42
1NS2.872
2.32%
311% 68721 MC3
WANN
01224
50:140
02:35
4:311
512.45
628/9
PhiiNex or WAA
0 51165
11881
17.75
6.544
30.005
5.17
45.885
60.335
99%
WiNoe WKLoss
0.4744
8.112
6.10
1,37
14.835
7.330
60.0%%
47,093
51%
015512
34,13
36.501
42.700
25.183
1007
Exuitomerr Regrewiee
M/A
6955
6.552
3.022
N/A
N14
Nis
NIA
UniMee Repreoass
9513
7.414
10438
NGA
N6A
NIA
AGA
Table presents comparative analysis of various models' performance on the PubChemQH dataset. Despite a higher Hamiltonian MAE, WANet with WALoss significantly outperforms the other models in practical utility, as evidenced by the System Energy MAE and the required SCF iterations. Specifically. the System Energy MAE for WANet with WALoss is dramatically reduced from 63579.233 kcal/mol to 47.193 kcal/mol. Additionally, the relative SCF iterations required for WANet with WALoss is only 82%, compared to 371% for QHNet and 334% for WANet without WALoss. This substantial reduction in SCF iterations demonstrates the effectiveness of WALoss in accelerating the convergence process. Furthermore, it is worth noting that the initial guess matrix, although not achieving as low an MAE as QHNet or WANet without WALoss, shows improved utility with a System Energy MAE of 374.313 13 kcal/mol. This finding reinforces the idea that elementwise losses are insufficient. Additional evidence supporting this idea is provided in Appendix C.2
5.2 RESULTS ON THE QH9 DATASET
The QH9 dataset is a comprehensive quantum chemistry resource designed to support the devel- opment and evaluation of machine learning models for predicting quantum Hamiltonian matrices. Built upon the QM9 dataset, QH9 contains Hamiltonian matrices for 130,831 stable molecular geometries, encompassing molecules with up to nine heavy atoms of elements C, N, O, and F. These Hamiltonian matrices were generated using pysct with the B3LYP functional (Lee et al. 1988
Beeke 1993] Vosko et al. 1980; Stephens et al. 1994) and the de basis set. Table 2 presents comparative analysis of the performance of our model, WANet, against the baseline model, QHNet, on the QH9 dataset in both stable and dynamic settings. In the Qt9-stable experiments, WANet
8

---

Published as a conference paper at ICLR 2025
Table Experimental results on the QI19 dataset. energy units are presented kcal/mol.
Model
Hamiltonian MAE
MAE
C Csimilarity
QHNer
QHI stable
QHNer w/WAloss
0.0780
0.4901
96.35%
WANet
0.0502
0.5231
96.86%
WANet w/Wlloss
0.0914
0.4587
96.95%
QHNer
0,0471
0.2744
93.
QH9 dynamic
QHNet w/WAloss
0.0495
0,2658
54%
WAnet
0,0469
0.2614
99,69 69.68%
WAnet w/WAloss
0,0812
demonstrates superior performance, achieving higher accuracy compared to QHNet. Specifically, WANet significantly reduces both the Hamiltonian and occupied energy MAE while improving the cosine similarity of C. For the QH9-dynamic dataset, WANet consistently outperforms QHNet, further enhancing prediction accuracy. These results underscore the robustness and effectiveness of WANet in both stable and dynamic scenarios.
5.3 COMPARISON WITH A PROPERTY REGRESSION MODEL
Conventional machine learning approaches typically employ property regression, mapping molecular features directly to the desired property value (Blum & Reymond 2009 Montavon al. 2013).
common question arises: why use Hamiltonians instead of a property regression model? We argue that property regression methods often fail incorporate underlying quantum mechanical principles, limiting their generalization capability. To illustrate this, we compared the performance of WANet with WALoss to model utilizing Equiformer V2 (Liao & Smidi 2023), UniMol+ (Lu et al., 2023), and UniMol2 (Xiaohong et al. (2024) with invariant regression heads, using identical training and test sets. As shown in Table I WANet with WALoss demonstrates significantly lower MAE values in predicting key quantum chemical properties. Specifically, WANet with WALoss achieves an 88.88% improvement in CLUMO MAE and 58.8 18.81% improvement in ca MAE. Moreover, the Hamiltonian predicted by WANet with WALoss is not limited to specific properties. It enables the accurate calculation of various critical properties, such as electronic densities, dipole moments, and excited-state energies, all from single model. Additionally, it can be applied to SCF acceleration. In contrast, the Equiformer V2 regression model constrained to predicting a narrow set of specific properties, necessitating the training of new model for each new property.
5.4 EFFICIENCY EVALUATION OF WANET
292.673
WANet exhibits efficiency advantages in two aspects: (1) its ap- plication to SCF relative to traditional DFT calculations, and (2) its efficiency and resource usage compared to existing state-of- the-art neural networks. Specifically, WANet can predict Hamil- tonians for large molecular systems significantly faster than tra- ditional DFT methods. Fi presents a wall-clock time com- parison between WANet-augmented SCF and traditional DFT calculations. Although the neural network evaluation introduces small overhead, WANet substantially reduces the number of required SCF iterations, resulting in a faster overall computation time. This notable speed-up makes WANet particularly advanta- geous for applications requiring rapid predictions for large molec- ular systems, such tighthhhouuhhut virtual screening. Fur- thermore, WANet outperforms QHNet in training and inference efficiency on the PubChemQH dataset, offering faster training times, improved inference speeds, and lower peak GPU memory usage (Figure 3b).
Time
200
5,100
Total Time NH Inference Time SCF
Figure 3(a): Wall-clock comparison of Augmenedde DFT with tra- ditional SCF iterations.
$0.03
q qllse Whizee
75
so-
25-
15:05
(b): Comparison of training and in- ference efficiency and resource usage between QHNet and WANet on the PubChemQH dataset.
5.5
MOLECULAR PROPERTIES BEYOND ENERGY PREDICTIONS
To validate the versatility of the Hamiltonian predicted by WANet with WALoss, we extended our exper-
Table 3: Additional property predictions. We evaluated WANet in
deriving dipole moment and electronic spatial extent.
Model
Dipole Moment MIN (D) Electronic Sputial Extent MAL (2...
Equioomen V2 Regiceson
A.121
0.109S
iments to include predictions beyond energy properties, specifically dipole moment and electronic spatial extent. As shown in Table 3 WANet with WALoss performs competitively in dipole moment prediction and excels in predicting electronic spatial extent. This demonstrates the model's ability to generalize across multiple molecular properties, highlighting its potential for broader quantum mechanical
EGNN Regression
3-12
0:0789
ViSNei Regensiin
4-5211
0.0982
Enccal Guess
4.0770
0.0161
WANet wi WALess
3,3928
0,6076
9

---

Published as a conference paper at ICLR 2025
calculations beyond energy- basey properties. Further details the derivation of these properties from the Hamiltonian are provided in Appendix FF
5.6 SCALABILITY IN ELONGATED CARBON CIAIN
To evaluate the scalability of our model trained on PubChemQH, we conducted inference us- ing elongated alkanes (2, series of saturated hydrocarbons. We compared three mod- els: our model with WALoss, our model without WALoss, and an initial guess algorithm. The results, shown in Fig- ure 4 demonstrate that our model with WALoss achieves enhanced perfor- mance in predicting LUMO and HOMO energies, partic- ularly in the "D2" region, where the atom count ex- ceeds the range of the Pub- ChemQH dataset. Notably, our model with WALoss also performs well on elon- gated alkanes with up to 182 atoms-thre times the average atom count of the PubChemQH training set (60). These findings high- light the effectiveness of WALoss in enhancing the scalability and applicability of our model for predicting electronic properties in scalable homogeneous series, demonstrating its potential for application to larger and more complex molecular systems. Additional analysis of the scaling performance is provided in Appendix
Mumther Alor
Humber Cartern
Number Alers
Momber Garbon Alore
Figure 4: Model performance in predicting HOMO and LUMO energies for elongated alkanes. The left panel shows the MAE for HOMO predictions, while the right panel shows the MAE for LUMO predictions. "DI" indicates that the atom count within the range of the PubChemQH dataset, whereas "D2" indicates that the atom count exceeds this range. The models compared include WANet with WAL.os (w/ WALoss), our model without WALoss (w/o WALoss), and the initial guess (init guess). Notably, our model with WALoss demonstrates superior performance in LUMO predictions and matches the best HOMO performance, particularly the "D2" region.
5.7 ABLATION STUDY ON WALOSS
To evaluate the effectiveness of our proposed WALoss, we conducted an ablation study with three variations: full WALoss, naive WALoss, and WALoss without reweighting. The naive WALoss applies the Frobenius norm to the eigenvalues leveraging backpropagation through eigensolvers, defined as Lnaive n Ei 1 |
where é,,, and 6", are derived from solving generalized eigenvalue problems. The WALoss without reweighting calculates the loss uniformly across all eigenvalues. As shown in Table 4 the full WALoss achieves the lowest MAEs across all metrics. The naive WALoss performs poorly, highlighting several challenges associated with optimizing the naive loss function. Removing reweighting also degrades performance, though not as drastically. Overall, these results validate the design choices in formulating WALoss to improve Hamiltonian prediction.
- ll.
(c)
Table 4: Ablation study of WALoss on the PubChemQH dataset. The best-performing models are highlighted in bold.
Model
Hamilistian MAE
FLONO MAE
C1
SCF
Natre Loss
0.4912
$0.174
55,630
3.534
632.220
456.322
5.36%
13562.7
300% srs
WiLos without Reweighting
0.9973
8:241
7.953
3.983
41.131
21.614
28 28.2615
55 492
WAlee Complste
0.4744
0,7122
0.730
1.327
18835
7,330
48.05%
4F,193
42%
CONCLUSION AND LIMITATIONS
In this work, we introduced WALoss, a loss function designed to improve the accuracy of predicted Hamiltonians. Our experiments demonstrate that incorporating WALoss achieves of-th-of-tte-rr performance by reducing prediction errors and accelerating SCF convergence. Additionally, we introduced new dataset, PubChem@H, and an efficient model, WANer. However, limitations remain, such as the high computational cost of generating large training sets. Despite these challenges, deep learning approaches incorporating WALoss show great promise in advancing computational chemistry and materials science.
Fllongated alkanes (C, H2.+2) are only present in the test
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
Tha authors would like to thank Gaoyuan Wang from Yale University for the discussion on the scaling of the eigenvalues. This work supported by ALW professorship funds to Mark Gerstein.
REFERENCES
Jan Almlöf, Knut Fægri Jr, and Knut Korsell. Principles for direct scf approach to licao-moab-inile calculations. Journal of Computational Chemistry, 3(3):385-399 1982.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
Nathan Argaman and Guy Makov. Density functional theory: An introduction. American Journal of Physies, 68(1):69-79, 2000.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35: 1423-114333 1436, 2022.
Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M Elena, David P Kováes, Janosh Riebesell, Xavier R Advincula, Mark Asta, William J Baldwin, Noam Bernstein, et al. A foundation model for atomistic materials chemistry. arXiv preprint arXiv:2401.000% 2023.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Axel D Beeke. 98(7):5648-66 6, 1993.
thermochemistry,. iii. the role of exact exchange. J. Chem. Phys,
KC Bhamu, Amit Soni, and Jagrati Sahariya. Revealing optoclectronic and transport properties of potential perovskites cs2pdx6 (x= cl, br): a probe from density functional theory (dft). Solar Energy, 162:336-343 2018.
L. C. Blum and J.-L. Reymond. 970 million druglike small molecules for virtual screening in the chemical universe database GDB-13. J. Am. Chem. Soc., 131:8732, 2009.
Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong- Graph networks as universal machine learning framework for molecules and crystals. Chemistry of Materials, 3199):3564-3577 2019.
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference on machine learning, pp. 4057-4086. PMLR, 2022.
Taco Cohen and Max Welling. Group equivariant convolutional networks. Maria Florina Balcan and Kilian Q Weinberger (eds.), Proceedings of The 33rd International Conference Machine Learning. volume 48 of Proceedings of Machine Learning Research, pp. 2990-2999, New York, New York, USA, 20-22 Jun 2016. PMLR. URL atttps://proceedeeee U48/cohenc16, 1ttmbennc6hh..mmll.l..
mlr press/
Andrea Dal Corso, Alfredo Pasquarello, and Alfonso Baldereschi. fensity-functional perturbation theory for lattice dynamics with ultrasoft pseudopotentials. Physical Review B, 56(18):R 1369, 1997.
Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference Machine Learning, pp. 5583-5608. PMLR, 2022.

---

Published as a conference paper at ICLR 2025
Alexandre Duval, Simon V Mathis, Chaitanya K Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D Malliaros, Taco Cohen, Pietro Liò, Yoshua Bengio, and Michael Bronstein. A hitchhiker's guide to geometric gnns for 3d atomic systems. arXiv preprint arX1v:2312.051112000 2023.
Robert Eisberg and Robert Resnick. Quantum physics of atoms, molecules, solids, nuclei, and particles. 1985.
Carrie A Farberow, James A Dumesic, and Manos Mavrikakis. Density functional theory calculations and analysis of reaction pathways for reduction of nitric oxide by hydrogen on (111). Acs Catalysis, 4(10):33073331,, 2014.
Thorben Frank, Oliver Unke, and Roburt Müller. So3krates: Equivariant attention for interac- tions arbitrary length-scales in molecular systems. Advances in Neural Information Processing Systems, 35:29400-22411 2022.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling, Se (3)-transformers:: 3d roto- translation equivariant attention networks. Advances in neural information processing systems, 33: 1970-1981, 2020.
Michael Gastegger, Adam McSloy, Mathis Luya, T. Schütt, and R. Maurer. A deep neural network for molecular wave functions in quasi-atomic minimal basis representation. The Journal of Chemical Physics, 134(4)334123, 7 2020. ISSN 0021-9606. doi: 10. 10.106//5.001211. URL
https://doi... www.:oo.orgg/// 10010..00060001011110110000
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03/23 2020.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34: 6790-6802, 2021.
Benyamin Ghojogh, Fakhri Karray, and Mark Crowley, Eigenvalue and generalized eigenvalue problems: Tutorial. arXiv preprint arXiv:1903.11240 2019.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013.
Ganesh Hegde and R. Chris Bowen. Machine-leamed approximations to density functional theory hamiltonians. Scientific Reports, 7(1):42669, 2017. doi: 10. 10.1038/srep42669. URL https
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-neural-nettorkk solution of the electronic schrödinger equation. Nature Chemistry, 12(10);811-897, 2020. doi: 10.1030/44157-000005444y URL http:////060..1g1100000/////////
Pierre Hohenberg and Walter Kohn. Inhomogeneous electron gas. Physical review, 136(3B):B864, 1964.
Ida-Marie Høyvik. The spectrum of the atomic orbital overlap matrix and the locality of the virtual electronic density matrix. Molecular Physics, 118:c1765034, 06 2020. doi: 10.1080/0026896.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzaioo with gumbel-softmax arXiv preprint 161 1.01 144, 2016.
Frank Jensen. Introduction to computational chemistry. John wiley & sons, 2017.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.0444 04088, 2024.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2025
R. O. Jones. Density functional theory: Its origins, rise prominence, and future. Rev. Mod. Phys.. 87:897-922 923, Aug 2015. doi: 10.100//Rw URL http:::ttt//ttt /ttn:k aws.kk.aps......... /rg/ doi/10.1103
87.897
Fusong Ju, Xinran Wei, Lin Huang, Andrew J. Jenkins, Leo Xia, Jia Zhang, Jianwei Zhu, Han Yang, Bin Shao, Peggy Dai, Ashwin Mayya, Zahra Hooshmand, Alexandra Efimovskaya, Nathan A. Baker, Matthias Troyer, and Hongbin Liu. Acceleration without disruption: Dft software-as--a service. 2024.
John Jumper; Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583
2021.
Kuzma Khrabrov, Ilya Shenbin, Alexander Ryabov, Artem Tsypin, Alexander Telepov, Anton
Alekseev, Alexander Grishin, Pavel Strashnov, Petr Zhilyaev, Sergey Nikolenko, and Artur Kadurin. nabladft: Large-scale scale conformational energy and hamiltonian prediction benchmark and dataset. Phys. Chem. Chem. Phys., 24:25853 2022. doi: 10.1039/DCCP03966.. URL http: 7/dx. doi. www.:ooo..gggoo..o 10. 1039 0D0/3/992CPP396655
Kuzma Khrabrov, Anton Ber, Artem Tsypin, Konstantin Ushenin, Egor Rumiantsev, Alexander Telepov, Dmitry Protasov, Ilya Shenbin, Anton Alekseev, Mikhail Shirokikh, Sergey Nikolenko, Elena Tutubalina, and Artur Kadurin. dft:: A universal quantum chemistry dataset of drug-like molecules and benchmark for neural network potentials, 2024. URL URL httpt::////// httpixxivvorg///// /
abs/ foo//406... 4347
Dmitrii Kochkov, Tobias Pfaff, Alvaro Sanchez-Gonzalez, Peter Battaglia, and Bryan K. Clark. Learning ground states of quantum hamiltonians with graph networks. 2021.
Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A1133) 133, 1965.
Walter Kohn, Axel D Becke, and Robert G Parr. Density functional theory of electronic structure. The journal of physical chemistry, 100(31): 19974-11974-,2980, 1996.
Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2747 2755-2755. PMLR, 10-15 Jul 2018. URL https: //tpsee///poccedingg mlr. pwwps/esssseee. v www.koondooooo.ooo html
Dávid Péter Kovács, Ilyes Batatia, Eszter Sara Arany, and Gabor Csanyi. Evaluation of the mace force field architecture: From medicinal chemistry materials science. The Journal of Chemical Physics, 159(4), 2023.
Chengteh Lee, Weitao Yang, and Robert G Parr. Development of the colle-salvetti correlation-eneryy formula into a functional of the electron density. Physicul review B, 37(2)7885 1988.
Ira N Levine, Daryle H Busch, and Harrison Shull. Quantum chemistry, volume 6. Pearson Prentice Hall Upper Saddle River, NJ, 2009.
He Li, Zun Wang, Nianlong Zou, Meng Ye, Runzhang Xu, Xiaoxun Gong, Wenhui Duan, and Yong Xu. Deep-learning density functional theory hamiltonian for efficient ab initio electronic- structure calculation. Nature Computational Science, 266)3367777 2022. doi: 1038/ URLpttts:///oi.ooo org. 1038 /3////43558888 02205588882222020-50000022
083588-022-002660
Yunyang Li, Yusong Wang, Lin Huang, Han Yang, Xinran Wei, Jia Zhang, Tong Wang, Zun Wang, Bin Shao, and Yin Liu. Long-short-range cesssag-passingg physics-informed framework to capture non-local interaction for scalable molecular dynamics simulation, 2023.
Yunyang Li, Lin Huang, Zhihao Ding, Chu Wang. Xinran Wei, Han Yang, Zun Wang, Chang Liu, Yu Shi, Peiran Jin, et al. E2former: A linear-time efficient and equivariant transformer for scalable molecular modeling. arXiv e-prints, pp. arXiv-2501, 2025.
13

---

Published as a conference paper at ICLR 2025
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990) 2022.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic
graphs. In International Conference on Learning Representations, 2023. URL https://
openreview net/ forum Kdw fARgOTD
Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. arxiv preprint arxiv:2306. 12059, 2023.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representations (ICLR), 2022.
Shuqi Lu, Zhifeng Gao, Di He, Linfeng Zhang, and Guolin Ke. Highly accurate quantum chemical property prediction with uni-mol+, arXiv preprint arXiv:2303. 2208822003.6682, 2023.
Erpai Luo, Xinran Wei, Lin Huang, Yunyang Li, Han Yang, Zun Wang, Chang Liu, Zaishuo Xia, Jia Zhang, and Bin Shao. Efficient and scalable density functional theory hamiltonian prediction through adaptive sparsity. arXiv preprint drXiv:2502.0117, 2025.
Norman H March. Electron Correlations in the Solid State. World Scientific Publishing Company, 1999.
Richard M Martin. Electronic structure: basic theory and practical methods. Cambridge university press, 2020.
Grégoire Montavon, Matthias Rupp. Vivekanand Gobre, Alvaro Vazquez-Mayagoitaa, Katja Hansen, Alexandre Tkatchenko, Klaus-Robert Müller, and O Anatole von Lilienfeld. Machine learning of molecular electronic properties in chemical compound space. New Journal of Physics, 15(9):
iop www.ppoogg/3331 095003, 2013. helLhttp////aaeee 2ttp...r//111122333////--------3 200111222011111--0000000
Albert Musaclian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kornbluth, and Boris Kozinsky, Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Maho Nakata and Toshiyuki Maeda. Pubchemqe b3lyp/6-31g*//pmd dataset: the electronic structures of 86 million molecules using b31yp/6-31g" calculations. arXiv preprint Stti::2300./8455. 2023.
Frank Neese. Prediction of molecular properties and molecular spectroscopy with density functional theory: From fundamental theory exchange-coupling. Coordination Chemistry Reviews, 253 (5-6):526-563, 2009.
Jörg Neugebauer and Tilmann Hickel. Density functional theory in materials science. Wiley Interdisciplinary Reviews: Computational Molecular Science, 3(5)438 448, 2013.
Maylis Orio, Dimitrios A Pantazis, and Frank Neese. Density functional theory. Photosyuthesis research, 102:443-553, 2009.
Robert G Parr and Weitao Yang. Density-functional theory of the electronic structure of molecules. Annual review of physical chemistry. 46(1):701-728, 1995.
Saro Passaro and C Lawrence Zimick. Reducing so (3) convolutions to so (2) for efficient equivariant gnns. In International Conference on Machine Learning, pp. 27420-27433 PMLR, 2023.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583-8595, 2021.
Soumya Sanyal, Janakiraman Balachandran, Naganand Yadati, Abhishek Kumar, Padmini Ra- jagopalan, Suchismita Sanyal, and Partha Talukdar. Mt-egenn: Integrating crystal graph convo- lutional neural network with multitask learning for material property prediction. arXiv preprint arXiv: 18/ I. .55660 2018.
14

---

Published as a conference paper at ICLR 2025
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323- 9332. PMLR, 2021.
Kristof T Schütt, Huziel Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018.
K. T Schütt, Michael Gastegger, Alexandre Tkatchenko, K. R. Müller, and R. J. Maurer. Unifying machine learning and quantum chemistry with a deep neural network for molecular wavefunctions. doi: 10. 00.10088/41447-019-111
Nature Communications, 10(1 ):024,, 2019.
|RLhttps
//doi.org/10. 1ttp:/000.003//41447777 019 12875-222722
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 1at.v:170..6533.. 06538,, 2017.
Guillem Simeon and Gianni De Fabritiis. Tensoreet: Cartesian tensor representations for efficient learning of molecular potentials. Advances in Neural Information Processing Systems, 36, 2024.
Alain St.-Amant, Wendy D Cornell, Peter Kollman, and Thomas A Halgren. Calculation of molec- ular geometries, relative conformational energies, dipole moments, and molecular electrostatic potential fitted charges small organic molecules biochemical interest by density functional theory. Journal of computational chemistry, 16(12): 1483)-141111100 1506, 1995.
Philip Stephens, Frank J Devlin, Cary F Chabalowski, and Michael Frisch. Ab initio calculation of vibrational absorption and circular dichroism spectra using density functional force fields. The 1621- 1167,, 1994.
Journal of physical chemistry, 98(45): (144))):113311111
Qiming Sun. Libeint: An efficient general integral library for aussian basis functions. Journal of computational chemistry, 36622))1664-1177)) 2015.
Qiming Sun, Timothy C Berkelbach, Nick S Blunt, George H Booth, Sheng Guo, Zhendong Li, Junzi Liu, James D McClain, Elvira R Sayfutyarova, Sandeep Sharma, al. Pysef: the python-based simulations of chemistry framework. Wiley Interdisciplinary Reviews: Computational Molecular Science, (())ee3340 2018
Qiming Sun, Xing Zhang, Samragni Banerjee, Peng Bao, Marc Barbry, Nick S Blunt, Nikolay A Bogdanov, George H Booth, Jia Chen, Zhi-Hao Cui, et al. Recent developments in the pysef program package. The Journal of chemical physics, 153(2), 2020.
Attila Szabo and Neil S Ostlund. Modern quantum chemistry: introduction to advanced electronic structure theory. Courier Corporation, 2012.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: cota translation translationeeuuivrraant neural networks for 3d point clouds. arXiv preprint arXiv: wtt.::8802.821111 2018
Julian Tirado-Rive Rives and William L Jorgensen. Performance of b3lyp density functional methods for large set of organic molecules. Journal of chemical theory and computation, 4(2))297-306, 2008.
Artur Toshev, Gianluca Galletti, Johannes Brandstetter, Stefan Adami, and Nikolaus A Adams. E (3) equivariant graph neural networks for particle-based fluid mechanics. arXiv preprint arXiv:2304.00 05000 2023.
Oliver Thorsten Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus Robert Muller. SE(3)-equivariant prediction of molecular wavefunctions and electronic densities. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021a. URL id auGY2U@Ehhu
Oliver Thorsten Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus Robert Muller. SE(3)-equivariant prediction of molecular wavefunctions and electronic densities. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021b. URL TRRLattss://ppeeeeeeeet http////pppp/eeeeeeeepepee net fwrrmormmm Id=auGY2UQFhSu
15

---

Published as a conference paper at ICLR 2025
JH Van Lenthe, R Zwaans, Huub JJ Van Dam, and MF Guest. Starting scf calculations by superposi- tion of atomic densities. Journal of computational chemistry, 27(8)) 922,2, 2006.
Tanja Van Mourik, Michael Bühl, and Marie-Pierre Gaigeot. Density functional theory across chemistry, physies and biology, 2014.
Seymour Vosko, Leslie Wilk, and Marwan Nusair. Accurate spin-dependent electron liquid correlation energies for local spin density calculations: a critical analysis. Canadian Journal of physics, 58(8): 1200:1220111111 1980.
Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. Comenet: Towards complete and efficient message passing for 3d molecular graphs. Advances in Neural Information Processing Systems, 35:650-666 2022.
Yi Wang, Mingqing Liao, Brandon J Bocklund, Peng Gao, Shun-Li Shang, Hojong Kim, Allison M Beese, Long- -ing Chen, and Zi-Kui Liu. Dittk: Density functional theory toolkit for high- throughput lattice dynamics calculations. Calphad, 75:102355, 2021.
Tim M Watson and Jonathan D Hirst. Density functional theory vibrational frequencies of amides and amide dimers. The Journal of Physical Chemistry A, 106(34)7858-786, 2002.
J Beau W Webber. A bi-symmetric log transformation for wide-range data. Measurement Science and Technology. 24(2):02700,, 2012
Florian Weigend, Accurate coulomb-fiting basis sets for h to m. Physical chemistry chemical physics, 8(9):1057-1065, 2006.
Florian Weigend and Reinhart Ahlrichs. Balanced basis sets of split valence, triple zeta valence and quadruple zeta valence quality for h to m:: Design and assessment of accuracy. Physical Chemistry Chemical Physics, 7(18):3297-3305, 2005.
J.H. Wilkinson. The Algebraic Eigenvalue Problem. Monographs on numerical analysis. Clarendon Press, 1988. ISBN 9780198534181. URL https: //books google. wwt.:ccccbbbo..c
5wsK1OP 7UF
Xiaohong, Wang Zhen, Gao Zhifeng, Zheng Hang, Zhang Linfeng, Ke Guolin, and E Weinan. Uni-mol2: Exploring molecular pretraining model at scale. arXiv preprint arXiv:2406.1499 14969, 2024.
Shi Yin, Xinyang Pan, Xudong Zhu, Tianyu Gao, Haochong Zhang, Feng Wu, and Lixin He. Harmonizing -qqivvarnne with neural expressiveness: hybrid deep learning framework oriented to the prediction of electronic structure hamiltonian. 2024.
Haiyang Yu, Meng Liu, Youzhi Luo, Alex Strasser, Xiofeng Qian, Xioning Qian, and Shuiwang Ji. Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules. In Advances in Neural Information Processing Systems, 2023a. URL http://aaaww..oo/////////666600
pdf
Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Efficient and equivariant graph networks for predicting quantum hamiltonian. In International Conference Machine Learning, pp. 40412-40424. PMLR, 2023b.
He Zhang, Chang Liu, Zun Wang. Xinran Wei, Siyuan Liu, Nanning Zheng, Bin Shao, and Tie-Yes Yan Liu. Self-consistnncy training for hamiltonian prediction. 2024.
Yang Zhong, Hongyu Yu, Mao Su, Xingao Gong, and Hongjun Xiang. Transferable equivariant graph neural networks for the hamiltonians of molecules and solids. npj Computational Materials, (1)::182, 2023.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models, arXiv preprint 222000200.088066 2022.
16

---

Published as a conference paper at ICLR 2025
Appendix Table of Contents
Notation Summary
18
Additional Related Work
19
Additional Experiments
20
Evaluation Metrics
20
SAD Phenomenon
20
Error Scaling Analysis
21
Overlitting Risk Evaluation
22
C,5 Energy Predictions for Elongated Alkanes
22
C.6
Additional Ablation Study
22
D
SCN convolution
23
Additional Background
26
E.1 Example
29
Dipole Moment and Electronic Spatial Extent
30
Group Theory
31
H
Exensive Details on PubChem(H
32
Extensive Details on WANel
32
J
Additional Theory
37
Proof of Theorem
37
J.2
Perturbation Analysis and Growth of Eigenvalues
39
J.3
Scaling Law of the Smallest Eigenvalue of Overlap Matrix
41
J.3.1
N-Nearest-Neighbor Overlaps
K Training
42
Discussion
44
Discussion on System Energy Error
44
Additional Conformational Energy Analysis
44
Discussion on SCF acceleration ratio
44
17

---

Published as a conference paper at ICLR 2025
A
NOTATION SUMMARY
Table 5: Summary of main notation used the paper.
Notation
Description
M = (2
Molecular system defined by muclear charges Z and positions R
Spatial coordinate electron
(x())
Single-clectron wbitfl/wwvcfuuttionn as function of electron spatial coordinate
p(r)
Electron density as a function of spatial coordinute r hasis
B
hasis functions
represent
Basis set consisting B basis functions (r)
C Rayy
Coefficient matrix where column contuins basis function coefficients orbital
H Ruxa
Hamiltonian
chosen basis representation
e RR
Overlap matrix with elements S.d f (.. (r)))(r)dr
Diagonal matrix containing orbital energies (eigenvalues of the Hamiltonian)
(M(M)
Machine leurning model purameterized 0 for predicting the Hamiltomian II from molecular structure M Wavefunction Alignment Loss function
Owa
é
Predicted and ground truth eigenvalues (orbital energies) of H and H Predicted and ground truth eigenvectors (basis enefficients) of H and H Ground truth Hamiltonian matrix
C.C
H'
M
Predicted transformed Hamiltonian matrix in an orthogonal basis Matrix ohtained from eigen decomposition the overlap matrix Number of occupied orbitals
G
Set of indices corresponding to the & + lowest eigenvalues
Hyperparameter controlling the weight of occupied and unoccupied orbitals in the loss fuunction Number of samples the
experts MoE nodel
D
Dataset used for training the machine learning model Number of electrons in the system / total number of experts in MoE model
N
II
Initial guess for the Hamiltonian matrix
Hamiltonian matrix at the k-th SCF iteration
C(A)
Coefficient matrix at the k-th SCF iteration
(R)
Diagonal matrix of orbital energies at the k-thh SCF iteration Convergence threshold for SCF procedure Degree of the irreducible representation (irrep) in SO(3) equivariant networks Wigner D-matrix representation of group element SO(3) of degree é Real spherical hurmonies of degree / evaluated at unit vector f
(i
(F)
Clebseh- Gordan coefficients coupling irreps of degree
& irrep of degree Eg
Tiv
Vector pointing from node node
Unit vector pointing from node node
Gyul
Total loss function combining Caiggn and mean squared error (MSE) loss
As
Hyperparumeters controlling weights of different terms in Lond
18

---

Published as a conference paper at ICLR 2025
B ADDITIONAL RELATED WORK
Predicting Kohn-Sham Hamiltonians Early work on predicting Kohn-Sham Hamiltonians used
kernel ridge regression (Hegde& & Bowen 2017), while newer approaches use neural networks (NNs).
Some NNs predict the wavefunction itself (Schütt et al.
2019,
Gastegger et al.,
2020
Hermann
2020), while others use equivariant (Unke al.
2021a Kochkov et al., 2021 Yu al.
2023b Li et al. 2022] Zhong et al.) (023) or hybrid architectures (Yin et al. (024) to predict the
molecular Hamiltonian. novel training method has been proposed to address the scarcity of labeled
data (Zhang et a (024), and two benchmark datasets aim to standardize evaluation of molecular
Hamiltonian prediction (Khrabrov al. 2022. Yu et 2023a Khrabrov al. 2024).
Equivariant Graph Neural Networks (EGNNs) It is often desired that machine learning (ML) models exhibit equivariance to rotations, translations, or reflections, which guarantee that they respect certain physical symmetries. Foundational work introduced group equivariant convolutional neural networks (ECNNs) (Cohen & Welling 2016), whose importance was underscored by a proof that equivariance and convolutional structure equivalent given certain ordinary constraints (Kondor & Trivedi 2018). One appealing way to implement convolution is with geometric graph neural networks (geometric GNNs), which apply naturally atomic systems by encoding them as graphs embedded in R Two important families of geometric GNNs are invariant GNNs and Cartesian equivariant GNNs. Over the past several years, invariant GNNs have achieved state-of-the-ee results in predicting proper- ties of molecules, crystals, and other materials (Schütt et al.) (2018) (Sanyal et al. (0188) (Chen et al.
2019) (Gasteiger al. 2020) (Liu et al. 2022 Gasteiger al.[[201]] (Wang et al.] as well as in predicting the folding structure of proteins (Jumper et al. 2021). Cartesian equivariant GNNs have
seen success in similar areas, benefiting from the greater flexibility their representations (Jing et al.)
(024)). Cartesian equivariant 2020) (Satorras et al.) 202T) (Du et al. 2022) (Simeon & De Fabritiis
GNNs have also seen recent innovation in Cartesian equivariant transformer layers (Frank et al. 2022). A third significant family of geometric GNNs is spherical equivariant GNNs, which use spherical tensors rather than Cartesian tensors. As as a result, spherical equivariant GNNs behave more naturally under rotations and avail themselves of many results of the representation theory of SO(3). They have shown dexterity in tasks in geometry, physics, and chemistry (Thomas et al. 2018); modeled dynamic molecular systems (Anderson al. 2019; enabled SO(3)- and SE -equuuaqqraari transformer layers (Fuchs et al.) 2020) (Liao & Smidi 2022): accuarately and efficiently calculated interatomic potentials (Batzner et al. 2022) (Batatia al. [022) (Musaelian et al. 2023); enabled -qquivarrr fluid mechanical modeling (Toshev et al. 2023); and improved efficiency by
reducing certain convolution computations in SO(3) to equivalent ones in (2) (Passaro & Zitnick 2023). Li et al. (2025) propose to use Wiger-6j Convolution to shift computation from edges to nodes.
19

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENTS
C.I EVALUATION METRICS
In this section, we provide a detailed description of the metrics in the main-text:
MAE for Hamiltonian The MAE for Hamiltonian assesses the accuracy of the predicted Hamilto- nian matrices. This metric crucial for evaluating the quality of the predicted electronic structure and its components, which are foundational in quantum chemistry calculations.
MAE for éHооO The MAE for the Highest Occupied Molecular Orbital (HOMO) energy (éномо) evaluates the precision of the predicted HOMO levels. Accurate HOMO predictions are essential as they influence a molecule's 's chemical reactivity and stability.
MAE for ELUMO The MAE for the Lowest Unoccupied Molecular Orbital (LUMO) energy (CLUMO) measures the accuracy of LUMO level predictions. Accurate LUMO predictions are critical for understanding molecule's electron affinity and chemical behavior.
MAE for cΛ The MAE for the energy gap between HOMO and LUMO (EA) assesses the precision of this important property, which determines the electronic properties and conductivity of materials.
MAE for 6ace The MAE for occupied orbital energies (Eace) assesses the accuracy of predicted energies for all occupied molecular orbitals, providing comprehensive measure of how well the model captures the electronic structure.
MAE for 6orb The MAE for all orbital energies (Eorb) measures the discrepancies in predicted energies for both occupied and unoccupied orbitals. This metric evaluates the overall accuracy of the model predicting the entire spectrum of orbital energies.
MAE on Total Energy The MAE on total energy assesses the accuracy of the predicted total energies derived from Hamiltonian matrices using pyscf. This metric is crucial for validating the model's accuracy in predicting the overall energy of the system, which is fundamental for understanding molecular stability and reactions.
Cosine Similarity for prrefuunttion/Eigevveetoo (C) The cosine similarity for wavefunc- tion/eigenvectors (C) measures the similarity between the predicted and actual wavefunctions or eigenvectors of the system. High cosine similarity indicates that the predicted wavefunction distribu- tion closely matches the actual distribution, which is important for accurately modeling electronic properties.
The SCF (Self-Consistent Field) iteration count evaluates the number of iterations required to achieve convergence in DFT (Density Functional Theory) calculations using the predicted Hamiltonian matrices when comparing with the inital guess. Mathmatically, it is defined as:
SCF Iteration
Predicted Hamiltonian Iteractions
Inital Hamiltonian Iteractions
This metric assesses the efficiency of the predicted matrices in expediting the DFT calculations.
C.2 SAD PHENOMENON
We present additional graphical illustration of the SAD phenomenon discussed in the main text, depicting the learning curve of the QHNet model using only elementwise loss on a subser of the PubChem@II dataset 5 that the MAE decreases, the system energy exhibits significant fluctuations. Notably, when the Hamiltonian's 's MAE is around 0.3, the system energy MAE reaches 30,000 kcal/mol. This highlights the monoonio relationship between the Hamiltonian MAE and the resulting system energy. This extreme instability in derived properties, despite seemingly small Hamiltonian MAE, a key characteristic of the SAD phenomenon, hindering the applicability of the predicted Hamiltonian.
"This differs from the dataset used in the maintexts.
20

---

Published as a conference paper at ICLR 2025
Hami-MAE vs Step
System Energy vs Stop
230
System Energy
Hamt-MAA MAE
2.55
2000
240
0.99
5.00
1000
1150
Step
Step
Figure 5: An additional graphical illustration of the SAD phenomenon discussed in the main text, depicting the learning curve of the QHNet model using only elementwise loss. The y-axis represents the metric, and the x-axis represents the steps. The figure shows that as the MAE decreases, the system energy exhibits significant fluctuations.
Distribution of the Smallest Eigenval for
PubchemQH
QH9
10a
Value (1/e) (Log Scale)
10
102
PubChemOH
QH9
Figure 6: the distribution of (1
on QH1 and PubChemQI1 dataset
C.3 ERROR SCALING ANALYSIS
In Figure and Table 6 we analyzed the error trend scaling for "intensive" properties such as еμомо and (LUMO with an increasing number of electrons. This challenge, widely recognized in the field (Yu et al., 2023a Zhang et al.) 2024). highlights a difficult out-of-distribution (OOD) generalization scenario: models trained on smaller molecules often struggle to generalize to larger ones, limiting their applicability. Previous state-of-the-art models have exhibited significant scaling errors when
extrapolating to larger systems (Yu et al. 2023a Zhang et al. 2024).
We conducted eиrоr scaling analysis for saturated carbon chains to assess the scalability of our model. The error scaling coefficients for the HOMO and LUMO energies, as well as the energy gap, presented in Table 6 Our results show significantly lower error scaling compared to baseline models for saturated carbon chains. Our model substantially outperforms the one without WALoss, as reflected by the scaling coefficients in Table R2, marking considerable improvement in extrapolation capabilities within the field.
We further analyze the smallest eigenvalue distribution which now included in Figure
21

---

Published as a conference paper at ICLR 2025
Table 6: Scaling Coefficients HOMO, LUMO, and Energy Gap with Carbon Atom Count
Model
HOMO Scaling Coeffciinnt
LUMO Scaling Coeliciiett
Gap Scaling Coefficients
WAt (with WALoss)
0.4211
0.4003
0.0206
Initial Guess
0.4221
0,4006
0.0214
Without WALoss
1.3639
2.7646
1.4007
C.4 OVERFITTING RISK EVALUATION
We evaluated the risk of overfitting the PubChemQH dataset. As shown in Table 7 the evaluation metrics on unseen data are comparable to those on the training set, indicating that our model generalizes well and does not overfit.
Table 7: Overfitting Risk Evaluation
PubChem
Hamihonian MAE
CROMO MAE
MAE
MAE
Ga
System Energy MAE
Training
0.4736
3.753
17,850
6.159
46.036
Test MAP
0.4744
0.7122
0.730
1.327
18,835
7,330
48.03%
47.193
C.5 ENERGY PREDICTIONS FOR ELONGATED ALKANES
We have incorporated additional plots (Figure 7 that illustrate the predicted and actual values for both LUMO and HOMO energies, as well as the HOMO-LUMO gap. As the number carbon atoms increases, both HOMO and LUMO energies exhibit an initial rise followed by a plateau with slight increases. Our model, utilizing the WALoss method, effectively captures this trend, outperforming baseline methods. Notably, Equiformer V2's HOMO predictions fail to generalize to saturated carbon systems, as these systems were not included in the training set. This observation underscores the advantage of Hamiltonian-basse models in capturing physical principles and achieving better generalization compared to property regression models.
For isolated molecular systems, such those studied here, computational chemists primarily concerned with per-electron energy levels. To demonstrate this, we plot the energy levels within reasonable window centered around the HOMO level and compare these with ground truth values obtained from DFT calculations. As shown in Figure 7 A. our model closely replicates the ground truth energy levels, outperforming baseline models. These results highlight the effectiveness of our approach in accurately predicting the electronic structure of molecules.
C.6 ADDITIONAL. ABLATION STUDY
In contrast to prior research, our study addresses the substantial challenges presented by the scale of the PubChemQC t-zvp dataset, particularly in providing the model with strong numeric starting point. To overcome this difficulty, we shifted our focus to a more tractable objective: making predictions based on an easily obtainable initial guess. Our ablation study shows that while predictions based on the initial guess significantly improve performance in Hamiltonian prediction, they struggle when applied to the prediction of physical properties.
Table 8: Performance Comparison of Baseline and Initial Guess Models on PubChemQH Dataset. The best models are bolded.
Modi
Hanioouse MABI
MARC
MAR:
Seetiatty
miate SCV
WANes Inaiaal Gazz
$4.107
56.33
3.65
A1k
481.15
4-22
90376.033
3009 47.093
WiNee Teltal Ganui WAl.o.
67M0
1.37
18A35
7,500
#15
To evaluate the contributions of each component in the WANet model, we conducted ablation studies. The results are summarized in the Table g The ablation studies confirm the importance of each architectural component in the WANet model.
22

---

Published as a conference paper at ICLR 2025
A
B
LUMO Number Sorted
H010 of
HOH0
205
Energy (Aligned HOMO)
Marber Coloon Steess
C
Energy
-10
11:
-20
Enagy/ 540
-25
Ereney Sao
WANet
Ground Truth w/o WA Losa
Model
Numberal Carbon
Figure 7: (A) The energy level for saturated hydrocarbon system Ha2) The energy is centered HOMO for comparison. (B) The predicted energy for differnet models on the satured hydrocarbon system. The y-axis indicates the energy in kcal/mol. The x-axis denotes the number of the carbon atoms in the elongated carbon-chains. v2 indicates Equiformer V2. 'wanet' indicates WANet with WALoss. 'groundtruth' indicates DFT calculations. 'init' indicates Fock matrix initialization algorithm using minao. 'without" indicates WANet without WALoss. (C) The HOMO-LUMO gap prediction for different models.
Table 9: Ablation study results for the WANet model. architectural components.
The table shows the impact of different
w/ SOC)) 1.R--MoB
Many Body
Inference (its))
GPU Memory
Hamilionian <HONO MAE.
CLuszo MAE. L
fa MAE
1.34
13.87
0.5895
15.39
27.50
4.651
L.13
12.80
0.4853
2.27
4.01
2.906
4
0.51
24.17
0,4792
0.75
0.73
1.594
15.86
0.4744
0.71
D ESCN CONVOLUTION
The equivariant Spherical Channel Network (cSCN) is graph neural network whose approach to SO(3) -equivariant convolutions significantly reduces their computational burden. To see this, let's brush up on the notation used the formula for the f. degree the message m, from source
node to target node t in SO(3) convolution: We, & is a learnable weight; z (8.) is the 4-tth degree of the irrep feature node s, the source node; the tensor product; (f(e)) is the fy degree spherical harmonic projection of fıs; and Prs is the (normalized) relative position vector. Then, the message is given by
(Ew) => Werey (a.(2) mia
(2. 60 () (Co)
23

---

Published as a conference paper at ICLR 2025
Traditionally, this involves ordinary (3(3) tensor product, which is decomposed into its inreps using the Trebsch-Gordor coefficients, cleurma)
(Pa)
Σ
We, -
E Th mi
(o(tt)) mas (tt
(r (x(())
mis
The tensor product, 8, has been exchanged for the proper Clebsch-Gordon coefficient. The direct sum, D included because the space tensors of degree Eo has a basis indexed by mag.
However, the full SO(3) tensor product is intpuseeintensive ( where T. is the tensor degree. In practice, this means that only tensors up to degree 2 or 3 are used, which is especially unfortunate because higher-order tensors allow more precise representation of angular information. To lessen this
computational burden, eSCN rotates the irrep features, (:(()) of node 2s, along with the relative position vector, fta. by a rotation chosen so that the relative position vector aligns with the y-axis. Thus, may be interpreted that eSCN changes the irreps into a more convenient basis, computes the convolution, and changes back into the original basis. This requires multiplication by change- of-basis Wigner D-matrix before and its inverse after, but is worthwhile because reduces the convolution O(L3) by sparsifying the Crebsch-Gordon coefficients. In particular, eSCN guarantees that the output coefficient is nonzero only if both m, im, and m = 0 hold. The output tensors of some order, therefore, are linear combinations of input tensors of that order. It is not necessary to perform tensor multiplication and decomposition. The non-zero Sroschh-Grrdon coefficients can c(fo.mm) (fy.
then be denoted (faim))
(fi,m),(t,.0) (k) with eSCN further guaranteeing that ml....m (f),,))
Clla--m) and (e, ()) -C(2, (f:.-m),(E,.)) Thus, eSCN guarantees that all but a few Clebsch-Gordon coefficients will be zero, and provides simple formulas involving those that remain.
(.fo.m)
More formally, let R be a rotation matrix chosen so that R fts = (0, 1), and let D( be a Wigner D-matrix representation of R of degree (which would more properly read (R), R), but which is truncated for readability). Then, the message be written
We (D (D & ( (R. (c))
(1w) (D(E)))) Σ mis
E.iks
Rewriting the tensor products using the Clebsch-Gordon coefficients yields
mia
(f.)
= (D())))
Σ
We f.
+ Σ (D()))(())
m: Ch (t) my)
(y (Y
Pts))
Exely
my my
Since the Mlebsch-Gordon coefficients are non-zero only when my O, there is need to my:
(Pw) mis
= (D(6)))-'
Σ We +
+Σ (D(t))(()))
mllo.ma) (limi))te (Y())) f.)))
Eids
Me m.
The rotation R has been chosen so that its yields simple result, so the spherical harmonic term drops out:
(Po) mas = (D(6))) -1
Σ Wene
+ +Σ
(DIG)
claumoo (emmailt,,o))
Ei.ly
me m.
Now, since of the rotations on the righthand side has already dropped out, make things simpler, define ((()) = ((t:): (e.)) yielding
m4a
(E.) (D(
Σ
We +
W
Σ (x(fe))
titts
mo mi;
Since the Clebsch-Gordon coefficients are non-zero only when m, -m. the second summation can be omitted:
m
(tw) (D((a))-1
Σ
W We
+
(a((e))
mo) (t.mm)))ty.o)
+ (z!(1))
mlo. -m, (t.-mm))(tx.o))
my
E,
24

---

Published as a conference paper at ICLR 2025
Now, to avoid the need to over Ey, rather than learning parameters & eSCN leams parameters Wo(lifo) defined for m 2 0 as
W Σ
Wenee Cle mllomm (e.m))
Σ
Weide Cla (f (f,, m),,m),0).0)
2,
and for m as
Wl(het)) Σ
Wenteres Wests (ta (L--m) = Westys clo--m) (lm))
(lom)
E
There exists a linear bijection between W and W, so this parameterization loses no information. Finally, defining
(vie ((id m.
-W(eeba) (x.5) -ma
m 0:
mo
-mo-
(ues (l(,fo))
Willffoo (1²:)
w/lee.oo) (5 -ma
+
m < 0;
m.
mo.
(yle (c),60)
Weleefo) (Fth) )
0,
me
the message equation can very concisely be written -D E,
m to
Σ
(yes (1)
mo
There is another way interpret this that is perhaps more intuitive. Fixing the direction of the relative position vector, Fee leaves single rotational degree of freedom: the roll rotation about this axis. Thus, eSCN reduces (3) convolution to (2) convolution. More formally, define the colatitude angle 0 [0, x| and longitudinal angle ó E [0, 2.T]. Now, using Legendre polynomials p (0), which depend only on the colatitude angle, 0, the real spherical harmonic basis functions can be written
m (0, Φ) p(() im)
Aligning the relative position vector with the y-axis fixes 6, leaving behind basis functions of the form eime which are the circular harmonic basis functions, used in SO(2) convolution. A convolution about á can take advantage of the Convolution Theorem, reducing convolution to point-wise multiplication, further harvesting efficiency gains.
25

---

Published as a conference paper at ICLR 2025
E ADDITIONAL BACKGROUND
Quantum mechanics is most often approached as the study of the Schrödinger equation, a linear partial differential equation whose solutions are called wavefunctions. In this paper, as is typical in molecular modeling, the time Lime-independen Schrödinger equation will be used:
H Ev,
where H. the Hamiltonian operator, corresponds to the total energy of the system, E' is an eigenvalue of corresponding to the energy of 4h, and an eigenfunction of H, also called a wavefunction, or solution to the Schrödinger equation.
As alluded to, the Hamiltonian operator contains all the information regarding the kinetic and potential energies for all particles of system. Thus,
Zo
H
MA
TiA
i=1
N ΣΣ
M
ZAZB
RAB
A=1
B>A
where the first summation is due to the kinetic energy of the electrons, the second summation is due to the kinetic energy of the nuclei, the first double summation due to the potential of the attraction between electrons and nuclei, the second double summation is due to the potential of the repulsion between electrons and electrons, and the third double summation is due to the potential of the repulsion between nuclei.
This formula is unwieldy, however, and it can simplified significantly without discarding much information that would be useful in chemical prediction. The mass difference between electrons and protons, which the minimum mass difference between electrons and nuclei, more than orders of magnitude. Thus, given the same kinetic energy, electrons will be traveling several times faster than nuclei. From the perspective of the electrons, the nuclei are nearly fixed, and from the perspective of the nuclei, the electrons change their position instantaneously. Therefore, it suffices to consider the positions of the nuclei fixed, setting
1
V2
MA
A=1
nuclear kinetic energy, 0, and
Z Zn
Σ
RAB
A=1 1>A
nuclear repulsive potential, to a constant. This is called the Born- Born-Oppenhemmer Approximation, and leaves behind the so-called electronic Hamiltonian,
Z
Helec
TiA
1>
which can more succinctly be written
Helec T + VNe + Vee.
The electronic Hamiltonian is still solved for its eigenfunctions, giving
Helee V (ri,....
= Frdeee V (ri. rN),
ignoring for simplicity's 's sake electron spin to write Veles as function the positions the electrons only. For readability, the subscripts of Helees Eglec, and Velec are henceforth omitted.
For a system of N electrons, is function of 3NN arguments, one for each spatial dimension of each electron. The Hartree-Fock Approximation, also called the Hartree Product, further simplifies
26

---

Published as a conference paper at ICLR 2025
this equation by decomposing the wavefunction into the product of N wave functions, each of three arguments, corresponding to each electron individually:
r r rw) = thi(() (ri) ... (N.()
is impossible to observe the wavefunction itself, but the wavefunction can be used to derive the probability of observing the system's 's electrons anywhere in space. In particular, the probability of observing the electrons at positions Fi through ry is the square of the amplitude of the wavefunction with rj through rn input:
(TI.
rN) r....,,N)F
v v (ri,................... , IN V (TI.....
Given the Hartree- Fock Approximation, this equation can converted into an electron density function that sums over the wavefunction of each electron individually,
n(r) = 2 Σ (r)).()))
multiplying by 2 to account for both spin up and spin down, which were neglected previously. This means that n(r) gives the density of electrons point in space, This is function of only inputs, but it contains much of the information that is observable from the full wavefunction, which, recall, is a function of 3NV inputs.
Density Functional Theory (DFT), which allows the electron density function, n(r), to be exploited. rests on fundamental theorems. First, the ground-state energy, Eo. which is the smallest eigenvalue of the Hamiltonian and corresponds to the energy of the lowest-energy wavefunction. is unique functional of the electron density function. Second, the electron density that minimizes the energy of this functional is the true electron density, which means that it is the electron density that corresponds to the full solution of the Schrödinger equation. Therefore, after the problem has been reduced from one of 3.N inputs to one of 3 inputs, significant information about the original problem can still be found.
It is worth examining the first statement in more detail. Recall that functional is a function that maps functions to scalars. For example, F, defined by
Flf(2)) S f(z)dz,
functional that takes in arbitrary real-valued function and gives out its integral from -1 1. Thus, for example, if (xx = x² + then F[](=)] = 33 The first statement, then, holds that there exists functional that uniquely determines the ground state energy of a particular electron density function. It promises that no more information is needed. Written as an equation, Eg = F F(n(r)).
However, this set-up still relies on some means of finding the individual-electron wave function. This is provided by the Kohn-Sham Equation:
! 1 + (i) + (H(()) + Vxc(r) (x) = (r))
The meaning of each term is follows. The value 1 V is kinetic energy. The value V(r) is the potential due to the interaction between the electron and the nuclei. The value VH(r), called the Hartree potential, is defined
Vn(rr)
ST.
and due to Coulomb repulsion between the electron and electron density function, defined by all electrons in the system. This means that the Hartree potential includes the Coulomb repulsion between the electron and itself, since the electron itself is included in the electron density function, n(r). This is unphysical result, and it is of several effects accounted for in the exchange-correlation potential, Vxc(r), which is defined
Vxc(r) = 6Exc (r)
ôn(r)
27

---

Published as a conference paper at ICLR 2025
This is the functional derivative of the exchange-correlann correlation energy with respect to electron density. Exchange-comee coreeltionol potential is due to quantum-chemice chemical effects, and its form not known because the exact form of (xc(r) is not known. The Kohn-Sham Equation can be summarized as
(r) = (1))
It is evident from the definitions of the terms in the Kohn-Sham Equation that the present approach circular. The Kohn-Sham Equation count on the Hartree potential, (r) The Hartree potential counts the electron density function, n(r). The electron density function counts the single- electron wavefunctions, (r). And the single-electron wavefunctions count on the Kohn-Sham Equation.
This is no problem. The following process leverages this circularity to check the soundness of particular n()::
1. Define an initial trial electron density function, (r).
2. Solve the Kohn-Sham Equation using the trial electron density function, Merial Merial to find the single-electron wavefunctions, (r))
3. Calculate the electron density function, rks(s) (r), implied by these single-electron wavefunc- tions, by wks (r) 2 > vi (r))
4. Compare the calculated electron density, riss (r), with the trial electron density, (rria If the two densities are the same, or nearly so, then this electron density is accepted as correct, and can be used to find the ground-state energy, Eg- If not, the trial electron density is updated somehow, and the process repeats.
To make this process simpler and more efficient, it is common to represent the single-electron wavefunetions as linear combinations of some predefined basis set {6) (r)) where B is defined as the cardinality of the basis. Often, the basis set is composed of atomic orbitals, especially Gassian-type orbitals, which are particularly convenient in calculation. The expansion coefficients of these wavefunctions can be organized in matrix C E R BxN, where each column contains the coefficients wavefunction Each wave function can then be recovered as
B
(r( = Σ (ri
a=l
This permits the Kohn-Sham Equation as whole to be written in matrix form. Recall the form HKsw,(r) = (r). Now, it is possible to rewrite vú (r) according to its decomposition in the basis:
B
HKS
Σ (ri =; Σ Carda (r).
0=1
Now, for any (r) in the basis, left left-multip both sides by ()) yields
(T) Σ - (.)
and integrating with respect to yields
B [ Σ (rid of(()); Σ (r)
a=1
a=1
More succinetly, defining (Kss)aa f fo HKsØa = (del Hkslda) and Sga = f @b (óolda)) (da), this is
Σ (HKs) ga C.i = Σ Spc.....
This allows the construction of the matrix H, with elements Hg. defined by Hşa. (HKk) (óg Hks)oa). and the matrix S, with elements Seaa (dalda)) The matrix H represents the
28

---

Published as a conference paper at ICLR 2025
Hamiltonian matrix, which depends the coefficient matrix C and is calculated using method called
whose time complexity is O(B)). The matrix the overlap matrix, which depends the basis set før (r))a= and accounts for its non-orthogonality This means that if the basis set is orthonormal, the identity. The Kohn-Sham Equation in matrix form is therefore
HC = SCe,
or
HCC)C SCc,
making clearer that H. depends on C, where e a diagonal matrix containing the orbital energies. This forms generalized eigenvalue problem, the principal difficulty of which is the dependence of H on C.
As result, traditional DFT uses the the coefficient matrix by approximating the Hamiltonian. Using superseripts to denote the iteration to which each matrix belongs, at each iteration, the Kohn-Sham Equation is
Field method (SCF), which iteratively refines
H( (C(k-1)) (()) = (())
Therefore, (k) is computed using C*-1, and the generalized eigenvalue problem is solved for C(k) and (R)) This process continues until convergence, which is formalized || H[ 8.
The objective of Hamiltonian prediction is to eliminate the need for this computationally expensive SCF iteration by directly estimating the target Hamiltonian, H", for given molecular structure, M. To predict the Hamiltonian, machine learning model (M(M) parameterized by 0, guided by an optimization process defined as
8" = arg min
Σ
dist H H)).
[D]
(M.HA,)D
where 'D is the dataset, |D] is its cardinality, and dist(., is predefined metric. Machine learning models hold the potential drastically improve the efficiency of DFT calculations without sacrificing accuracy.
E. EXAMPLE
In the example section, we explore the molccular structure and basis set details of water (H₂O), which consists of one oxygen and two hydrogen atoms. We expand each electron's wavefunctions using a basis set, which, in practice, means each atom is represented using predefined basis set that collectively describes set of electrons. Specifically, oxygen is described by five basis functions (1s, 2s, 2p 2p,, 2p2), and each hydrogen atom has two basis functions (Is for each). Altogether, this totals seven basis functions.
the STO-3G minimal basis set, where
2py
2s
Iso
2pz0
2p.0
1881
Iso
Fii
Fi2
F13
Fi4
Fi5
Fi6
F17
2s0
Fz1
Fza
F1
Fat
Fas
Fan
Fsr
Fa3
F34
Fas
Fac
Fst
2p,
2px
2ps
F41
F42
F48
F44
F45
F46
F47
2p:
Fsi
Fsa
Fsa
Fs4
Fes
Fas
Fst
F6a
F64
F65
Feo
FGt
1882
F71
F7a
FTB
F74
FTS
Fro
Fyr
H
H
Is
Is
Figure 8: Molecular structure and basis functions for water (H₂O).
When considering the effect of rotations on the Hamiltonian matrix, is important to understand how the matrix elements transform under such operations.
29

---

Published as a conference paper at ICLR 2025
For instance, if consider rotation that affects only the 2pz, 2py, and 2p, orbitals of the oxygen atom, the corresponding submatrix of the Hamiltonian Hp within the 2p orbital space would transform
- R,,,H,
where Hp is the submatrix containing the interactions between the 2p,, 2p and 2p, orbitals. The rotation matrix Rp specific to the rotation in the 2p orbital space.
1so
280
2p
2p
2pz
18
18H2
180
F11
F1
F
P14
P55
F
F16
F17
280
F21
F22
F23
F24
F25
F26
F27
2pz
F31
F32
F38
F34
Fi5
F36
F37
2p
F41
F42
F4
P44
F45
F46
FA7
2ps
F51
F52
F53
P64
F55
F66
Fsr
18
F61
F62
F63
F84
F65
F66
F67
18
F71
F72
F73
P74
F75
F76
F77
2p,
2p
2pz
2p
33
24
35
2p
F
F41
F45
2p
53
1554
155
This transformation ensures that the physical properties described by the Hamiltonian remain con- sistent under rotational operations, fundamental requirement for accurately modeling molecular systems.
F DIPOLE MOMENT AND ELECTRONIC SPATIAL EXTENT
In this section, describe how to compute the dipole moment and the electronic spatial extent using the Kohn-Sham orbitals derived from the Kohn-Sham Hamiltonian.
The Kohn-Sham equations defined as
ht2
22 + Ven(r) vi(r) = (,))
2m
where (r) denotes the Kohn-Sham orbitals, are the orbital energies, and Vof(r) is the effective potential. The electron density then expressed
DCC
p(r)
Σ (2)
The dipole moment d is computed as the sum of electronic and nuclear contributions. The electronic contribution is given by
dloc
(rp(r)
(6)
and the nuclear contribution
nic - Σ ZARA.
(7)
where represents the elementary charge, ZA denotes the atomic number nucleus A, and R is the position vector of nucleus A. The total dipole moment therefore the sum of these components:
30

---

Published as a conference paper at ICLR 2025
d d + d
(8)
Next, we define the electronic spatial extent, denoted by (r²), which provides measure of the spatial
distribution of the electron density. It is calculated as
(r2) = f p()) dr.
(9)
To compute these quantities in practice, one often works basis set representation. The electron density matrix Pnu is defined
PAu = 2 2 CipCis
(10)
where Cij. the coefficients of the molecular orbitals in terms of the basis functions &,. The dipole integrals paa for Cartesian direction a are expressed as
P =
(11)
The electronic contribution to dipole moment the basis set representation is then given by
Jelee
Σ
(12)
Similarly, the electronic spatial extent in the basis set representation computed using the integral
(13)
leading to the final expression
(1²)= Pnuktr
(14)
G
GROUP THEORY
If a function is equivariant the action group, it does not matter whether the group acts on the function's 's input or output. More formally, for vector spaces V and W equipped with arbitrary group representations Dy (g) and Dw (g), function f V + W is equivariant to
f(Dv(g)v) = (w(g)f(r)
for all E G and for all v E V. the case of SF((), a function f is equivariant if the output the whether the input put through
same whether the input is slid and rotated, then put through then slid and rotated in the same way. A function f is invariant to group G if
(g =
the identity element in W, for all g G. In the case of SE(3), function is equivariant if it gives the same output no matter how its input is slid and rotated. Equivariance is a fundamental notion in the modeling of physical systems. In the context of this paper, it is desirable that the function that predicts the Hamiltonian matrix be ( eqquivarra,,,,,, reflecting that the molecule's s energetic properties are equivalent if the molecule is rotated or translated.
31

---

Published as a conference paper at ICLR 2025
Group representations instance of the general notion of group homomorphisms. Given group G with group operation and another group H with group operation *, a group homomorphism map p G H such that
p - g2) = p(gi) * p(92).
group homomorphism, then, must preserve the structure the group, which means that it does not matter whether the group operation is performed in G H. Note, however, that homomorphism might respect the group structure only trivially. For example, G + H, defined by p(g) = e, is trivial group homomorphism.
group representation simply group homomorphism where H = V is some vector space. In plainer language, this means that group representation is group written as set matrices, whose group operation is matrix multiplication. always the case that V C GL, since any matrix that is degenerate or non-square lacks inverse and therefore fails to satisfy the inverse axiom.
Two group representations D and D' are equivalent if there is fixed matrix P such that D(g) p-1 D') for all E G. In this case, D and D' can be interpreted as the same representation defined with respect to different bases. A representation D is reducible if it acts independent subspaces of V; otherwise, is irreducible. An irreducible representation is called an irrep. More formally, D is reducible if
(g)
D(g) P-1
(g))
P= P= P-P-I(AD ( ((D
which means that it D is block- diagoodia with respect some basis. It is convenient decompose group representations into irreducible representations because this reduces the group operation calculation several smaller independent caleulations. Irreps the atoms of group representations in the sense that arbitrary representations can be composed with the direct sum of irreps. The irreps of (3) are called Wigner D-matrices, with ( (g) denoting a Wigner D-matrix representation of g of degree f. Wigner D-matrices of size (2l 1) (2 + 1), with higher-degree representations allowing for more precise handling of angular information.
H EXENSIVE DETAILS ON PUBCHEMQH
Here, we present a detailed comparison between the PubChemQH dataset and the curated QH9 dataset. This comparison aims to highlight the key differences and similarities. Additionally, we provide a comprehensive analysis of the atom number distribution within the PubChemQH dataset, supported by Figure 9
Table 10: Comparison of PubChemQH and QH9
Feature
PubChemQH
QH9
Source
PubChem Database
QM9
Number of Molecules
50,321
130,831 (QH9-stable)
Functional
B3LYP
B3LYP
Basis Set
Def2TZVP 10-8 8
De 10 10-13 13
SCF Convergence Tolerance
SCF Gradient threshold:
10-4
3.16 10-4
Grid Density Level
3
3
Mean of Node Number
61.85
18
Mean of Hamiltonian Size
1025
141
I EXTENSIVE DETAILS ON WANET
We here provide extensive details on the motivations of the WANet architecture. The WANet architec- ture builds upon vell-established equivariant neural network designs, which have been extensively studied and validated in the field.
32

---

Published as a conference paper at ICLR 2025
Number of atoms in each molecule
4000
Number of molecules
3000
2000
1000
0
30
40
50
60
70
80
90
100
atoms
Figure 9: Node atom distribution for PubChemQH Dataset
We provide a visual representation illustrating WANet's architecture and its key components, which given in FigureTO]. This visual aid will help readers better understand how WANet enhances the scalability of Hamiltonian prediction for large molecular systems. Each architectural element was carefully chosen address specific computational and physical challenges scale:
First, the motivation behind WANet's architecture stems from a critical limitation in existing Hamil- tonian prediction methods- inability scale to larger molecular systems. When using larger basis sets, such as Def2-TZVP, higher-order irreducible representations are required to accurately capture angular dependencies in molecular orbitals. This poses severe computational challenge for traditional SE(3) -Eqiivraaa methods, including QHNet, which become extremely expensive due to their computational complexity. WANet addresses this bottleneck by introducing SO(2) convolutions, which reduce computational complexity from OLL)) to O(L)). This improvement enables WANet to process high tensor degrees efficiently. Without SO(2) convolutions, handling the scale of our PubChemQH dataset and larger systems would be significantly more challenging and computationally expensive.
Second, large molecular systems exhibit fundamentally different physics at various distance scales. As molecular size increases, long-range interactions become more prominent. WANet's Mixture-of- Experts architecture is designed model this complex physics efficiently. It employs specialized experts for different interaction ranges, capturing both short-range effects (like covalent bonding) and long-range phenomena (such electrostatics). By sparsifying these experts, WANet achieves rich representation of molecular interactions while maintaining computational efficiency, making it particularly well-suited for large-scale systems.
Third, WANet's architecture is designed to accurately capture the intrinsic properties of molecular systems, particularly for large molecules. The Hamiltonian matrix, which fully characterizes the quantum state and electron distribution, presents unique challenges in prediction due to complex electron correlation effects as the system size grows. To address this, WANet incorporates the MACE architecture's 's density trick, enabling efficient computation of many-body interactions without explicit calculation 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 precision in Hamiltonian prediction for large systems.
Pair Construction Layer The objective of the pair construction layer is to extend the model's capacity to consider non-diagonal node pairs by introducing a tensor product filter. This filter modulates the projection of irreducible representations onto the space of node pair irreducible representations, denoted by ftx. It is important to note that, in contrast to the Node Convolution Layer, which performs graph convolution radius graph KNN graph, the pair construction
33

---

Published as a conference paper at ICLR 2025
WANet Architecture
90(3) Convolution
Efficient Conv
Convelation
Many- Body
Micture LSR Experts
Marry Body
Challenges for Scalability
Interoction Layer
Expansion Blook
Long-Range Interactions
Imapa
CC Coeft
Sphenical Harmonree
Nondloçonal Bloce
Oiegondl Block
Food sis
slau oou
Assemble
Pradicted Hamiltonian
1.
Mobate
SO(3) Conv
X Y(rx)
WignerO Matrix D(R)
Atom Pair
D(R)
DY(R)
Short Ronge (Bond Langin)
RBF((.)
Van Der Weals
SO(2) Conv
CG
Rotelled
V(ral
Rotated Irrepe
Spherical
Sofmma
my
D(R)h & ((R)
Hormonics
Experts ta
Convolution
ovar Expersd
Density Trick
D*(R)
A +m++ my
Expert
Expert
Espert
Expert
A (mi mm+ m. m.)
DT(R)
Pick Topk Experte
m. + m + m, m²²
2mim + 2mimy 2m-
Trained vit Backsmp
2m,m, 2m3m4 2m;m,
DF(R)
Implicit 3-body term
Repeesentation on
Inapa
Sparse Matiik
Figure 10: Scalable architecture for predicting the Hamiltonian matrix (ft). The framework addresses three key scalability challenges in quantum many-body systems: (1) efficient convolutions, (2) many-body interactions, and (3) long-range interactions. (1) Efficient convolutions are achieved through an S0(3)-equivariant convolution reduced to SO(2) for computational efficiency. (2) Many- body interactions captured using density trick, which implicitly includes three-body terms by combining pairwise interactions quadratically. (3) Long-range interactions are modeled with a mixture of StrruSSttucture Representation (LSR) experts, where softmax-basee gating mechanism allocates pairwise inputs to top-performing experts. Finally, the diagonal and non-diagonal components are assembled into the predicted Hamiltonian matrix, providing scalable and accurate representation of the system dynamics.
34

---

Published as a conference paper at ICLR 2025
Short (Bond Length)
Ven Der Wesls
(BF(r
Exports
Softmax
Gates
Distribution over Experts
Expert n-1
Expert
Expert
Expert
Pick Topk Experts
Trained Backprop
Cutput
Figure 11: illustration of the Mixture of the Song-Short-Rang Pair experts
layer considers all possible node interactions by operating on complete graph. The mathematical formulation of this layer is given as follows:
fiz Ewi.ssht (ol) 8 )
(15)
where u/ and v/? are the 14-th and 1y-th irreducible representations of source node and target node t, respectively, and Wi,,,, are the learned weights that couple these representations into the output representation Lo
To improve the efficiency of the tensor product, we employed the thannel-grouped tensor product, where the channels of the first and second tensors tied to collectively construct the output channel path Additionally, to accommodate the symmetry inherent in the Hamiltonian matrix, we implement symmetric structure in the pair representations. Specifically, the representations for node pairs (s, 1) and (1, 8) must be identical, reflecting the symmetrical nature of physical interactions. This is formulated as:
- 5 + f22),
where ffz and for denote the initial, unsymmetrized tensor products for the node pairs and f before the application of symmetry. The primed notations pto' and fly represent the final, symmetrized outputs.
Graphical Illustration of MoE Here we provide an additional illustration of the Mixture of the Long-Short-Range Pair experts described in the maintext, which is shown Figure II
Expansion Block In the construction of final Hamiltonian blocks that encompass full orbital information using pair irreducible representations and many-body irreducible representation, a tensor expansion operation is employed alongside the filtering process. This expansion is defined by the following relation:
Thii also called "uuw" tensor product in e3nn implementation.
35

---

Published as a conference paper at ICLR 2025
(m) ) (Err))) me
cltsmmaa (e.mm)))))mms)))) mos
(16)
where C' denotes the Clebsch-Gordan coefficients, and @ symbolizes the tensor expansion which is converse operation of tensor product. u @ o can be expressed as sum over tensor expansions:
4 3 v Σ W 85
(17)
subject to the coupling constraints |e, 6 2
Loss Weighting. We determined the loss weights through principled ablation study on validation set. While keeping λ and X2 fixed at 1, we varied λ from 0.5 to 3 to understand the impact of WALoss weighting. We found the method to robust across a reasonable range of values, with A d2 I, and λ3 2.5 providing consistently strong performance. The table below summarizes the performance metrics for different values of λ1, 22 and Äg.
Table 11: Performance metrics for with different A1, X2, and A values.
A
As
Humiloomia MAE
Funnan MAE]
G10so MAE
Ca.
Excr MAE 1
fa1MAB...
Ci
System Energy MAE [
0.5258
0.653
19116
2.0077
21.92
8.18
48.00%
54.77
2.5
0.4744
0.7122
0.73
1.327
18.84
7.33
48.03%
47.193
0.4586
1,50225
3,29005
3.1123
23.57
8.56
47.29%
64.857
0.4807
L,159
3.45767
3.2068
23.16
9.2
47.15%
62.61719
0.4223
2.38836
6.83532
5.69888
28.96
11.78
45.30%
72.9
36

---

Published as a conference paper at ICLR 2025
J ADDITIONAL THEORY
J.1 PROOF OF THEOREM
Theorem 2. Let H,
H. E m be symmetric matrices, and let S Rnxx be symmetric positive definite matrix. Consider the generalized eigenvalue problems:
HC = SCe, =C=SCe,
where and are diagonal matrices of eigenvalues, and C and C are the corresponding eigenvector matrices. Define AH H H. and let be the minimum distance between the eigenvalue of interest and the rest of the spectrum of SHS-H. Then, the following bounds hold:
1. Eigenvalue Differences:
((()) S
K(S) SS|
(H. (H.
|AH
F S
|AH
1,1
where k(S) 2 ||2 is the condition number of S with respect the spectral norm, 11.1 denotes the element-wint l norm.
|1- || denotes the Frobenius norm, and
2. Eigenspace Angle:
<(S) S
AHH
((S) S
|AH
sin
5
2
where 0 is the angle between the eigenspaces corresponding to (H. (H, S) and (H,
Proof. Consider the generalized eigenvalue problems for H and H with respect S:
HC SCC, HČ SČ
Since is symmetric positive definite, it invertible. We can transform the generalized eigenvalue problems into standard eigenvalue problems by multiplying both sides by S-1.
A S S A+E,
where E S-IAH
Weyl's Perturbation Theorem
Theorem 3 (Weyl's Perturbation Theorem). Let A, A E Ruxa be symmetric matrices with eigenval- ues \ < A2 S ... S An and \₁ S Se S ... S Say respectively. Then, for all = 1,...,
R A < 1 All2-
Applying Weyl's Theorem A and Å, we have:
(A( x((A/ < |E||2 = SS-'AHH
Bounding the Eigenvalue Differences
Using the sub-multipicaative property of the spectral norm:
S AH < ||1
Since the spectral norm bounded by the Frobenius norm:
AH < || r-
Combining these inequalities:
S |S-AA | AH
37

---

Published as a conference paper at ICLR 2025
The condition number k(S) is defined as:
:((S)
<(S) | 1S--112 112 = ||-
|
Substituting back:
k(S)
|S |S-'AH|| -
AH F.
|||
Therefore:
x(S)
(((A) - MAA) S
||H|
|
Since the eigenvalues of A and Å correspond to the generalized eigenvalues of (H. S) and (H, respectively, we have:
((S)
(H, - X (H.))
|HH
S
Since |AH|| p < ||AH| 1,1, can further bound:
.(())
(H. - -H, -
|A 1,1-
|S|||
Davis-Kahan sin Theorem
Theorem 4 (Davis-Kahan sin A Theorem). Let A, Å RaxR be symmetric matrices, and let U and ú be the invariant subspaces of A and Å corresponding to eigenvalues in intervals I and I, respectively. f dist(I, 0, then
A A-Alp
sin 011 S
6
where O is the matrix of principal angles between U and ú.
Bounding the Eigenspace Angle
Applying the Davis-Kahan Kahan Theorem our case:
sin 0 E
|S |S-1AH|
Using the previously established bound:
k(S)
|
2 <
AA F.
| 2
obtain:
x(S)
AA
sin 0
||S|||
5
Similarly, since ||HH]| < | ||H]]
k(S)
AH
sin <
S 1
8
This completes the proof.
38

---

Published as a conference paper at ICLR 2025
J.2 PERTURBATION ANALYSIS AND GROWTH OF EIGENVALUES
Theorem 5 (Generalized Bai-Yin's law). Let A be an n x n random matrix with independent entries having mean It and variance d², Then, with high probability, the spectral norm of A is bounded by:
||l|2 | + 20 nn.
Proof. Consider the matrix A decomposed into its mean and fluctuation components:
= Hu + B,
where J is the matrix of all ones and B is random matrix with zero-mean entries and variance o². First, we bound the spectral norm of the mean component JJJ. The matrix J/ is rank-1 matrix with all entries equal to 1. Its largest singular value n, so:
|ل 1 |||
Next, we bound the spectral norm of the fluctuation component B. Since B has i.i.d. entries with zero mean and variance x², using Bai-Yin's 's law, we get:
BB| 20,
with high probability.
Combining these bounds using the triangle inequality, we have:
|| ||l ||2 |B S Iu|n 20 20
Thus, with high probability, the spectral norm of A is bounded by:
||lla S |u/n + 20 n
Theorem (Perturbation Sensitivity Scaling). Consider the generalized eigenvalue problem HC SCe, where H R BxB is symmetric Hamiltonian matrix, S E Bx is positive definite overlap matrix, C E Bx is the eigenvector matrix, and E E R is a diagonal matrix of eigenvalues S). Let H = H AH be the perturbed Hamiltonian, with AH E BxB a perturbation matrix whose entries have mean μ and variance o². Suppose the smallest eigenvalue of S satisfies:
Amin (S) =c c+
No
where 0, A>> 0, a > 0, and No > 0 are fixed constants. Then, for each eigenvalue, there exists a constant K > such that:
( < K (oBVB+Blul)
for all sufficiently large B.
Proof. Consider the unperturbed generalized eigenvalue problem Hc; = X,Sc,, where H and are symmetric, S is positive definite, and cf Sc, = 1. The perturbed system is Hè, = S,Se, with H = H AH, é; c; + óc, and Å A; + 8A, Substituting and expanding yields:
(H + (H) + oc;) (A; + 8X,) + &c,). Subtracting the unperturbed equation and projecting onto cf noting cf (H 2,S) = 0, we obtain:
of AHc; of AHHC; = 5\; + dx Söc,.
Solving for the perturbation:
c A'Hc + of AHóc:
8A
1 + c. Sdc,
39

---

Published as a conference paper at ICLR 2025
Reciprocal Minimum Eigenvalue vs. Number Atoms 29)
25000
Reg regseessoon
x
1/ M/mmmmmm Elgenvalue
11000
X
1000
Poealum: 4 4.50-220
2
RI
10
15
25
Number of Carbon Atoms
Figure 12: The scaling of the smallest eigenvalue of the overlap matrix in satured hydrocarbons.
For small AH, dc, is ((]HH])) so c AHóc, and c Soc, are higher-order terms. Here, we only consider the first-order term arises from evaluating the perturbation with the unperturbed eigenvector. Higher-order terms account for eigenvector shifts. We claim that analyzing the higher- order terms gives rise to unnecessary complication and does not alter the nature of catastrophic scaling. Specifically:
8b R cf AHC;
Bounding this term:
|oA: S lcf AHc, ||:
|HH
|
AH 2-
Since c/ Sc; = 1 and S - Amiu (S)I, we have:
+
B
No
1c Amiu
c+
+
B
+
1+
No
No
For AH as random matrix, apply a spectral norm bound (by Bai- Yin law):
|HH|| S Ci BB,),,
with constant C₁, holding with high probability. Combining:
|5A <
No
1/
c
1
A
No
Thus, the lcading-ordee order bound is:
(.H..))) (HH =0
+
No
(oB1/2 + (oB)/++lIV) +BI)++ (OBB12++B))
cel
c
No
L
Remark 1. If higher-order terms like c AHóc; and the denominator correction are included,
the perturbation includes an additional term of order O reflecting the quadratic contribution of the perturbation strength.
(
(aB1/2 + Blul)2
40

---

Published as a conference paper at ICLR 2025
J.3 SCALING LAW OF THE SMALLEST EIGENVALUE OF OVERLAP MATRIX
Figure shows the relationship between the number of carbon atoms in saturated hydrocarbons and the reciprocal of the smallest eigenvalue of the overlap matrix. The x-axis represents the number of carbon atoms, ranging from 0 to 29. The y-axis represents the reciprocal of the smallest eigenvalue, ranging up 25,000. Specifically, the regression analysis yields a p-value of 4.58 10 an R value of 0.73, and a Pearson coefficient of 0.85, suggesting strong linear relationship.
We also compare another scaling law suggested by Høyvik (2020). As shown in Figure investigated the relationship between 1 and the number of basis functions for various molecules.
Aui
Our results indicate a clear power-law scaling, which can be expressed as y parameters for each molecule are as follows: Eicosane (C20H42): y = 0.000107 2..43, Graphene (C24): y 0.00158 x2283 and Diamond (C10): y 0.0178 44116 This power-law behavior emphasizes the significant impact of the number of basis functions on Amn with different scaling exponents for each molecule.
zb The fitted
J.3.1
OVERLAPS
Here, we analyzed system where each basis function overlaps with its k-nearest neighbors. The overlap matrix S > 0 in this case is symmetric banded matrix with bandwidth 2k: + 1:
if = j, if
Sij
8K--1 0,
il k li -
To derive scaling equation for the smallest eigenvalue, consider the case where k 1, which corresponds to the nearest-neighhoo 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:
0
0
1
S =
0
8
0
1.
The eigenvalues An of this tridiagonal matrix can be derived using the known solutions for such matrices:
nI B B+1
Ay = 1 + 2s cos
n
1,1,, B.
The smallest eigenvalue corresponds B:
BT B ++
Amin = 1 + 2s cos
Using the trigonometric identity cos( - 8) cos(f) and approximating for large B:
BT
π B+1
π
cos
= cos
~
B+1
B B+1
Substituting back:
41

---

Published as a conference paper at ICLR 2025
Elcosane
4.
4.0
1.4
10
26
lnghturterr Function))
lighteererr
Figure 13: The scaling of the smallest eigenvalue of the overlap matrix in three systems. Data adapted from Høyvik 2020).
π
Amin 2 1 2s +
Since O, this yields 2s 0. Here, we investigated the smallest eigenvalue Amia as - oc, identifying two distinct cases:
1. Saturating Case: When < 0.5 the smallest eigenvalue saturates at a finite value - 2s as
- 00.
2. Non-Saturating Case: For 0.5, Amin decreases indefinitely with increasing B.
K TRAINING
Predicting the Gap from Initial Guess In contrast to prior research, our study tackles the significant challenges posed by the scale of the PubChemQC t-zvp dataset. It is very challenging for the model to develop a good numeric starting point. To address this, we have shifted our focus towards more tractable objective: making predictions based initial guess, which is casy to obtain. This adjustment is formalized in our objective function as follows:
= arg min
Σ
dist (Ho(M), H - H
(18)
D
(M.
where [D] denotes the cardinality datasct D, and dist(..)) is a predefined distance metric, H (o) is the initial guess of the Hamiltonian. We conducted ablation study, showing that although predictions based on the initial guess significantly improve performance in Hamiltonian prediction, they struggle with predicting physical properties.
Total Loss Function The total loss function used to train our model a combination the orbital alignment loss and the mean squared error (MSE) loss between the predicted and true Hamiltonian matrices:
Liotal A1
1
H.
H (0)
H
+ d2
H. H
Hill
|
+ AsCallgn (19)
where λ1, A2, Ag are hyperparameters that control the relative importance of each loss component.
Experimental Setting The experimental settings presented in Table 12 For the PubChemQH dataset, we used an 80/10/10 train/validationtees 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 patience 1,000 steps. WANet converges at 278,391 steps, QHNet 258,267 steps, and PhisNet 123, 170 steps. All models used the Adam
42

---

Published as a conference paper at ICLR 2025
optimizer with learning rate of 0.001 for PubChemQH, along with polynomial learning rate scheduler with 1,000 warmup steps. We used gradient clipping of 1, and radius cutoff of 5 A. For QHNet and PhisNet, we referred to the official implementation for these two models.
Table 12: Hyperparameter settings for the experimental study using the PubChemQH and QH9 datasets.
Host-pannccs
PiDCaemrod
(WAND WWANDY
QUe (WANe)
PLlCizarOH (QRNEO
QHA colNee
PalCcesroft (HeaNes
Doietitiin
S000
SC-4
Scheshikr
Polywamod
Polynomisl
Polunerral
Polenumeal
Pelymaii
1R Wamnep Suep.
1,000
1930
1,000
1030
1.001
Narther erpe liserly trmooe the louning ror
Mai Styn
300,000
300,000
300,000
300/000
300L000)
Masiama sumber unikiig stess
spackkal
Enredding Dinemsion
125
129
Disenioon
Bartle Hider Size
NIS
NIA
32
32
NIA
Size of thz lodles layes the botleeeck
Numberr of GNS Layere
5
Nerzher grsph roural networt hyyrs Matimars distance
nugghhoring aioera
128
Eder Chaznels
128
128
32
32
Number channels
fealures
Doyp Park Ryoe
8.1
6.1
NIA
NOA
NIA
Probbiiity ofdepping pan the aatanes Deport can projection laper
Projecton Deop
0.5
00
The training details for the HOMO, LUMO, and GAP predictions are presented in Table 13 For 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 models: Adam optimizer with learning rate 0.001 and polynomial learning rate scheduler with 1,000 warmup steps. WANer followed the same training setup described in Table 12,
The complete regression model hyperparameters are detailed Table [3] V2, Uni-Mol+, and Uni-Mol2 models, used their original implementations.
For Equiformer
Table 13: The training details for the HOMO, LUMO, and GAP predictions.
Hyperparameter
Uni-Mols In 2033
Uni-Matl (inohong at al 2014
Equifamna
Deseription
Lecamin Rate
0:01
Batch Size
32
32
32
Schednler
Polynomial
Polynomial
Pulynomial
LR Wammup
1,000
1,000
1,000
Number of lincudy
leuming
Max
300,000
300,000
3000000
Maximun number of training
Embldding Dimeusion
128
Dinession injut
Enbelding
3072
512
Enkeiding dimeusino for feed-Gowwad network
Number Fnooder Layars
Number layers
encader
Drop
Probubiliyy dropping
43

---

Published as a conference paper at ICLR 2025
L DISCUSSION
L.1 DISCUSSION ON SYSTEM ENERGY ERROR
The mean absolute error (MAE) of 47 kcal/mol for system energy prediction is notably above the threshold of chemical accuracy (typically kcal/mol). This highlights the inherent challenges of accurately predicting system energies for large molecular systems. To our knowledge, this work is the first to evaluate Hamiltonian predictions on system energy. which presents a non-trivial implementation challenge. Previous studies have typically focused on metries such as cosine similarity of cigenvectors or MAE of occupied energy levels. By directly assessing system energy, our approach provides a more comprchensive and practical evaluation of Hamiltonian accuracy, which is critical for large-scale molecular simulations.
Our model demonstrates significant improvement over baseline models in predicting system energies for these large systems, indicating enhanced prediction accuracy. The primary goal of this work to showcase the scalability and applicability of our approach to large molecular systems, where achieving absolute precision is inherently difficult due to their complexity. Despite this, our results represent meaningful step toward more accurate and efficient predictions in this challenging domain. The advances we have made underscore the potential of our approach to be further refined and applied across a range of complex molecular systems.
Moreover, other critical molecular properties such as the highest occupied molecular orbital (HOMO), the lowest unoccupied molecular orbital (LUMO), and dipole moments are predicted with reasonable accuracy. These properties, often vital in quantum chemistry workflows. demonstrate the practical utility of our model beyond system energy prediction. The ability to predict multiple important molecular properties with competitive error rates reinforces the broader applicability of our model to a variety of quantum chemistry and matcrials science applications
L.2 ADDITIONAL CONFORMATIONAL ENERGY ANALYSIS
In quantum chemistry applications, relative energy differences between conformational states are typically more relevant than absolute energies. 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 AE between each conformation and a reference structure using both DFT and WANet.
3. Calculated the MAE of these energy differences.
For sampling, we employed Gaussian perturbation 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
0. 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.8499800 (.9.844ccll/oo)) depending the molecule's 's flexibility. For baseline models (included the Table 14), the PhisNet model withWALos. achieved 9.90 kcal/mol MAF, the QHNet model with WALoss achieved 10.98 kcal/mol MAF, and the WANet model without WALoss achieved 48.92 kcal/mol MAE. QHNet without WALoss yielded MAEs of 50.56 kcal/mol. Interestingly, prediction, demonstrating the effectiveness of WALoss. It worth noting that while more complex sampling strategy could be used, those methods more computationally expensive they require DFT calculations.
observed that the model without WALoss produce meaningless
L.3 DISCUSSION SCF ACCELERATION RATIO
We report 18% reduction in SCF cycles, which is notable improvement, especially considering that our work targets significantly larger and more complex molecular systems compared to previous studies such QH9 and QHNet, where SCF reductions of 18-35% were achieved. Given the
44

---

Published as a conference paper at ICLR 2025
Table 14: Relative energy differences between conformational states.
Model
MAE of Energy Differences (kcal/mol)
WANct w/WALoss
1.12
PhisNet w/WALoss
9.90
QHNet w/WALoss
10.98
WANet
48.92
QHNet
50,56
increased complexity and size of our systems, this 18% reduction represents a considerable improve- ment, 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.
45


Reviews:
[{"review_id": "XwvAU6jsC4", "reviewer": "Reviewer_Rcb6", "summary": "This work presents a novel machine learning framework for predicting Kohn-Sham Hamiltonians on large-scale molecular systems. The proposed novel method includes a novel loss function WALoss based on basis change, and a novel SO(3)-equivariant model WANet adopting eSCN operation and MACE based many-body interaction in its architecture. A large-scale molecule dataset PubChemQH is also curated from PubChemQC benchmark for experiments. The proposed method is experimentally shown to achieve good performance on QH9 and PubChemQH benchmarks.", "questions": "See Weaknesses part.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Originality:  \nThe originality contribution of this work is excellent by proposing a physics-inspired loss function and model for Hamiltonians prediction problem. Also, the curated dataset is a significant originality contribution by providing a large-scale testbed in the field.\n\nQuality:  \nThe quality of this work is evidenced by solid theoretic analysis and strong performance on benchmark experiments.\n\nClarity:  \nThe writing of this work is good, clear and well-organized.\n\nSignificance:  \nThis work makes significant contribution in not only proposing an insightful machine learning model, but also having benefits in reducing computational cost for broad quantum physics and chemistry community.", "weaknesses": "(1) In Table 2, 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)? Also, authors are encouraged to report the results of QHNet with WALoss on QH9 datasets.\n\n(2) 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].\n\n[1] Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. 2023. Efficient and equivariant graph networks for predicting quantum hamiltonian. In Proceedings of the 40th International Conference on Machine Learning (ICML'23), Vol. 202. JMLR.org, Article 1692, 40412–40424.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HRZM1Jq56R", "reviewer": "Reviewer_YtXY", "summary": "The paper identifies a problem with training Hamiltonian models using only element-wise loss, which can lead to unphysically large errors in ground-state properties, especially for larger systems. To address this, the authors propose a WALoss function that aligns the eigenspaces of the predicted Hamiltonian with the ground truth, along with a new architecture, WANet. Additionally, they introduce PubChemQH, a new Hamiltonian dataset containing substantially larger molecules than previous datasets like QH9. Experiments demonstrate that the proposed WALoss significantly improves the accuracy of properties derived from the predicted Hamiltonian.", "questions": "1. The computational cost of WALoss could be substantial for large molecular systems, especially those containing heavy elements, as it involves eigendecomposition and matrix operations for each training step. 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?\n\n2. How were the loss term weights (λ1, λ2, λ3) in Equation 20 determined? 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.\n\n3. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality:\n- The observation and formalization of the Scaling-Induced MAE-Applicability Divergence (SAD) phenomenon is both insightful and important, representing a novel identification of a fundamental limitation in current methods.\n- The WALoss approach offers an original solution to maintaining physical accuracy in Hamiltonian predictions, particularly for larger molecular systems.\n\nQuality:\n- The theoretical analysis is rigorous, supported by eigenvalue perturbation theory and clear mathematical derivations.\n- The experimental validation is comprehensive, with thorough ablation studies and detailed comparisons with baseline methods.\n\nClarity:\n- The technical content is well-organized and clearly presented, with supporting diagrams and mathematical formulations.\n- The experimental results effectively demonstrate the practical benefits of the proposed methods.", "weaknesses": "The paper's main weakness lies in its insufficient differentiation from prior work, particularly regarding the use of eigenvalues in Hamiltonian prediction loss functions. The methodology section lacks critical details about computational costs and efficiency analysis of the proposed method, particularly regarding the training overhead of WALoss compared to traditional approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0yFf3nGU0E", "reviewer": "Reviewer_hKHK", "summary": "This paper tackles the problem of predicting the Kohn-Sham Hamiltonians for molecules. It first proposed a new Hamiltonian dataset for large molecules derived from the PubChemQC dataset. The paper then identifies the problem with the robustness of prior models when applied to molecules with large numbers of atoms. The cause of the problem is explained both empirically and theoretically. To improve the robustness of Hamiltonian prediction models to larger molecules, the authors introduce a new loss called WALoss that explicitly penalizes deviations from the expected eigenstructure of the ground-truth Hamiltonian. To further improve the quality of Hamiltonian prediction, the paper proposed a novel architecture dubbed WANet. The experiments demonstrate that WALoss significantly improves the quality of the baseline model (QHNet) and the proposed model (WANet) on downstream tasks such as System Energy,  $\\epsilon_\\text{LUMO}$ and $\\epsilon_\\text{HOMO}$ prediction. The proposed WANet model outperforms the baseline model.", "questions": "### **Questions and remarks**\n- Equation (3) might benefit from additional details. Energies associated with occupied molecular orbitals (and LUMO) make up the most energy of the molecule.\n- The $\\epsilon_\\text{LUMO}$ and $\\epsilon_\\text{HOMO}$ results in Table 1 seem strange. Why is the difference between QHNet and WANet so significant only for these targets? Moreover, the $\\epsilon_{\\Delta}$ MAE for WANet is larger than  $\\epsilon_\\text{LUMO}$ and $\\epsilon_\\text{HOMO}$ MAE. This implies that the model is consistently mistaken in opposite \"directions.\" Could you please double-check the results or explain why this is happening?\n> The theorem highlights that the difference between the predicted and actual Hamiltonian matrices when only considering the element-wise norm, can lead to unbounded differences in eigenvalues/eigenvectors due to a significant $\\frac{\\kappa(S)}{||S||_2}$ratio\n- It would be interesting to see the distribution of these values for the PubChemQH/QH9 datasets. If I understand correctly, $||\\Delta \\mathbf{H}||_{1, 1}$ is the $B^2 * \\operatorname{MAE}(\\mathbf{H}, \\hat{\\mathbf{H}})$ between real and predicted Hamiltonians, and its values can be relatively small.\n- line 219: mistake in \"spectral\"\n- Is the Algorithm 1 novel? If not, please provide clear citations. Also, it is better to include the final version of the algorithm, including the changes from the last paragraph on page 5.\n\n### **Closing remark**\nOverall, the paper presents a valuable technique that greatly improves the applicability of Hamiltonian prediction models for downstream tasks. However, the WANet architecture seems poorly motivated, has no ablations, and is not properly compared with baselines. I would consider raising my score if the concerns with the WANet model are resolved.\n\n[1] Khrabrov, K., Ber, A., Tsypin, A., Ushenin, K., Rumiantsev, E., Telepov, A., ... & Kadurin, A. (2024). ∇ DFT: A Universal Quantum Chemistry Dataset of Drug-Like Molecules and a Benchmark for Neural Network Potentials. CoRR.\n\n[2] Unke, O., Bogojeski, M., Gastegger, M., Geiger, M., Smidt, T., & Müller, K. R. (2021). SE (3)-equivariant prediction of molecular wavefunctions and electronic densities. Advances in Neural Information Processing Systems, 34, 14434-14447.\n\n[3] Zhong, Y., Yu, H., Su, M., Gong, X., & Xiang, H. (2023). Transferable equivariant graph neural networks for the Hamiltonians of molecules and solids. npj Computational Materials, 9(1), 182.\n\n[4] Lu, S., Gao, Z., He, D., Zhang, L., & Ke, G. (2023). Highly accurate quantum chemical property prediction with uni-mol+. arXiv preprint arXiv:2303.16982.\n\n[5] Hu, W., Fey, M., Ren, H., Nakata, M., Dong, Y., & Leskovec, J. OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).\n\n[6] Eastman, P., Pritchard, B. P., Chodera, J. D., & Markland, T. E. (2024). Nutmeg and SPICE: Models and data for biomolecular machine learning. Journal of Chemical Theory and Computation.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The problem with robustness is explained nicely. The authors demonstrate it empirically in Figure 1 and then provide two theoretical results that explain poor System Energy prediction quality even when relative MAE is low.\n- The effect of the proposed WALoss on the performance of models on downstream tasks such as System Energy,  $\\epsilon_\\text{LUMO}$ and $\\epsilon_\\text{HOMO}$ prediction is remarkable.", "weaknesses": "### **Major weaknesses**\n- 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.\n- The results in Table 1 do not provide enough evidence that WANet architecture is superior to baselines. A recent benchmark $\\nabla^2$DFT [1] compares Hamiltonian prediction models on a dataset of molecules with more atoms than QH9. The results in [1] show that PhisNet [2] performs significantly better than QH9. A comparison to other Hamiltonian predicting models, such as [2] and [3], is necessary to validate the claims about the proposed architecture.\n- 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]. The System Energy MAE is also ~50 kcal/mol, almost 50 times larger than metrics published in [1] and [6]. Moreover, the acceleration rate seems relatively small (18%), and most SCF iterations are still required when starting from the predicted Hamiltonian.\n\n### **Minor weaknesses**\n- Relatively large molecules (40 - 100 atoms) are used, but the functional (B3LYP-Def2TZV) does not contain the dispersion correction. The absence of the dispersion correction can potentially lead to inaccurate Hamiltonians.\n- The notation across the paper is unclear. A matrix with a subscript (i.e., $C^*_i$) usually defines the i-th row of the matrix, whereas in this paper it defines the $C^*$ matrix for the i-th molecule in the batch. Moreover, $C^*$  often denotes the conjugate transpose of matrix $C$. This notation complicates understanding Algorithm 1, as the Shur algorithm for complex-valued matrices operates with conjugate transposition. Undefined symbols in the equation on page 4. In equation (2), $n$ denotes the batch size but was never introduced.\n- Related work is in Appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UWaLdRcqsi", "reviewer": "Reviewer_LpfH", "summary": "This work encompasses 4 contributions. First, it introduces a new, larger dataset for predicting Kohn-Sham Density Functional Theory (KS-DFT) Hamiltonians, surpassing previous datasets in size. Second, it finds that training with Mean Absolute Error (MAE) on the Hamiltonian is inadequate for larger datasets, often producing results inferior to traditional minao guesses. Third, the authors propose a wave function alignment loss as an alternative to MAE, focusing on optimizing the eigenspace rather than the Hamiltonian itself. Finally, they present a new model that achieves state-of-the-art performance.", "questions": "* 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is very well written and easy to understand.\n* The empirical analysis of SAD is valuable and motivates this work well.\n* The WALoss appears like a natural choice and comes at essentially no additional cost.\n* Both the WALoss and WANet significantly improve upon previous works.", "weaknesses": "* The asymptotic scaling in Collary 1 for the lowest eigenvalue seems unfit to describe Gaussian-type orbitals. Due to their fast decay, the overlap is spatially limited. After a certain system size, more atoms are unlikely to affect the lowest eigenvalue.\n* l.62/63, \"complex transformation\" sounds overly complicated given the symmetric nature of the Hamiltonian.\n* Unfortunately, the real-world DFT speed-ups remain limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)