Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Authors: Erpai Luo
Xinran Wei
Lin Huang
Yunyang Li
Han Yang
Zaishuo Xia
Zun Wang
Chang Liu
Bin Shao
Jia Zhang
Venue: ICML 2025 poster (2025)
Keywords: Equivariant network
Hamiltonian Matrix
Computational Chemistry
Efficiency
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: SPHNet, an efficient and scalable equivariant network, which introduces adaptive SParsity into Hamiltonian prediction networks.

Abstract: Hamiltonian matrix prediction is pivotal in computational chemistry, serving as the foundation for determining a wide range of molecular properties. While SE(3) equivariant graph neural networks have achieved remarkable success in this domain, their substantial computational cost—driven by high-order tensor product (TP) operations—restricts their scalability to large molecular systems with extensive basis sets. To address this challenge, we introduce **SPH**Net, an efficient and scalable equivariant network, that incorporates adaptive **SP**arsity into **H**amiltonian prediction. SPHNet employs two innovative sparse gates to selectively constrain non-critical interaction combinations, significantly reducing tensor product computations while maintaining accuracy. To optimize the sparse representation, we develop a Three-phase Sparsity Scheduler, ensuring stable convergence and achieving high performance at sparsity rates of up to 70\%. Extensive evaluations on QH9 and PubchemQH datasets demonstrate that SPHNet achieves state-of-the-art accuracy while providing up to a 7x speedup over existing models. Beyond Hamiltonian prediction, the proposed sparsification techniques also hold significant potential for improving the efficiency and scalability of other SE(3) equivariant networks, further broadening their applicability and impact.

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Erpai Luo" ! Xinran Wei 22
Lin Huang? Yunyang Li Han Yang² Zaishuo Xia*
Zun Wang? Chang Liu"
Bin Shao? 2 Jia Zhang?
Abstract
1. Introduction
The Kohn-Sham Hamiltonian matrix is a fundamental quan- tity in computational chemistry, as it enables the prediction of all molecular properties obtainable through traditional Density Functional Theory (DFT) calculations (Hohenberg & Kohn, 1964; te Vrugt et al., 2020). Accurately predicting the Hamiltonian matrix allows for the rapid determination of system energy, HOMO-LUMO gaps, electron density, and other important properties (Fang et al., 2022; Zang et al., 2023; Batzner et al., 2022; Batatia et al., 2022; Wang et al., 2022; Li et al., 2024; Chen et al., 2023). Consequently, recent efforts have focused using machine learning tech- niques to predict the Hamiltonian matrix (Schütt et al., 2017; Unke et al., 2021: Yu et al., 2023; Gong et al., 2023; Zhou et al., 2022; Zhong et al., 2023). Among these, methods based on SE(3) equivariant graph neural networks (Unke et al., 2021; Yu et al., 2023; Gong et al., 2023; Yu al., 2023; Li al., 2025b) have achieved superior accuracy due to their consistency with the equivariance of the Hamilto- nian matrix. These methods typically project molecular features into spherical harmonics space and use tensor prod- uct operations to interact features of different orders, which are crucial for maintaining equivariance.
Hamiltonian matrix prediction is pivotal in com- putational chemistry, serving as the foundation for determining a wide range of molecular prop- erties. While SE(3) equivariant graph neu- ral networks have achieved remarkable success in this domain, their substantial computational cost driven by high-order tensor product (TP) operations-restricts their scalability to large molecular systems with extensive basis sets. To address this challenge, we introduce SPHNet, an efficient and scalable equivariant network, that in- corporates adaptive SParsity into Hamiltonian pre- diction. SPHNet employs two innovative sparse gates to selectively constrain non-critical inter- action combinations, significantly reducing ten- sor product computations while maintaining ac- curacy. To optimize the sparse representation, we develop Three-phase Sparsity Scheduler, ensur- ing stable convergence and achieving high perfor- mance sparsity rates of up to 70%. Extensive evaluations QH9 and PubchemQH datasets demonstrate that SPHNet achieves state-of-the. art accuracy while providing up 7x speedup over existing models. Beyond Hamiltonian pre- diction, the proposed sparsification techniques also hold significant potential for improving the efficiency and scalability of other SE(3) equiv- ariant networks, further broadening their appli- cability and impact. Our code can be found
Tensor product number
Time cost for single tensor penduct
led
Y N
0.5
y x L
0.3
0.1
50
90 Atoms
0
2
Order
Figure I. (A) The number of tensor products grows quadratically with the number of atoms N. as the Hamiltonian includes features for possible atomic pair combinations. tensor products grows with the sixth power of their order L,, where the increase in order corresponds to the expansion in the number of orbital types in the DFT basis For example, the def2-SV basis set requires a maximum order of 4, while def2-TZV demands an order of 6.
Equaal contribution 'Deparmeent of Automation, Tsinghua University, Beijing, China 'Microsoft Research Al for Science, Beijing, China Department of Computer Science, Yale Univer- sity. CT. USA 'Department of Computer Science, University of California, CA, USA. Correspondence ran @ wwwrosoctt.coommm com'>, Lin Huang <huang. lin @ wwwemmmmot.ccmmm. Jia Zhang <jia.zhang @microsoft.com.
(B) The time cost
Xinran Wei weixxn-
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
However, the efficiency issues introduced by tensor product operations have become a significant bottleneck. On the one hand, as the number of atoms in the system increases, the

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
number of tensor product operations grows quadratically, significantly increasing computational cost and memory re- quirements, as illustrated in Fig. (A). On the other hand, since each order in the output representation corresponds to different orbitals in the DFT basis set, the order of tensor product operations must increase as the DFT basis set ex- pands. As shown in Fig. I(B,, the computational complexity of tensor products scales with the sixth power of the order, causing the time cost of a single tensor product operation to grow rapidly with larger basis sets. Therefore, optimizing both the number and efficiency of tensor product operations is critical for extending Hamiltonian prediction to larger systems and larger basis sets.
due to its unique advantage in predicting quantum tensors (Fuchs al., 2020; Du et al., 2022; Musaelian et al., 2023; Liao & Smidt, 2022; Liao et al., 2023; Batzner et al., 2022; 2023). Key models include
(Fuchs et al., 2020), which introduced a robust self-attention mechanism for 3D point clouds and graphs, and Equiformer (Liao & Smidt, 2022), which predicted molecular properties using SE(3) Transformer architecture. EquiformerV2 (Liao et al., 2023) improved on this by employing efficient eSCN convo- lutions (Passaro & Zitnick, 2023), outperforming traditional networks like GemNet (Gasteiger et al., 2022) and Torchmd- Net (Thölke & De Fabritiis, 2022). Allegro (Musaclian et al., 2023) used a local, equivariant model without atom-centered message passing, showing excellent generalization.
The recent advancements in leveraging sparsity in numerical DFT solvers (Burow& & Sierka, 2011; Laqua et al., 2022) and message-passing networks (Liu al., 2023; Peng & Zhang, 2022; Passaro & Zitnick, 2023) inspire our exploration of incorporating adaptive sparsity into Hamiltonian prediction models. To address the aforementioned challenges, we developed SPHNet, and the main contribution of this work can be summarized as follows:
Hamiltonian Matrix Prediction. Hamiltonian predic- tion has advanced with neural networks in recent years. SchNOrb (Schütt et al., 2019) extended SchNet (Schürt et al., 2017) for high-accuracy molecular orbital predictions, while PhiSNet (Unke et al., 2021) successfully introduced SE(3) networks, significantly improving accuracy while fac- ing inefficiency due to tensor product operations. QHNet (Yu al., 2023) improved efficiency by reducing the num- ber of tensor products and (Li et al., 2025b) further extends the scalability and applicability of Hamiltonian matrix pre- diction on large molecular systems by introducing a novel loss function. At the same time, methods such as DeepH (Li et al., 2022) are focused on the Hamiltonian prediction of the periodic system.
We propose a scalable and efficient network for Hamil- tonian prediction, termed SPHNet, which introduce adaptive sparsity into equivariant networks by intro- ducing two sparse gates. Specifically, we employ the Sparse Pair Gate to filter out unimportant node pairs, reducing the number of tensor product computations, and the Sparse TP Gate to prune less significant in- teractions across different orders in tensor product, thereby improving the efficiency of tensor operations.
Network Sparsification. Network sparsification enhances computational efficiency by removing redundant neural net- work components. Early methods like Optimal Brain Dam- age (LeCun et al.. 1990) and Optimal Brain Surgeon (Has- sibi & Stork, 1993) pruned weights based on importance, followed by retraining to restore performance. Iterative pruning and retraining (Han et al., 2015) became widely used for reducing complexity while preserving accuracy. The Lottery Ticket Hypothesis (LTH) (Frankle & Carbin, 2019) later showed that sparse subnetworks in dense models could independently match full model performance. Spar- sification has since expanded to target weights, nodes, and edges, proving effective in molecular property prediction (Liu et al., 2023; Peng & Zhang, 2022). Recent work has also extended to pruning weights of the Clebsch-Gordan tensor product (Wang al., 2023), demonstrating potential despite unproven efficiency.
To optimize the sparse representation, we develop a Three-phase Sparsity Scheduler, which ensures effi- cient weight updates for all combinations through three stages: random, adaptive, and fixed. This approach fa- cilitates stable convergence to the optimal set while maintaining high accuracy at sparsity rates up to 70%.
We demonstrate the enhancements of SPHNet across multiple datasets, showing that it outperforms existing models in accuracy on QH9 and PubchemQH, while achieving up to 7x speedup and reducing memory usage by up to 75%. It also demonstrates compara- ble performance on small molecular trajectories on the MD17 dataset. Moreover, the proposed adaptive sparsification techniques exhibit strong promise in im- proving the computational efficiency and scalability of other SE(3) equivariant networks, paving the way for broader applications across diverse tasks.
3. Preliminary
DFT Hamiltonian. Density Functional Theory (DFT), as formulated by Hohenberg and Kohn (Hohenberg & Kohn, 1964) and further developed through the Kohn-Sham equa- tions by Kohn and Sham (Kohn & Sham, 1965), represents key computational quantum mechanical methodology for
2. Related Works
SE(3) Equivariant Neural Network. The SE(3) equiv- ariant neural network is widely used in AI for chemistry
2

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
studying the electronic structure of many-body systems, including atoms, molecules, and solids (Jain et al., 2016; te Vrugt et al., 2020). The core principle of DFT lies in solving the matrix C via the Kohn-Sham (KS) equations, summarized briefly as (Szabo & Ostlund, 2012):
Three-phase Sparsity Scheduler, three-stage selection func- tion that always selects retained elements from the set based learnable weight matrix W. Thus, for any unsparsified set U, the scheduler can be defined as:
RANDOM 1 11-k), k),,
if epoch < t, if epoch = t,
HC SCE,
TSS(W, =
TOP(W, 1 1-k), k), TOP(WPt', 11- k),,
where C representing the coefficients of molecular orbitals, H signifies the Hamiltonian matrix, S refers the over- lap matrix, and is the diagonal matrix encapsulating or- bital energies. Note that C E Raxxo, H,S E RTx=, and E R"*x*., where n denotes the number of basis func- tions and n. denotes the number of atomic orbitals used in DFT calculations. Primarily, an iterative self-consistent field (SCF) methodology (Payne et al., 1992: Cances & Le Bris, 2000; Wang & Baerends, 2022; Kudin et al., 2002) is uti- lized to solve the Kohn-Sham (KS) equations and to obtain the Hamiltonian matrix, requiring a temporal complexity of (nn)).
if epoch t,
where l represents the sparsity rate of the tensor product, t is the round of the training epoch, and (SS(.) always returns a subset USS of U, containing (1 - k)|U elements selected based on W.
Specifically, in the first phase, W is initialized as an all-ones vector, and RANDOM(.) denotes a random function that selects elements from the set independently of W with given probability 1 - k. This ensures that all parameters are unbiasedly selected and updated. In the second phase, (OP(-) is used to select the elements whose weights are within the top - A percent of all elements, ensuring that the optimal combinations are chosen globally. In the third phase, WPz' represents a fixed vector without gradients, which is frozen after the final backward update of the second phase. From this point onward, the selected combinations remain unchanged. This design prioritizes efficiency, as static connections typically result in faster computation.
SE(3) Equivariant and Tensor Product. The coordi- nate system in 3D Euclidean space is the most commonly used description for molecular systems. In this space, the group of 3D translations and rotations forms the SE(3) group. A function f(.) is SE(3) equivariant if it satisfies f(D(g)=)= = (xr), where D(g) is representation of SB(3) acting on x. Therefore, a neural network in which every operation satisfies the definition of (()) is referred to as an SE(3) equivariant neural network. Among these operations, the tensor product operation is the most crucial. Typically, higher-order irreducible representations (irreps) are constructed using spherical harmonics. The tensor prod- uct operation combines two irreps, and y. with rotational orders& & and lg respectively, using the Gorsann (CG) coefficients (Griffiths& & Schroeter, 2018), resulting in a new irreducible representation of order lg as:
4.2. Sparse Tenor Product Gate
The Sparse TP Gate is designed to accelerate individual tensor product computations by filtering out unnecessary cross-order combinations in traditional tensor product opera- tions, as illustrated in FBB. By doing so, tensor product operations implemented with e3nn (Geiger et al., 2022) can achieve near-linear acceleration, as demonstrated in the Ap- pendix B.3. Moreover, this approach can be easily adapted to various tensor product implementations by simply modi- fying the instructions that declare cross-order combinations in the tensor product operation.
( mm
mls d (a) ma Mmaa
E2
(2) where m denotes the m-th element in the irreducible repre- sentation and satisfies -f < m < e. The output order & is restricted by 141 - &21 S &y < ( + 12). Since each order has to interact with every other order, the computational complexity of the full tensor product using irreps up to or- der L have computational complexity of O(L%), which constrains its applicability for systems with higher degrees.
Generally, the classical tensor product operation includes the interaction of every element m in every irrep order f. The aim of Sparse Tensor Product Gate is learn a set of the most valuable combinations U:TSS from the complete set U. = {(f1, f2, fa) la E [[6] - f2). & + f2]} in the tensor product and remove the redundant combinations ac- cordingly. UTSS is defined by a one-dimensional learnable value score c
.2f of length (Ucl as:
4. Methodology
UTSS = {(11, (2, Wai E TSS( k)), (4)
4.1. Three-phase Sparsity Scheduler
where TSS(-) is the selection scheduler defined in Equation 3. Accordingly. we can obtain the updated combination weight from the original weights e defined in classical tensor
To ensure stable selection of combinations in the Sparse Tensor Product Gate and Sparse Pair Gate, we propose the
3

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
A
Z
B
D
Three-phase
&
Sparsity Scheduler
Vectorial Node
Sparse Pair
Random Phase
Interaction
Inner product
Gate
8
Lo
2,000
L2
Sparse Pair Gate
Concat
Lo
L1
L
Spherical Node Interaction
Linear
Sigmoid
12
Lp
Pair weight
All pair
Adaptive Phase
Pair Construction
Select
Lo
Lo
&
Sparse Pair
L2
Gate
0
02 03
10
12 13
La
Diag
Non-Diag
0P 13
Optimal pair
L ...
l2
Sparse
Sparse
TP Gate
TP Gate
C
Fixed Phase
Tensor product
Tensor product
Original weight
Sparse TP Gate
Lo
La
12
fj
Select
1o
[2 ...
Expansion Block
Learnable value
Random selected
Optimal comb
Adaptive selected Redundant
Predicted Hamiltonian
Figure 2. (A) The overall architecture of SPHNet. Atomic numbers and atomic coordinates first passed through the Vectorial Node Interaction Blocks to obtain atomic features x Subsequently, the Sparse Pair Gate selects the key pair set (i.j) for the Spherical Node Interaction Blocks, where the irreps x elevated from the é to Lmux during the interaction process. Next, the Sparse Tensor Product Gate in the construction block identifies the key cross-order combinations (E, , l2, {a) for the diagonal blocks, yielding diagonal pair features f,, For non-diagonal blocks, both the Pair Gate and Tensor Product Gate applied select the critical pairs and tensor product combinations, producing non-diagonal pair features f,. Finally, these features are fed into the expansion block to construct the predicted Hamiltonian matrix. (B) Sparse Pair Gate: It takes pairwise features as input, computes weights for each pair (i, j). and selects a optimal subset using the sparsity scheduler. (C) Sparse Tensor Product Gate: Similarly, it utilizes the sparsity scheduler to identify an optimal subset of cross-order combinations (Ea, 82, Ea) based on learnable weights, (D) Three-Phase Sparsity Scheduler: Designed the sparse gates, it operates in three phases; random, adaptive, and fixed.
product by:
selects subset from all the node pair U, ((,
d= tetiiats Whi, (4,, lx) E T...).
(5)
and uses linear layer Fş(-) to learn the weight of each pair W in this fully connected graph. This can be formulated as below:
In practical experiments, the sparsity rate is set based on the details in Section 5.4.
I (x?1) (x,,
(6)
W = Sigmoid (F,))))))
4.3. Sparse Pair Gate
Since the output of the Hamiltonian matrix encompasses interactions between all pairs of atoms, previous models typically define inter-atomic relationships as fully con- nected. This results in tensor product computations that scale quadratically with the number of nodes. However, observed that not every block requires such dense connec- tivity. To reduce the number of tensor product operations in the network, we developed the Sparse Pair Gate, which adaptively selects the most important edges.
where x and x are the zero-order features of the irreps, is concatenation operation in the last feature dimension and (.--) stands for the inner product operation. Thus, the sparse pair gate can formulated as:
UTSS {(i,)) W ETSS(W, p. k)).
(8)
where TSS(:) is the same Equation Accordingly, we can get the pair weight Way of x;, Xy used in tensor produet:
Similarly with the tensor product gate, the Sparse Pair Gate
Wij = F, (RBF( F. (WÚ) I,),
(9)

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
where (.)) and (.) are linear layers, (BF(.) is the radial basis function. The extra computational cost introduced by above sparse gates and the scheduler is minimal and has negligible impact on the overall time, which is discussed in detailed in Appendix D.7.
where Lmax represents the allowed max order of irreps in the network and Way is the weight defined in Equation 9.
Pair Construction Block. Consists of the Non-Diagons block and the Diagonal block. The Non-Diagons block uses these irreducible representations to compute the interaction from node j to node i, generating the pair-wise features f.y for the -on-Diagonal blocks of the Hamiltonian matrix, The Diagonal block calculates the self-interaction of node i to produce the node-wise feature f. for the diagonal blocks of the Hamiltonian matrix.
4.4. SPHNet
With the above two sparse gates, SPHNet can be formu- lated as four modules. The model takes atomic numbers Z and 3D coordinates of molecular systems as input features, initializes the node representation zi with a linear func- tion as 2; = (Z,). processes them through the Vectorial Node Interaction Block, and sequentially passes through the Spherical Node Interaction Block, Pair Construction Module, and Expansion Module to output the predicted Hamiltonian matrix H. as illustrated in Fig.2.
As shown in F(g. use Sparse Tensor Product Gate before both the Diagonal and Non- -on-Doagonal block, and use Sparse Pair Gate before Non- -inaooa block, therefore the fai fiy can be reformulated as:
Σ
W W(( x (X)) 8 0 x/2), (13)
Node Interaction Block. The Node Interaction Block ag- gregates information from neighboring nodes through a pessagg-passing mechanism, thereby extracting irreducible representations of nodes. Specifically, the irreducible rep- resentation I; can be updated by aggregating the message mij as,
Σ
(2. x (2.) of wi (allall))))
(12 (2.
S.L. (i.1) E U,SS t, S Lmaxs
(14)
where Wa is learnable parameters for each combination in UTSS c' is an updated combination weight defined in Equation and Wiy is the weight defined in Equation 9, It's noteworthy that the Sparse Pair Gate is only employed from the second Spherical Node Interaction Block and the Non-Diagonal Pair Construction block to ensure complete information flow between nodes, and more details could be found in Appendix D.5.
Fm(x + Σ m
(10)
where (-) is a linear layer.
To achieve sufficiently interactive high-order information while minimizing the associated computational cost, we propose a mechanism that gradually increases the maximum order of node representations. Initially, four Vectorial Node Interaction Blocks are applied to obtain vectorial represen- tations m where& & < 1, as follows:
Expansion Block. Finally, the expansion block utilizes the tensor expansion operation to obtain the non-diagonal Hamiltonian block h:j and diagonal Hamiltonian block has The detailed formulation can be found in Appendix D.6. Considering the symmetry of the Hamiltonian matrix, where for any sub-block h:j, there exists corresponding sub- block hys such that hy; = h2y. So we can further enhance the model's efficiency by constructing the node pair features f,y from the upper triangular part of the Hamiltonian ma- trix, i.e., {fay i < j}. Leveraging this symmetry allows to halve the number of tensor products required when predicting the node pair irreps f.y. significantly reducing the computational burden associated with constructing h.j-
m c o rig). s.t. 151.
(11)
where Wij is the weight defined in Equation 9, and riy the distance vector between atoms. This representation does not involve interactions between high-order tensors, thereby eliminating tensor product operations, as detailed in Appendix D.3.
Next, we apply two Spherical Node Interaction Blocks to capture interactions for high-order irreducible representa- tions, increasing the highest order of the irreducible repre- sentation to match the maximum required order determined by the basis set. Specifically, these blocks project distance information Fīg into the high-order spherical space using spherical harmonics function y=((). (-). Subsequently, mi, in this block is computed under the constraints imposed by Sparse Pair Gate as:
5. Result
We conducted a series of experiments to compare the overall performance of SPHNet with the previous ofte-offthhean model, including SchNOrb (Schütt et al., 2019), PhiSNet (Unke et al., 2021), QHNet (Yu et al., 2023), and WANet (Li et al., 2025b), and the results listed in the table are sourced from their papers. Note that the results for PhiSNet and
m
Σ f
x Biie whi (r7g);
(12) SchNOrb can only be conducted on the MD17 dataset, as these models are specifically designed for trajectory datasets
s.t. (i,j) E UTS,
ei, i1l2, ls s Lmax.
5

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Table Experimental results
QH9 dataset.
V
Memory 4 [GB/Sample]
Speed f [Sample/Sec]
Speedup 1
Dataset
Model
[10- [10~
[10-3]
Ratio
QHNet
76.31
798.51
95.85
0.70
19.2
1.0x
QH9-stable iid
WANet SPHNet
80.00
833.62
96.86
N/A
N/A
N/A
45.48
334.28
97.75
0.23
76.80
4.0x
QHNet SPHNet
72.11
644.17
93.68
0.70
21.12
1.0x
QH9-stable ood
43.33
186.40
98.16
0.23
78.72
3.7x
QHNet
70.03
408.31
97.13
0.70
24.68
1.0x
QH9-dynamin geo
WANet
74.74
416.57
99.68
N/A
N/A
N/A
SPHNet
52.18
100,88
99.12
0.23
82.56
3.3x
QHNet
120.10
2182.06
89.63
0.70
23.04
1.0x
QH9 diy-ddyamicc mol
SPHNet
108.19
1724.10
91.49
0.23
80.64
3.5x
in conformational space. Additionally, since WANet has not been open-sourced, we only report results on a subset of datasets and metrics based on the information provided in their paper. Detailed experimental settings are described in Appendix A.
four different split sets and compared the results with the baseline models QHNet and WANet. As shown in Table 1, for two stable split sets, SPHNet attained a speedup of 3.7x to 4x over QHNet while improving the accuracy of the Hamiltonian Mean Absolute Error (MAE) and the occupied energy MAE. Furthermore, SPHNet significantly reduced GPU memory usage, requiring only 30% memory usage of the baseline model. For two dynamic split sets, SPHNet maintained a speedup of 3.3x to 3.5x, simultaneously en- hancing prediction accuracy and decreasing GPU memory usage by over 70%. These results underscore SPHNet's effi- ciency, achieving substantial speedups and resource savings without sacrificing precision, indicating that SPHNet can effectively learn molecular features and the Hamiltonian matrix within the small-scale dataset, establishing solid foundation for calculations in larger-scale systems.
Datasets. Three datasets were used in the experiments: The MD17 dataset includes Hamiltonian matrices for the trajectory data of four molecules: water, ethanol, malon- dialdehyde, and uracil, containing 4,900, 30,000, 26,978, and 30,000 structures, respectively. The QH9 dataset (Yu et al., 2024), consisting of Hamiltonian matrices for 134k small molecules with no more than 20 atoms, and the Pub- ChemQH dataset (Li et al., 2025b), containing Hamiltonian matrices for 50k molecules with atom counts ranging from 40 to 100. The QH9 and PubChemQH use the B3LYP exchange-correlation function, and the MD17 uses the PBE exchange-correlation function. The orbital basis DeP2-SVP used for MD17 and QH9, while the D is used for PubChemQH. Compared the MD17 and QH9 datasets, the maximum number of orbitals in the Hamiltonian of the PubChemQH dataset increases by about 15 times, result- ing in a 200x difference in the number of elements in the Hamiltonian matrix.
5.2. Performance on PubChemQH Dataset
We further validated our approach on the larger-scale molec- ular system and trained SPHNet on the PubChemQH dataset. It was observed that SPHNet trains over 7.1x times faster than the baseline QHNet while achieving a better Hamilto- nian prediction accuracy, as shown in Table 2, which has a higher speedup ratio compared that in the smaller-scal datasets. Moreover, GPU memory consumption of SPH- Net was significantly lower, at only 25% and 37% of the baselines.
Evaluation metrics. The evaluation metrics include Hamil- tonian MAE (mean absolute error between predicted and Hamiltonian matrices calculated by DFT), occupied energy MAE (mean absolute error of energies of occupied molecu- lar orbitals calculated from predicted Hamiltonian matrices), C manilarityccosinn similarity of coefficients for occupied molecular orbitals), training GPU memory (GB per sample), training speed (training samples per second) and speedup ratio compared to QHNet.
Table 2. Experimental results on PubChemQH dataset.
Model
QHNet
WANet
SPHNet
H [10 1100 [EA]] 4
123.74
99.98
97.31
3.33
1.17
2.16
+ [10 -21
2.32
3.13
2.97
5.1. Performance on QH9 Dataset
Memory [GB/Sample]] 4
22.5
15.0
5.62
We first evaluated the performance SPHNet on the QH9 dataset. Note that the QH9 dataset has two subsets and four different splits, including the stable-iid, stable-ood, dynamic-geo, and dynamic-mol. We trained SPHNet on
Speed [Sample/Sec]]
0.44
1.09
3.12
Speedup Ratio
1.0x
2.4x
7.1x
We observed that SPHNet achieves a higher speedup on
6

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
the PubChemQH dataset. As described in the dataset sec- tion, the number of orbitals in the Hamiltonians of Pub- ChemQH is 15 times greater than in QH9. Additionally, since PubChemQH was computed using the Def2-TZVP or- bital basis, the network requires a maximum output angular momentum order of Lma.. = 6, whereas datasets computed with the Def2-SVP basis only require Liuxx = 4. These make the number and time consumption of tensor products much higher than that in small-scale datasets. Thus, we can reasonably infer that the higher speedup achieved on the PubChemQH dataset is due to the markedly greater com- putational complexity of its molecules compared to those in the MD17 and QH9 datasets. This increased complexity introduces higher information redundancy, making it more amenable to adaptive sparsification. Our sparsity ablation experiments across different datasets support this inference. As molecular size and maximum angular momentum or- der increase, the rising sparsity rates of sparse pair gates and sparse tensor product gates generally result in smaller accuracy losses. Therefore, higher sparsity rates can be applicd to SPHNet, further accelerating its performance on large-scale datasets. For more details on this ablation study, please refer Section 5.4.
spaces. On the one hand, baseline models generally already perform well across all these datasets, leaving limited room for SPHNet to achieve significant improvements in either accuracy or speed. On the other hand, taking the water molecule with only three atoms as an example, the num- ber of possible interaction combinations within the system is inherently small, which limits the potential benefits of adaptive sparsification.
5.4. Ablation Study on Sparsity Rate
The key idea of SPHNet is the implementation of adaptive sparsity to reduce computational demands in SE(3) equiv- ariant networks. The most critical hyperparameter in this adaptive sparse system is the sparsity rate. To determine the optimal sparsity rate and assess its impact across various molecular systems, we conducted a series of experiments. We varied the sparsity rate from 0% to 90% in 10% intervals and trained SPHNet on three datasets representing different molecular scales (Here we chose Ethanol to represent MD17 dataset.). As illustrated in Fig.3, the predicted Hamiltonian MAE across all three datasets remained relatively stable at lower sparsity rates but increased significantly when the rate reached a specific threshold. This finding suggests that an appropriate range of sparsity has minimal impact on model accuracy while simultaneously enhancing computational speed.
Table 3. Experimental results on MD17 dataset.
Dataset
Model
H [10- [
[10-9Ex]
& 1
SchNOrb
165.4
279.3
100.00
PhiSNet
15.67
85.53
100,00
Water
QHNet
10.79
33.76
99.99
200
PubchemOH
QH9
MD17
X 1336.00
SPHNet
23.18
182.29
100,00
Hamitanion MAE/ 10e En
X Speed: Samples / Sec
5.13
X 4887.00
SchNOrb
187.4
334.4
100.00
150
PhiSNet
20.09
102.04
99.81
Ethanol
QHNer
20.91
81.03
99.99
100
3.12
SPHNet
21.02
82.30
100.00
0.78
SchNOrb
191.1
400.6
99.00
60
36:00
76,80
PhiSNet
21.31
100.6
99.89
Malon- aldchyde
132.00
204.00
QHNet
21.52
82.12
99.92
SPHNet
20.67
95.77
99,99
0%
10%
20%
30%
60% Sparsity Rata
40%
30%
70%
80%
80%
SchNOrb
227.8
1760
90.00
PhiSNet
18.65
143.36
99.86
Uracil
QHNet
20.12
113.44
99.89
Figure 3. The effect of different sparsity rates on the model per- formance on the datasets with different scale of molecules. See the detailed computational cost scaling in Appendix B.2 of Puh- ChemQH Dataset.
SPHNet
19.36
118.21
99.99
5.3. Performance on MD17 Dataset
Additionally, we observed that the threshold at which the MAE begins to rise significantly becomes larger as the com- plexity of the molecular system increases. For the smallest MD17 system, the critical turning point is at 30%, while for the QH9 system, it 40%, and for the largest PubChemMH system, it rises 70%. In the experiments conducted in this study, we adopted the above critical turning point as the spar- sity rate for training and testing for different datasets. This observation aligns with our earlier conclusion that the adap- tive sparsity strategy offers greater potential in large-scale systems, as there are more calculations of lower importance that can be optimized. Consequently, the system size can
We also evaluated the performance of SPHNet on the smaller-scale MD17 dataset. As presented in Table 3, we compared SPHNet's performance with four baseline models across four MD17 water, ethanol, malondialde- hyde, and uracil coneillcomiiinnggg 3 to 12 atoms. The results demonstrate that SPHNet achieves accuracy comparable to other models, indicating its suitability for small molecular trajectory datasets.
is worth noting that predictions on the MD17 dataset rep- resent relatively simple task compared to other datasets, as focuses solely on small systems and their conformational
7

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
serve as a rough estimate for determining the applicable sparsity rate, thereby reducing the need for extensive param- eter searches.
Then, for the selected optimal combination set, we ana- lyze the order of output irreps of each selected cross-order combination in tensor product operations. For classical tensor product operation, each irreps of the output zaa is con- tributed by (( cross-channel combinations. As shown in Fig.4(B) and Fig.7, in the QH9 dataset, the proportion of each order that selected decreases monotonically as the order increases in both Sparse Combination Gates, In the PubChemQH, the trend is the same except the order zero has a relatively lower selected ratio. These results sug- gest that although higher orders' output irreps contain more combinations, the importance of a single combination is weakened as the number of combinations increases, making them easier to filter out.
5.5. Analysis of Optimal Selected Set
Experiments demonstrated that the two adaptive sparse gates effectively select the most important tensor products and their optimal combinations. Here, we examined the selected pairs and combinations to understand the learning outcomes of the sparse gates.
For the optimal pair set, we analyzed the length(atomic dis- tance) of the selected pairs, illustrated in Fig.4(A). The results indicate that sparse gating differs from the commonly used RBF- cutoff strategy, as it evaluates the impor- tance of a pair based solely on its contribution to the output irreps, rather than its distance. Interestingly, as the pair length increases, the probability a pair being selected also rises. This likely because the features of long-range pairs are more challenging to learn, requiring as many samples as possible achieve accurate representations. Notably, the selection probability increases nearly linearly in the distance range of 16A to 25A. This range corresponds to the influ- ence of electrostatic interactions and may also include weak van der Waals forces. These findings highlight the impor- tance of long-range interactions in Hamiltonian prediction and demonstrate the effectiveness of sparse pair gating in dynamically integrating both long-range and short-range in- teractions, ultimately enabling the model to achieve superior performance.
5.6. Scaling with Increasing Size of Hamiltonian
To further validate the limits of Hamiltonian prediction scale achievable by different models under restricted GPU mem- ory, as well as the scaling of speed and memory consump- tion with increasing system size, we randomly selected 6 molecules of varying sizes from the PubChemQC PM6 dataset. We computed the labels for these molecules with the setting of PubChemQH, and conducted tests with sin- gle A6000 card based on the outcomes. The results are illustrated in Fig.5 and Fig.8.
Our method demonstrates better scaling than the baseline model. As the Hamiltonian size (number of atomic orbitals) increases, the speedup ratio improves while memory con- sumption decreases. Thanks to significant savings in both memory and time, our model can train on larger systems (around 3000 atomic orbitals) within the same memory con- sumption, whereas the baseline model is limited to approxi- mately 1800 atomic orbitals.
A 1000
Numbe gi
800
Selocted pairs
600
400
200
100
Saldted Peggted Propotton
0.5
Speed Sample/Sec (logr)
50
SPHNcl
CHNNat
0.3
20
0.2
10:
0.1
5
5
10
15
20
25
2
Distanne AA
w6.69
QHH
PusCHemmH
19.00
0.84
0.8
Prganigrs
0.0-
0.7.
Selected Proportoons
0.5g
166
1182
1796
2104
2781
2895
Number of orbitais Hamiltonian
0.5-
0.3
Figure 5. Comparison of the training speed with the increasing size of Hamiltonian.
0.3
0.2
0.1
0.1
5.7. Additional Ablation study on SPHNet
0.0
0
2
3
5
Combination Output Order
Combiniion on Outpul Ordwr
To evaluate the effectiveness of the Three-phase Sparsity scheduler and the Sparse Gate on the model performance, we carried out a series of ablation studies. The results showed that these strategies have significant improvements in model efficiency and accuracy. The details are as follows.
Figure 4. (A) The distribution of distance between selected node pairs and the selected proportions of different pair distances in the pair gate. (B) The proportions of each output order being selected the first Non-Diagonal block's 's tensor product gate, See the full cases of selected pairs in Appendix B.4.
8

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Three-phase Sparsity Scheduler. We evaluated the effect of the three-phase sparsity scheduler on model performance by training SPHNet under various ablation settings: ex- cluding the random stage, the adaptive stage, or the fixed stage. Fach configuration was trained five times to assess performance stability and accuracy, As shown in Table 4, models trained without the random stage often converged to local optima, resulting in suboptimal performance. Models trained without the adaptive stage exhibited greater variance in outcomes and significantly lower accuracy compared to SPHNet with the full scheduler. Although omitting the fixed stage did not substantially affect accuracy, it reduced the speedup ratio to 5.45 due to the additional computational overhead from the sparse gate and the non-static graph im- plementation of tensor products following the Sparse Pair Gate, in contrast the optimized implementations in the e3nn library. These findings highlight the importance of the three-phase sparsity scheduler in enhancing model perfor- mance and efficiency.
Applying Sparse Gate to QHNet model. To further as- sess the effectiveness of the Sparse Gates, we conducted additional experiments by integrating them into the QH- Net model. Specifically, the Sparse Pair Gate was applied to the second non-diagonal pair block, while the Sparse Tensor Product Gate was applied to the node-wise interac- tion blocks as well as both diagonal and non-diagonal pair blocks. As shown in the table below, the sparse gates sig- nificantly improved the training speed of the QHNet model and reduced computational resource usage. These results demonstrate the generalizability of sparse gating mecha- nisms across different SB(3)- networks.
Table 6. The effect of sparse gates on QHNet model.
Sparse Pair Gate Sparse TP Gate
/
/
II [10-*E,] 4
123.74
128.16
126.27
128.89
Memory [GB/Sample] 4
22.50
12.68
10.07
8.46
Speed [Sample/Sec]1
0.44
0.90
0,73
1,45
Table 4. Ablation Study on Three--hase- Sparsity Scheduler.
Speedup Ratio f
1.00x
2.04x
1.66x
3.30x
Random Stage
x
Adaptive Stage
/
x
Fixed Stage
/
/
H 11100 F
97.31
112.68
122.79
97.11
6. Conclusion and Future Work
+0.52
=10.75
19.02
1.33
Memory [GB/Samplel L
5.62
5.62
5.62
5.62
In this work, we introduced SPHNet, efficient and scal- able SE(3) equivariant neural network for Hamiltonian pre- diction. By incorporating two adaptive sparse gates and a corresponding three-phase learning scheduler, SPHNet optimized both the number of tensor product operations and the efficiency of individual tensor computations. As result, SPHNet achieved exceptional efficiency and per- formance in predicting Hamiltonians for larger molecular systems. Moreover, the proposed sparsification techniques demonstrated significant potential for extension to other SE(3) equivariant networks and broader prediction tasks.
Speed [Sample/Sec]
3.12
3.12
3.12
3.12
Speedup Ratio 1
7.09x
7.09x
7.09x
5.45x
Sparse Gates. To further evaluate the impact of the Sparse Pair Gate and the Sparse Tensor Product Gate on overall performance, we trained SPHNet under three configurations: without the Sparse Pair Gate, without the Sparse Tensor Product Gate, and without both gates, As shown in Table 5, the Sparse Pair Gate achieved 78% speedup and a 30% reduction in GPU memory usage, with minimal impact model accuracy. The Sparse Tensor Product Gate, which removes 70% of the combinations in the tensor product, yielded a 160% speedup with an acceptable loss in accuracy. These results underscore the contribution of the Sparse Gate to improving model efficiency. Users may refer to Section 5.4 for guidance on choosing a sparsity ratio that balances speed and accuracy based on system size.
This study opens several promising avenues for future work. First, the adaptive sparsity technique shows considerable potential for generalization to wider range of tasks. Specif- ically, the sparse tensor product gates can be readily ex- tended to any SE(3)-equivarian network architecture based on Clebsch-Gordan tensor products, while the sparse pair gates could be adapted to more types of pairwise interac- tions, such as those found in MACE (Batatia et al., 2022) for multi-body interactions. Second, integrating advanced loss functions, such as those introduced in (Li al., 2025b), could further improve the accuracy of downstream proper- ties of the Hamiltonian matrix, particularly in enhancing the applicability for large molecular systems. These directions will guide future efforts with further experimental valida- tion, expanding the impact of adaptive sparsification and optimizing the performance of SE(3) (equivrriv networks.
Table Ablation Study on Sparse Gates.
Sparse Pair Gate
/
X
/
X
Sparse TP Gate
/
/
X
X
H [10-6 Ex]
97.31
94.31
87.70
86.35
Memory [GB/Sample] +
5.62
8.04
6.98
10.91
Speed [Sample/Sec] 1
3.12
1.75
1.20
0.76
Speedup Ratio 1
7.09x
3.98x
2.73x
1.73x
9

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Impact Statement
Frankle, J. and Carbin, M. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In Interna- tional Conference on Learning Representations (ICLR), 2019.
This paper presents work whose goal is to speed up the Hamiltonian matrix prediction process and advance the SE(3) equivariant network. There are many potential soci- etal consequences of our work, none of which feel must be specifically highlighted here.
Fuchs, F., Worrall, D., Fischer, V., and Welling, M. Se (3)-transformers: 3d roto- cnto-translation equivariant attention networks. Advances in neural information processing systems, 33:1970-1981, 2020.
References
Atz, K., Grisoni, F., and Schneider, G. Geometric deep learning on molecular representations. Nature Machine Intelligence, 1022)1102--103,, 2021.
Gasteiger, J., Shuaibi, M., Sriram, A., Günnemann, S., Ulissi, Z., Zitnick, C. L., and Das, A. Gemnet-oc: devel- oping graph neural networks for large and diverse molecu- lar simulation datasets. arXiv preprint arXiv:220.002882. 2022.
Batatia, I., Kovaes, D. P., Simm, G., Ortner, C., and Csányi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Ad- vances in Neural Information Processing Systems, 35: 11423-11436, 2022.
Geiger, M., Smidt, T., M., A., Miller, B. K., Boomsma, W., Dice, B., Lapchevskyi, K., Weiler, M., Tyszkiewicz, M., Batzner, S., Madisetti, D., Uhrin, M., Frellsen,
J., Jung, N., Sanborn, S., Wen, M., Rackers, J., Rød, M., and Bailey, M. Euclidean neural networks: e3nn, April 2022. URL https: /tt.//ooiiiooooooo http//i..ogg/////// org/10..g...oo org/10.5211/ zenodo. 6459381.
Batzner, S., Musaclian, A., Sun, L., Geiger, M., Mailoa, P., Kornbluth, M., Molinari, N., Smidt, E., and Kozinsky, B. E (3)-equívariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Gong, X., Li, H., Zou, N., Xu, R., Duan, W., and Xu, Y. General framework for e (3)-cquivariant neural network representation of density functional theory hamiltonian. Nature Communicutions, 14(1):2848, 2023.
Batzner, S., Musaclian, A., and Kozinsky, B. Advancing molecular simulation with equivariant interatomic poten- tials. Nature Reviews Physics, 5(8):437 438, 2023.
Griffiths, J. and Schroeter, D. Introduction to quantum mechanics. Cambridge university press, 2018.
Burow, A. M. and Sierka, M. Linear scaling hierarchi- cal integration scheme for the exchange-correlation term in molecular and periodic systems. Journal of Chemi- cal Theory and Computation, 7(10):3977-3144 2011. doi: 10. 1 URL https:ttttt.mm /dwi..ioo/////o/oooo .org/ 10. 1021 ct204412r.
Han, S., Pool, J., Tran, J., and Dally, W. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
PMID: 26598153.
Cances, E. and Le Bris, C. On the convergence of scf algo- rithms for the hartree-fock equations. ESAIM: Mathemat- ical Modelling and Numerical Analysis, 34(4):749-774, 2000.
Hassibi, B. and Stork, D. G. Second order derivatives for network pruning: Optimal brain surgeon. Advances in Neural Information Processing Systems (NeurIPS), 1993.
Hohenberg, P. and Kohn, W. Inhomogeneous electron gas. Physical review, 136(3B):B864, 1964.
Chen, T., Luo, S., He, D., Zheng, S., Liu, T.-Y., and Wang, L. Geomformer: A general architecture for geometric molecular representation learning. 2023.
Jain, A., Shin, Y., and Persson, K. A. Computational predic- tions of energy materials using density functional theory. Nature Reviews Materials, 1(1):1-13, 2016.
Du, W., Zhang, H., Du, Y., Meng, Q. Chen, W., Zheng, N., Shao, B., and Liu, T.-Y. Se (3) equivariant graph neural networks with complete local frames. In Interna- tional Conference on Machine Learning, pp. 5583-5608. PMLR. 2022.
Kohn, W. and Sham, L. J. Self-consistent equations includ- ing exchange and correlation effects. Physical review, 140(4A):A1133, 1965.
Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H., and Wang, H. ecometry-enhaneee molecular representation learning for property prediction. Nature Machine Intelligence, 1(2):177-134, 2022.
Kudin, K. N., Scuseria, G. E., and Cances, E. A black- box self-consistent field convergence algorithm: One step closer. The Journal of chemical physics, 116(19):8255- 8261, 2002.
10

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Laqua, H., Dietschreit, C. B., Kussmann, J., and Ochsen- feld, C. Accelerating hybrid density functional theory molecular dynamics simulations by seminumerical in- tegration, iidettitttt approximation, and graphics processing units. Journal of Chemical The- ory and Computation, 18(10)6010-60200 2022. doi: 10. 2..1021/aessjccc:220000000 URL https:/ :
Passaro, S. and Zitnick, C. L. Reducing so (3) convolu- tions to (2) for efficient equivariant gnns. In Inter- national Conference Machine Learning, pp. 27420- 27438. PMLR, 2023.
Payne, M. C., Teter, M. P., Allan, D. C., Arias, T., and Joannopoulos, a. J. Iterative minimization techniques for ab initio total-energy calculations: molecular dynamics and conjugate gradients. Reviews of modern physics, 64 (4): 1992.
10. 1el.1021/aaa 10/21/acss 2c00502 PMID: 36136665.
LeCun, Y., Denker, J. S., and Solla, S. A. Optimal brain damage. In Advances in Neural Information Processing Systems (NeurIPS), pp. 598 603, 1990.
Peng, Y. and Zhang, R. Towards effective sparsification for molecular graph property prediction. Journal of Chemical Information and Modeling, 2022.
Li, H., Wang, Z., Zou, N., Ye, M., Xu. R., Gong, X., Duan, W., and Xu, Y. Deep-learning density functional theory hamiltonian for efficient ab initio electronic-structus cal- culation. Nature Computational Science, 2(6):367-377, 2022.
Schütt, K. Kindermans, P.-J., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R. Schnet: A continuous- filter convolutional neural network for model- ing quantum interactions. Advances in neural information processing systems, 30, 2017.
Li, Y., Wang, Y., Huang, L.,, Yang, H., Wei, X., Zhang, J., Wang, T., Wang, Z., Shao, B., and Liu, T.-Y. Long-short- range message-passing: A physics-informed framework to capture non-local interaction for scalable molecular dynamics simulation. In International Conference on Learning Representations, 2024.
Schütt, K. T., Gastegger, M., Tkatchenko, A., Müller, K.- R., and Maurer, R. J. Unifying machine learning and quantum chemistry with a deep neural network for molec- ular wavefunctions. Nature communications, 10(1):5024, 2019.
Li, Y., Huang, L., Ding, Z., Wang, C., Wei, X., Yang, H., Wang. Z., Liu, C., Shi, Y., Jin, P., et al. E2former: A linear-time efficient and equivariant transformer for seal- able molecular modeling, arXiv e-prints, pp. arXiv-2501, 2025a.
Sun, Q., Liu, J., Zhang, X., and et al. Recent developments in the pyscf program package. The Journal of Chemical Physics, 1 153()):024109 2020.
Szabo, A. and Ostlund, N. S. Modern quantum chem- istry: introduction to advanced electronic structure the- ory. Courier Corporation, 2012.
Li, Y., Xia, Z.. Huang, L., Wei, X., Harshe, S., Yang, H., Luo, E., Wang, Z., Zhang, J., Liu, C., Shao, B., and Gerstein, M. Enhancing the scalability and applicability of kohn-sham hamiltonian for molecular systems. In International Conference on Learning Representations, 2025b.
te Vrugt, M., Löwen, H., and Wittkowski, R. Classical dynamical density functional theory: from fundamentals to applications. Advances in Physics, 6212)122 2020.
Liao, Y.-L. and Smidt, T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint w-XXv:2206.11990 2022.
Thölke, P. and De Fabritiis, G. Torchmd-net: Equivariant transformers for neural network based molecular poten- tials. arXiv preprint arXiv:2202.025:: 02522, 2022.
Liao, Y.L.,, Wood, B., Das, A., and Smidt, T. Equiformerv2: Improved equivariant transformer for scaling to higher- degree representations. arXiv preprint arXiv:2306. /2059, 2023.
Unke, O., Bogojeski, M., Gastegger, M., Geiger, M., Smidt, T., and Müller, K.-R. Se (3)-equivariant prediction of molecular wavefunctions and electronic densities. Ad- vances in Neural Information Processing Systems, 34: 14434-14447, 2021.
Liu, C., Ma, X., Zhan, Y., Ding, L., Tao, D., Du, B., Hu, W. and Mandic, D. P. Comprehensive graph gradual prun- ing for sparse training in graph neural networks. IEEE Transactions on Neural Networks and Learning Systems, 2023.
Wang. J. and Baerends, E. J. Self-consistent-fie fiiel method for correlated many-electron systems with an entropic cumulant energy. Physical review letters, 1 128(1) 12((1):013001, 2022.
Musaelian, A.. Batzner, S., Johansson, A., Sun, L., Owen, C. J., Kornbluth, M., and Kozinsky, B. Learning local equivariant representations for large-scale atomistic dy- namies. Nature Communications, 14(1):579, 2023.
Wang. N., Lin, C., Bronstein, M., and Torr, P. Towards flexible, efficient, and effective tensor product networks. NeurIPS 2023 Workshop: New Frontiers in Graph Learning, 2023.

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Wang, Y., Li, S., He, X., Li, M., Wang. Z., Zheng, N., Shao, B., Liu, T.-Y., and Wang, T. Visnet: equivariant geometry geometry--ehanceed graph neural network with vector- scalar interactive message passing for molecules. arXiv preprint arXiv:2210. 16518, 2022.
Wang, Z., Liu, C., Zou, N., Zhang, H., Wei, X., Huang, L., Wu, L., and Shao, B. Infusing self-consiseency into density functional theory hamiltonian prediction via deep equilibrium models. arXiv preprint arXiv:2406.03794 2024.
Yu, H., Xu. Z., Qian, X., Qian, X., and Ji, S. Efficient and equivariant graph networks for predicting quantum hamiltonian. In International Conference on Machine Learning, pp.
PMLR, 2023.
Yu, H., Liu, M., Luo, Y., Strasser, A., Qian, X., Qian, X., and Ji, S. Qh9: A quantum hamiltonian prediction benchmark for qm9 molecules. Advances in Neural Information Processing Systems, 36, 2024.
Zang. X., Zhao, X., and Tang, B. Hierarchical molecular graph self-supervised learning for property prediction. Communications Chemistry, 6(1):34, 2023.
Zhang, H., Liu, C., Wang, Z., Wei, X., Liu, S., Zheng, N., Shao, B., and Liu, T.-Y. Self-consistency training for fuuetional-thhorry hamiltonian prediction. In In- ternational Conference on Machine Learning, pp. 59329- 59357. PMLR, 2024.
Zhong, Y., Yu, H., Su, M., Gong, X., and Xiang, H. Transfer- able equivariant graph neural networks for the hamiltoni- ans of molecules and solids. npj Computational Materials, 9(1) 91): 2023.
Zhou, G., Lubbers, N., Barros, K., Tretiak, S., and Neb- gen, B. Deep learning of dynamically responsive chemi- cal hamiltonians with semiempirical quantum mechanics. Proceedings of the National Academy of Sciences, 119 (27):221203331 19, 2022.
12

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
A. Additional Experimental Settings
A.I. Training Setup.
The bascline model QHNet was using its default setting. The training setting for SPHNet shown in Table 7, most training settings were set to align with the QHNet's experiment to make a fair comparison. The batch size for training 10 on the MD17 dataset (except Uracil which was set to 5), 32 on the QH9 dataset, and 8 on the PubChemQH dataset. The training step was 260,000 for the QH9 dataset, 200,000 for the MD17 dataset, and 300,000 for the PubChemQH dataset. There was a 1,000-step warmup with the polynomial schedule. The maximum learning rate was 1e-3 for the QH9 dataset and PubChemQH dataset and was 5e-4 for the MD17 dataset. The training and test sets in the QH9 dataset were split in the same manner as their official implementation, including four different split sets. The MD17 dataset was randomly split into the train, validation, and test sets of the same size as the QHNet experiment. The PubChemQH datasets were split randomly into train, validation, and test sets by 80%, 10,, and 10%. The batch size for speed and GPU memory test was set to the maximum number that the GPU can hold maximize its capability for fair comparison.
Table 7. Model Training Settings for Different Datasets,
Dataset
Batch Size
Training Steps
Warmup Steps
Learning Rate
Lana
Sparsity Rate
TSS Epoch
Train/Val/Test Split Method
MD17 Water
10
200,000
1,000
5e4
4
0.1
3
Random (500/500/990))
MD17 Ethanol
10
200,000
L,000
5e-4
4
0,3
Random (25000/5044500)
MD17 Malondialdehyde
10
200,000
1,000
5e-4
4
0.3
3
Random (2500)/5001477)
MD17 Uracil
200,000
1,000
Se-4
4
0.3
3
Random (25000/500/4500)
QH9
32
260,000
1,000
lo-3
4
0.4
3
Official (4 splits)
PuhChemQH
300,000
1,000
le-3
6
0.7
3
Random (40257/5032/533))
A.2. Hardware and Software.
Our experiments are implemented based on PyTorch 2.1.0, PyTorch Geometric 2.5.0, and e3nn (Geiger et al., 2022) 0.5.1. In our experiments, the speed and GPU memory metrics are tested on a single NVIDIA RTX A6000 46GB GPU. The complete training of the models is carried on 4 x 80GB Nvidia A100 GPU. Our code is available in supplementary material.
A.3. Problem Formulation.
The loss function of SPHNet is the MAE plus MSE between predicted matrix Hpred and ground truth matrix Href-
loss = AME(H,r Hprea) + ..MS H Hprea).
(15)
SPHNet predicts the gap between the Hamiltonian matrix and the initial guess. The predicted target can be written as:
AH = Hres Hinit:
(16)
where Hini is the initial guess of Hamiltonian matrix get with the pyscf (Sun et al.. 2020) by function init guess by minao(.). On one hand, we observed that the scale and variance of AH are approximately an order of magnitude smaller than those of Hee on large datasets such PubchemQH, which facilitates more stable learning across different datasets. On the other hand, the computational cost for calculating the initial guess is very cheap, it would not take long to obtain all the initial guesses of molecules in the dataset. Thus, we switched the training target to the AH and added the Hnst back to the predicted AH to get the full Hamiltonian matrix.
A.4. Evaluation Metric
Mean Absolute Error (MAE) of Hamiltonian Matrix H: This metric quantifies the Mean Absolute Error in relation to ground truth data obtained from Density Functional Theory (DFT), accounting for both diagonal and Non-Diagonal elements that reflect intra- and inter-atomic interactions, respectively. This can be presented as:
H mean
H peed
H Hg.))
(17)
where Hpod is the predicted Hamiltonian matrix and Hgr is the reference Hamiltonian matrix.
Mean Absolute Error (MAE) of Occupied Orbital Energies e: This metric evaluates the MAE of the occupied orbital energies, specifically the Highest Occupied Molecular Orbital (HOMO) and Lowest Unoccupied Molecular Orbital (LUMO)
13

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
levels. The accuracy of these critical properties is assessed by comparing the energies derived from the predicted Hamiltonian matrix against those obtained from the reference DFT calculations. The MAE can be expressed as:
pred
ck'l
(18)
pred and ex are the predicted and reference energy of the k-th occupied where M represents the number of occupied orbitals, orbital.
Cosine Similarity of Orbital Coeflicients v: To evaluate the similarity between the predicted and reference electronic wavefunctions, we calculate the cosine similarity of the coefficients corresponding to the occupied molecular orbitals. This similarity metric pivotal for understanding and predicting the chemical properties of the system. The cosine similarity can be defined as:
Σ; prod of
Siebreed 2)
(19)
where ||vl|| denotes the norm of the vector vs viset and vfte are the Coefficient of the i-th occupied molecular orbital in the predicted and reference wavefunction
B. Additional Experimental Results
B.1. Ablation Study for blocks of SPHNet
conducted an ablation study to evaluate the effect of different modules in the SPHNet architecture. Specifically, the standard SPHNet model has 4 Vectorial Node Interaction blocks, Spherical Node Interaction blocks, and Pair Construction blocks. We removed all the sparse gates and reduced the number of these three kinds of modules to 1, respectively, and observed the model performance. As shown in the table below, we found that both the Vectorial Node Interaction block and the Spherical Node Interaction block significantly affect the model performance, indicating that the design of architectures with progressively increased irreps orders has an important positive impact on the models. Interestingly, we found that removing one Pair Construction block would not strongly affect the model accuracy, suggesting that there actually room to further speed up the model. We will explore this further in our future work.
Table 8. The effect of blocks' numbers in SPHNet the PubChemQH dataset.
Sparse Pair
Sparse Gate
Vectorial Node Interaction block
Spherical Node
Pair Construction block
H (10- E..
Model
Gate
Interaction block
SPHNet
X
X
2
2
86.35
SPHNet
X
2
2
96.01
SPHNet
2
97.35
SPHNet
X
2
89.17
14

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
B.2. The Effect of Sparity Rate the Computational Cost
We analyze the effect of sparsity rate on the model performance in Section 5.4. Here put the complete results of training speed and GPU memory in the PubChemQH dataset under every sparse rate. The results showed that the computational cost reduced linearly the sparsity increased, which was in line with our expectations.
Table 9. The effect of sparsity rate v.s. computational cost on PubChemQH dataset
Sparsity
H [10 6 EEn]
Speed [Sample/Sec]
Memory [GB/Sample]
78.48
0.78
17.02
10
79.90
0.96
12.49
20
81.29
1.05
11.21
30
85.33
1.24
10.93
40
86.65
1.50
6.84
50
90.52
1.86
6,22
60
91.90
2.31
5.12
70
97.31
3.12
5.62
80
158.70
3.92
4.68
90
183.52
5.13
3.38
B.3. Single Tensor Product Time Scaling with Sparsity Rate
We test the time consumption for a single tensor product under different sparsity rates. As shown in Fig.6, the time consumption decreases linearly with increasing sparsity, in line with our expectations. This experiment suggested that the sparsity strategy can accelerate the tensor product speed on a linear scale.
Singlc Tensor Product Time Scaing with Sparsity Rate
0.150
Time :eec
0.125
0.100
0.075
0.050
0.025
0%
10%
20%
30% 40% 50%
60% 70% 80% 90% Sparsity Rate
Figure 6. The time consumption for single tensor product with scaling sparsity rates.
B.4. Selected Combination Set Details
Here presented the details of the selected combination sets in all Sparse TP Gate in the SPHNet. We first analysis the output order of the selected combinations. As shown in Fig.7, the selected combinations' output order in the Non- Don-Daagonal block shows a monotonically decreasing trend as analyzed in section 5.5, suggesting the weakening of a single combination as the number combinations within the same order increase. In the Diagonal block, we didn't observe any apparent pattern, this might be because all the combinations have similar effects on the model performance. However, this requires further analysis the future.
15

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
PusCtem@rt Lays? Day
Rechen Laven! Deg
PLSCsendee Layar Nas
6-6-
0.00
0.10
008
67-
25
64--
13.2
Combrater Deturiate
Comer O.00.
MOFFlawed
VE-17 MannDag
3.0
Combraten Ootarodder
Cembiaate Arde
Figure 7. The proportions each output order being selected in difffernt tensor product gates.
We then check the combination weights for each the selected combinations. The details are shown in Table 10, Table 11 and Table 12. These results are from the Sparse Tensor Product Gate before the first Non-Diagonal block. For the PubchemQH dataset, there are a total of 175 combinations in the tensor product (maximum irreps order 6), and the Sparse Pair Gate selects 30% of these combinations. For the QH9 dataset and MD17 dataset, there total of 65 combinations in the tensor product (maximum irreps order 4), and the Sparse Pair Gate selects 60% of these combinations for QH9, 70% of these combinations for MD17.
Table 10.
Weight of selected combinations in PubChemQH
Rank
Comb
Weight
Rank
Comb
Weight
Rank
Comb
Weight
(4, 2, 3)
1.1422
19
(2, 5)
1.0736
37
(4, 4, 2)
1.0488
(5,4 2)
1.1398
20
(3, 0)
1.0725
38
(3, 1, 2)
1.0530
3
(0, 4)
1.1160
21
(4,3,5)
1.0710
39
(4.4.
1.0488
(0, 3)
1.1093
22
(0,5,5)
1.0710
2,,
1.0486
(3,
1.1081
23
(2,3.3)
1.0682
41
1.0485
(6, 6,
1.1078
24
6,,
1.0664
42
(6.
1.0482
2,)
1.1031
25
6,)
1.0651
43
33,3,4)
1.0478
8
(6, 3,
1.1028
26
(6.3.3)
1.0647
44
5,
1.0464
(4.0.4)
1.1015
27
1,
1.0607
(,
1.0452
10
(1,4,3)
1.0978
28
(3,5,4)
1,0600
46
(5, 5)
1.0417
11
(2)5.3)
1,0917
29
(6,4.
1.0596
47
4)
1.0377
12
(5,
1.0903
30
2,
1.0579
48
(4,
1,0300
13
1,
1,0856
31
(3,2.1)
1.0566
49
(1.
1.0283
(5,
1.0854
32
(5,4.4)
1.0545
(6. 6.
1.0281
15
44
1.0842
33
65,
1.0537
51
35,3,4)
1.0274
16
(2,4.2)
1.0835
34
(4,3,3)
1.0535
52
5,
1.0248
17
(3.
1.0829
35
5,))
1.0534
53
(2, 2,
0.9987
18
(1,
1.0799
36
(3, 1,2)
1.0530
16

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Table H.
Weight of selected combinations in QH9
Rank
Comb
Weight
Rank
Comb
Weight
Rank
Comb
Weight
(1,
1.1748
14
(3, 1,
1.0758
27
(3, 0, 3)
1.0154
(2, (2,4.2)
1.1722
15
(2, 0,
1.0715
28
(2,3,4) 4)
1.0150
1)
1.1661
16
(2,
1.0696
29
(4,4,
1.0145
1)
1.1433
17
(3.2.3)
1.0577
30
(1,0,1)
1.0145
(3,4,3)
1.1284
18
3,
1.0530
31
(0,
1.0133
6
1,3)
1.1240
19
(0, 0,
1.0406
32
(3,4,2)
1.0117
(3,3.2)
1.1233
20
(4,
1.0383
33
(1,3,3)
1.0108
8
(2,
1.1163
21
(4,
1.0336
34
2,
1.0099
1,)
1.1157
21
1.0319
35
(4,
1.0097
10
(4,4
1.1085
23
1,
1.0276
36
(2,
1.0090
11
(0,3.3)
1.0972
24
(3.4.4)
1.0213
37
(3,
1.0059
12
1)
1.0853
25
(0,4,4)
1.0198
38
1, 4)
1.0057
13
(0,
1.0788
26
(3,
1.0154
39
(4, (4,0,4)
0.9912
Table 12. Weight of selected combinations in MD17 (Ethanol)
Rank
Comb
Weight
Rank
Comb
Weight
Rank
Comb
Weight
(3, (3,2,1 I)
1.0555
16
(3, I)
1.0213
31
(3, 3, 0)
0.9894
(3,4,
1.0674
17
4.4,2
1.0206
32
4, 1)
0.9891
(2, 1)
1.0608
18
2,
1.0141
33
1,
0.9886
2,
1.0592
19
1,)
1.0119
34
(2,
0.9878
5
1)
1.0449
20
0,
1.0090
35
(3,
0.9856
(0,
1.0447
21
(4,4.0)
1.0037
36
(4,3,3)
0.9851
(0,3,3)
1.0446
22
1,3)
1.0016
37
(3,3,3)
0.9811
22,4,
1.0418
23
(3,
1.0003
38
(3, 4,44
0.9802
12)
1.0417
24
(3,4.2)
0.9995
39
(3, 4,3)
0.9795
10
(1,3,2)
1.0408
25
4,)
0.9991
40
(1,3,4)
0.9787
(0,
1.0399
(3, 0,
0.9978
(2, 2,
0.9769
12
1,,),2)
1.0349
27
2,
0.9972
42
(3, 1,3)
0.9763
13
(1,
1.0330
28
(3,2,3)
0.9928
43
(2, (2,4,3)
0.9736
14
(4,2.2)
1.0301
29
3,43,4)
0.9920
44
(2, 0.9730
15
(1,4,3)
1.0295
30
(0, 0,
0.9911
45
2, 0.9719
B.5. GPU Memory when Scaling with Increasing Size
We used molecules with different sizes to test the GPU memory consumption of SPHNet and the baseline model QHNet. As shown in Fig.8, we observed that the GPU memory consumption SPHNet increased much slower than the baseline model. When the molecule had more than 1800 atomie orbitals, the baseline model reached the maximum GPU memory, which is times the memory SPHNet needs, while the SPHNet can handle molecules with around 2900 atomic orbitals, making possible to train larger molecule systems.
C. Additional Preliminary
C.1. Spherical Harmonic Function
The spherical harmonic function in SPHNet is used to project the atomic vector rīj into spherical space. real basis for spherical harmonics Yim : 52 + R can be expressed in relation to their complex analogues Y!! S2 - C. This relationship
formulated through the following equations:
17

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
130
SPHNet QHNect
GPU memonny Sempele (lope))
50-
17
20
x5.19
10
x0.50
2:04
Number aromic orbiiate of Hamilooiian
Figure The GPU memory consumption with scaling molecule size.
V2 -|m| [m]
Y
(-1)
/m/ m < 0.
y!
Yo:
m = 0.
(20)
yi v2 [m] + (-1) if m > 0.
To ensure both consistency and standardized treatment throughout the analysis, the Condon-Shortley phase convention is frequently utilized. The inverse relationships that define the complex spherical harmonics Y S - in terms of the real spherical harmonics Yim 2 >> are outlined as follows:
(Yll)
iY i
if m < 0.
v2 mm]
y! m
Yi,0,
if 0.
(21)
(-1)m
1 (Y2. + ii),-imm) vi
if m > 0.
Established theories regarding the analytical solutions the hydrogen atom indicate that the eigenfunctions corresponding to the angular component of the wave function are represented as spherical harmonics. Notably, the absence of magnetic interactions, the solutions to the non-relativistic Schrödinger equation can also be represented as real functions. This aspect highlights the common implementations by eliminating the need for complex algebra. It is crucial to acknowledge that these real-valued functions occupy a functional space identical to that of their complex counterparts, thus preserving generality and completeness in the solutions.
of real-valued basis functions in electronic structure calculations, as simplifies software
C.2. Irreps and Tensor Products
Irreps Representation: SPHNet employs the special orthogonal group SO(3) to capture the essential 3D rotational symmetries intrinsic to molecular structures. It leverages the irreducible representations (irreps) of SO(3), indexed by an integer which are associated with spherical harmonic functions Ytm. These spherical harmonics impart rotational characteristics to the feature vectors, thereby ensuring that the model maintains rotational invariance and facilitates consistent assessment of geometric properties.
Tensor product is the core operation in SPHNet. It facilitates interactions among irreps linked to distinct order (angular momentum) / and enhances expressiveness within the model. This operation merges two irreps with order 1 and 12 to generate new irrep characterized by order 1g. The expansion is accomplished by utilizing Clebsch-Gordan coefficients, weighted by :m..
Tensor Product:
2
Es m3
w L2.22 (ll.z)) (m) (2) I Veng
Wm..ma (ex.mm)).((2,mm))
(22)
E mg
18

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
This mathematical approach allows for the amalgamation of complex features derived from simpler ones, effectively capturing the nuanced interactions of angular momentum in the molecular context.
D. Additional Model Details.
A
Z
Norm
Norm
Short range madule
x4
Sph Linear
Sph Linear
X0
Sparse
Long module
x4
Tensor Product
Sparse Pair Gate
Norm
Linear
Sph
W00
RBF
Linear
Linear
Yi
Sparse Pair Gate
B
Product
Spherical harmonica
r'u
x1
RBF
f
Sph Lineer
Tensor expansion
hi
Norm
Sph Linear
Linear
X7
Linear
Sparse TP Gate
Symmetry
fy
fo
Sph Linear
Tensor expansion
hr
Tensor Product
X: Linear
Norm
Linear
Sph
f
Figure 9. The components of SPHNet. (A) The Vectorial Node Interaction Block, which uses a long-short range message-passing mechanism. (B) The Expansion block. (C) The Spherical Node Interaction Block. (D) The Diagonal block in the Pair Construction block. (E) The Non-Diagonal block in the Pair Construction block.
D.1. RBF
The RBF here refers to the Radial Basis Functions and used to construct nonlinear maps of vectors r. We used Exponential Bernstein Radial Basis Functions here to process the vector of two atoms, and the mathematical formulation as below:
RBF(r) Fest (r: cutoff) exp (log + n (-ar) v log(-exp(< 1)),
(23)
and the cutoff function F Four limit the range of action of RBF, beyond which inputs will be ignored or reduced in impact:
Feur (x, cutoff)- = exp
2 (cutoff - (cutoff+
if cutoff.
(24)
0,
otherwise.
D.2. Sph Linear
The Sph Linear used in the SPHNet was implemented through the Liner function in the e3nn (Geiger et al., 2022). It is an 0(3) linear operation that takes an irrep as input and outputs the linear combination of input. Specifically, the feature of each order in the input irreps will be mapped to the feature of specific order of the output irreps through a learnable matrix. We didn't 't set the instruction attribution in the e3nn's 's (Geiger et al., 2022) function, which resulted in fully connected instruction and each order in the input and output irreps will have a connection.
19

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
D.3. Vectorial Node Interaction Block
As shown in Fig, the Node Vectorial Node Interaction Block receives the atomic numbers Z and 3D coordinates of the molecular system Fij imput and extracts the representations x; for node i. In this block, we adopted the long-short range sassage-passing mechanism (Li ct al., 2024) to capture the long-range interaction between each atom and obtain more informative node representations. There are 4 short-range message-passing modules and 4 long-range message-passing modules. Specifically, for the short-range message-passing module, the module passes the neighboring atoms' features to the center atom as defined below:
x = Fshort (x} +
Σ (EV(()
(25)
where Fshort is the linear layer, m, is the message from atom j to atom i, and the (i) is the neighboring atoms of atom i. For long-range message- message-passing module, the module passes the neighboring groups' feature to the center atom as defined below:
x = Fiong (xf +
E m
(26)
JEN(()
where Flong is the linear layer, (() is the neighborhood group of atom i on the atom-fragment bipartite graph, and maj is the message from neighboring group j to atom
D.4. Spherical Node Interaction Block
As shown in FCg. the Spherical Node Interaction Block takes the node feature x; and 3D coordinates of the molecular system Fig as input and outputs the new high-order feature through a tensor product operation. Specifically, the input atom feature X; and its neighboring atom feature are first fed the Sparse Pair Gate and obtain the selected important pair set (x², (rT))) and their pair weight W This weight feature is then multiplied with the RBF of rīj and gets the final tensor product weight defined in Equation 9 and 10, The atom feature is finally doing tensor product with the high-order spherical harmonics projection of the selected (rij), as defined in Equation 12. Last, the ascended node feature è; is outputted after normalization and sph linear:
Layernorm(x) xx))
ifl =
x
x
Tomen
E/er
Frna
2) 2). 21+1 Em=_, x
ifl >
(27)
x; = Sphlinear(x:).
(28)
There are two Spherical Node Interaction Blocks in the SPHNet. The first block increases the maximum order input irreps from zero to the required order without utilizing the pairs selected by the Sparse Pair Gate. The second block inputs and outputs irreps with the same maximum order, and employs the pairs (x/, (17))) selected by the Sparse Pair Gate.
D.5. Pair Construction Block
The Hamiltonian matrix contains the relational information of each pair of atoms in the molecular system. The objective of the Pair Construction Block is extend the model's capacity to consider diagonal node pairs f,z and non-diagonal node pairs fú in the Hamiltonian matrix. Therefore, there two subblocks in the Pair Construction block. the Diagonal block for the diagonal feature and the Non-Diagonal block for the Non- -iagonal feature. There are two Pair Construction blocks in the SPHNet, to ensure that each subblock h:g in the Hamiltonian matrix is covered at least once, the first Pair Construction block does not use the pair selected by Sparse Pair Gate.
D. DIAGONAL BLOCK
As shown in FDg.9D, the Diagonal block first uses two separate sph linear layers to get two different representations from the input node feature, and then uses a self-tensor product operation obtain f., Here, the combinations in the tensor product selected by the Sparse Tensor Product Gate:
8; Fs, (X:), Xi, Fsz (x::,
(29)
20

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
Σ (G
W (all. (R co 2
(30)
where x; is the atom feature from the Spherical Node Interaction Block, and Fe, and Fs, are the sph linear layer. W is the combination weight defined in Equation 13. Last, the diagonal feature fü is outputted after normalization and sph linear the same as in the Spherical Node Interaction block.
D.5.2. NON-DIAGONAL BLOCK
For the non-diagonal pairs, the module first uses Sparse Pair Gate select the most valuable pair (x',,',)) and their tensor product weight Waj. as defined in Equation 9. The tensor product weight is then multiplied with the RBF of ri and gets the final tensor product weight as defined in Equation 12. Then, the module calculates the tensor product of x', and x'j, the process is presented in Fig.9E. Note that the combinations in the tensor product are selected by the Sparse Tensor Product Gate. The result is regarded the desired f,y. defined in Equation 14. Last, the non-diagonal feature f.y is outputted after normalization and sph linear the same in the Spherical Node Interaction block.
There are two separate Pair Construction Blocks that receive atom representations from the two SO(3) Convolution Blocks respectively, and the final node pair feature is the addition of these two Pair Construction Blocks' output.
D.6. Expansion Block
Once we have gathered the irreps for both diagonal and non-diagonal pairs, the subsequent step involves constructing the complete Hamiltonian matrix. There are many matrix elements in the Hamiltonian matrix, each matrix element represents the interaction between two orbitals. Specifically, the block denoted as hij captures all interactions between atoms i and j. Since atoms have varying numbers of orbitals, the shape of the pair block hig is different, and must be determined during the construction of the Hamiltonian matrix.
Here, same as the previous work (Yu et al., 2023), introduce an intermediate block My containing the full orbitals, from which we can derive ha by extracting the relevant components based on the specific atom types. For instance, in the MD17 dataset, there are four atoms- C, N, and O-each with its own set of orbitals. The full orbital set consists of 1s, 2s, 3s. 2p, 3p, 3d orbitals, where M E R 1214x14 14, When dealing with the hydrogen atom (H), only the Is, 2s, and 2p orbitals are selected to contribute to the Hamiltonian matrix. By using this strategy, each node's irrep fúy can be converted into an intermediate block Mij with a predetermined structure, irrespective of the atom type. This technique is particularly advantageous as can be easily adapted to various molecules.
To construct the intermediate pair blocks with full orbitals using pair irreducible representations, we apply tensor expansion operation in conjunction with the filter operation. This expansion is defined by the following relation:
to
(Fammo)) (e.me))(e)mm)) Cin::
(31)
where C' denotes the Clebsch- Gorsann coefficients, and symbolizes the tensor expansion which is the converse operation of the tensor product. Specifically, xf @ S can be expressed sum over tensor expansions:
Σ W......ff.
(32)
subject to the angular momentum coupling constraints 1e, - Egl S to S li+ll The filter then takes the atom types as input, producing a weight for each path (fo1, Loz, l:n):
Fffarhaa = s(2,, (loiiRoaddin) = [(Za),
sj
(20)
where Z denotes the embedding of the atom types. The intermediate blocks M are generated by the filter, using the node pair irreducible representations as follows:
M (Mloiihua)
Σ
F (lorskozzin) sp(linee)
(21)
P
21

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
M
Σ Pinie
F ffo Jrlfin
(33)
Here, o' indicates the channel index in the input irreducible representations, and e represents the channel index in M. For instance, there nine channels with (f1, f2) (0,0) and four channels with (41, 62) (1, 1). A bias term is added for the node pair representation when L.n 0.
Last, we derive hiy by extracting the relevant components based on the specific atom types and use these matrix elements to construct the complete Hamiltonian matrix. As shown in FBg. for the Von-Diagonal block in the Hamiltonian matrix, we only construct the node pair features fly from the upper triangular part of the Hamiltonian matrix. The lower triangular part then obtained by symmetrizing the sub-blocks from the upper triangular part.
D.7. Computational Overhead of the Sparse Gates and Scheduler
Compared to the tensor product operation, Sparse Gate and the three-phase learning schedule only introduce minimal computational overhead and have little impact on the overall speed. We would like to discuss this part of computational cost here.
In the three-phase sparsity scheduler, for given unsparsified set U, the additional computational overhead in the first phase has complexity of O(IT)) )) contributed by the R RANDOM()) operation. The second phase has a computational overhead of O(UU] log [U]), arising from the TOP(-) operation. Since we fix the learnable weight matrix and the selected elements, there is no additional computational overhead in the third phase. For detailed information, please refer to Equation 3.
For the sparse TP gate, the computational overhead comes from the element- wise multiplication of two weight vectors (Equation 5), its complexity is O(UU,) O(L*).
For the sparse pair gate, the additional computational overhead mainly comes from the linear layer F2(-) in Equation 7, with its time complexity being the square of its hidden dimension. Other operations, including the inner product (Equation 6) and the weight calculation (Equation 9), are necessary operations in our framework, even without the sparse pairwise gate.
E. Additional Related works
E.1. SE(3) Equivariant Neural Network
The SE(3) equivariant neural network is one of the most used models in the field of AI for chemistry (Fuchs et al., 2020; Du al., 2022; Musaelian al., 2023; Liao & Smidt, 2022; Liao al., 2023; Batzner al., 2022; 2023). The fundamental feature of SE(3) equivariant neural network is that all the features and operations in the model are SE(3) equivariant. This is achieved by using irreducible representations (irreps) and functions of geometry built from spherical harmonics. In the network, the equivariant features propagate through each layer and interact with other features by equivariant operations, and finally obtain the desired SE(3) equivariant quantity.
Transfoereer (Fuchs al., 2020) proposed a novel self-attention mechanism for 3D point clouds and graphs that ensures robustness through continuous 3D -oo-translation equivariance, which is widely used in the field of quantum chemistry. The Equiformer (Liao & Smidt, 2022) introduced novel graph neural network leveraging SE(3) Transformer architecture based on irreducible representations to predict molecule property. It demonstrated strong empirical results by incorporating tensor products and new attention mechanism called equivariant graph attention. The EquiformerV2 (Liao et al., 2023) is an improved version of Equiformer, which scales effectively to higher-order representations by replacing SO(3) convolutions with efficient eSCN convolutions (Passaro & Zitnick, 2023), and outperforming the traditional network such as GemNet (Gasteiger et al., 2022) and Torchmd-Net (Thölke De Fabritiis, 2022) molecular energy and force prediction and other downstream tasks. Allegro (Musaelian et al., 2023) used a strictly local, equivariant model to represent many-body potential using iterated tensor products of learned representations without atom-centered message passing. It demonstrates remarkable generalization to out-of-distributio data and accurately recovers structural and kinetic properties of amorphous electrolytes in agreement with ab initio simulations.
These SE(3) equivariant neural networks greatly improve the performance of Artificial Intelligent in the field of quantum chemistry. However, the computational complexity of tensor product operation greatly reduces the efficiency of these SE(3) equivariant models. The eSCN (Passaro & Zitnick, 2023) presents an efficient method to perform SO(3) equivariant
22

---

Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity
convolutions. It reduces the computational complexity by aligning node embeddings' primary axis with edge vectors, transforming the SO(3) convolutions into mathematically equivalent SO(2) convolutions, which decreases complexity from O O(L*). E2Former (Li et al., 2025a) introduces equivariant and efficient transformer architecture that incorporates the Wigner 6j convolution (Wigner 6j Conv). By shifting the computational burden from edges to nodes, the Wigner 6j Conv reduces the complexity from O Fl) to O(VI) while preserving both the model's expressive power and rotational equivariance.
E.2. Hamiltonian Matrix Prediction
Predicting the Hamiltonian matrix (also known as electronic wavefunctions) is gradually gaining more and more attention because of the wide range of potential application scenarios of the Hamiltonian matrix. There are more and more works trying to solve the problem predicting the Hamiltonian matrix using deep learning techniques.
The SE(3) equivariant neural network is of the most used models in the field of Hamiltonian matrix prediction (Unke al., 2021; Yu al., 2023; Gong et al., 2023; Atz et al., 2021). The PhiSNet (Unke et al., 2021) is the first method that primarily focuses on Hamiltonian matrix prediction. It leverages SE())--quiaarian operations through its whole architecture to predict molecular wavefunctions and electronic densities, and can reconstruct wavefunctions with high accuracy. The main problem for PhiSNet is its inefficiency due to the large amount of tensor product operations. To solve this problem, QHNet (Yu et al., 2023) proposed a model with careful design to greatly reduce the number of tensor products and improve the efficiency of Hamiltonian prediction. The DeepH-E3 (Gong et al., 2023) an E(3--qquivrriant model that preserves Euclidean symmetry even with spin-orbit coupling. The method allows for efficient and accurate electronic structure calculations of large-scale materials by learning from small-sized DFT data, significantly reducing computational costs.
There are also other works that try to enhance the prediction ability from different aspects. DeepH (Li et al., 2022) introduces a local coordinate system defined for each edge according to its local chemical environment to handle the gauge (or rotation) covariance of the DFT Hamiltonian matrix. This allows the Hamiltonian matrix blocks to be invariant under rotation when transformed into the local coordinate system. Self-connistency training (Zhang et al., 2024) leverages the self-consistency principle of density functional theory (DFT) to train a model without requiring the labeled Hamiltonian matrice. This enhances the generalization and efficiency and reduces reliance on costly DFT calculations for supervision. The DEHQ (Wang et al., 2024) integrates Deep Equilibrium Models (DEQs) to predict Hamiltonians. It inherently captures the self-consistent nature of Hamiltonians, a critical aspect often overlooked by traditional machine learning methods. By employing DEQs, the model circumvents the need for iterative DFT calculations during training. WANet (Li et al., 2025b) introduces a scalable deep learning model for Hamiltonian prediction. By leveraging novel Wavefunction Alignment Loss (WALoss), the model notably reduces total energy error derived from the predicted Hamiltonian matrix and accelerates SCF calculations with the predicted Hamiltonian matrix.
23


Reviews:
[{"review_id": "wDyc7UNNjq", "reviewer": "Reviewer_PzTP", "summary": "This paper introduces SPHNet, an SE(3) equivariant graph neural network designed to efficiently and scalably predict Density Functional Theory (DFT) Hamiltonian matrices. The core contribution is the incorporation of adaptive sparsity to address the significant computational cost associated with high-order tensor products (TP) operations in equivariant networks, which limits their application to large molecular systems. SPHNet employs two novel mechanisms: a Sparse Pair Gate to filter unimportant node pairs and a Sparse TP Gate to prune less significant interaction components within the tensor products. To manage the sparsity dynamically during training, a Three-phase Sparsity Scheduler (random, adaptive, fixed phases) is proposed to ensure stable convergence while achieving high sparsity levels (up to 70%). The main findings reported are that SPHNet achieves state-of-the-art accuracy on the QH9 and PubchemQH datasets while demonstrating significant computational improvements both in speedup and memory usage.\n\n## Update after rebuttal\n\nThank you for your thoughtful rebuttal. I appreciate the clarifications provided regarding your methodology. I look forward to the revised manuscript.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your valuable comments and suggestions. Below are our detailed responses.\n\n**Other Strengths And Weaknesses**\n\nThank you for pointing that out—your notation is the more precise way to express it. $C$ should indeed be $n \\times n_0$ in Equation 1, where $n_0$ corresponds to the number of occupied orbitals. We sometimes write $C$ as $n \\times n$ because the KS equation is typically solved for all eigenstates, including virtual orbitals. We appreciate your keen attention to notation and will make sure to clarify this in the revised version.\n\n**Other Comments Or Suggestions**\n\n1. Thank you for pointing this out. We are carefully considering this, and while we may not be able to provide a definitive new name in the short term, we will update this in the subsequent versions.\n\n2. The higher coefficient rates imply that as the molecular system grows larger, we can apply a higher sparsity rate to further accelerate the model without sacrificing accuracy. We apologize for the confusion and have revised this to 'higher sparsity rates' for better clarity.\n\n3. As you mentioned, compared to the tensor product operation, all these operations only introduce minimal computational overhead and have little impact on the overall speed. However, this discussion is still meaningful, and we will include it in the subsequent version of the manuscript.\n\n    In the three-phase sparsity scheduler, for a given unsparsified set $U$, the additional computational overhead in the first phase has a complexity of $\\mathcal{O}(|U|)$, contributed by the RANDOM(·) operation. The second phase has a computational overhead of $\\mathcal{O}(|U| \\log{|U|})$, arising from the TOP(·) operation. Since we fix the learnable weight matrix and the selected elements, there is no additional computational overhead in the third phase. For detailed information, please refer to Equation 3.\n\n    For the sparse TP gate, the computational overhead comes from the element-wise multiplication of two weight vectors (Equation 5), so its complexity is $\\mathcal{O}(|U_c|)=\\mathcal{O}(L^3)$.\n\n    For the sparse pair gate, the additional computational overhead mainly comes from the linear layer $F_p(\\cdot)$ in Equation 7, with its complexity being $\\mathcal{O}(d_{hidden}|U_p|)$, where $|U_p|$ is always the square of the number of atoms and $d_{hidden}$ is the hidden feature dimension. Other operations, including the inner product (Equation 6) and the weight calculation (Equation 9), are necessary operations in our framework, even without the sparse pairwise gate.\n\n4. -7. Thank you for the comments. We have revised the manuscripts accordingly.\n\n**Questions For Authors**\n\n1. Thank you for the question. We agree with your point. In fact, in our experiments, the number of selected short-distance pairs still constitutes the majority of all selected pairs, which aligns with the nearsightedness of electronic matter, as shown in [Figure A](https://imgur.com/a/0EGNbF8). Each bar represents the fraction of selected pairs within the range from k to k+1 Å relative to the total number of selected pairs. Therefore, when we say *“as the pair length increases, the probability of a pair being selected also rises,”* we mean that pairs at longer distances are indeed retained at a higher proportion in [Figure B](https://imgur.com/a/0EGNbF8)—though their absolute count remains significantly lower compared to short-distance pairs. For example, among 10 pairs at 25–26 Å, 6 may be selected, whereas among 300 pairs at 4 Å, 100 may be selected.\n\n    There are two possible reasons for the observed tendency to retain long-distance pairs. First, long-range interactions, including electrostatic interactions and weak van der Waals forces, are crucial for accurately describing large molecules. Consequently, incorporating such interactions could be beneficial for overall accuracy, a property that has been studied in previous research (Knörzer J,  et al. https://arxiv.org/pdf/2202.06756; Li Y, et al. https://arxiv.org/pdf/2304.13542.). Second, since long-distance pairs constitute only a small fraction of all pairs, once sufficient data has been collected to characterize short-range interactions, selecting more long-distance \"hard samples\" could contribute to improving final accuracy.\n\n    However, due to the lack of large molecular system datasets, we conducted a similar experiment on the QH9 dataset (Figure A-B), where we found that the selected ratio was similar to the proportion of short-distance pairs in the PubChemQH dataset. However, since the maximum atomic distance in this dataset is less than 8 Å, we did not observe a preference for retaining long-distance pairs. This is reasonable, as long-range atomic interactions typically become significant only when the interatomic distance exceeds 12 Å. As larger molecular datasets become available in the future, we hope to conduct further experiments to validate this hypothesis.\n\n2. We have discussed that in the previous responses.", "claims_and_evidence": "The main claims are:\n1.  SPHNet significantly improves computational efficiency (speed and memory) for Hamiltonian prediction compared to existing SE(3) equivariant models.\n2.  This efficiency gain is achieved through novel adaptive sparsity mechanisms (Sparse Pair Gate, Sparse TP Gate, Three-phase Sparsity Scheduler) that reduce tensor product operations.\n3.  SPHNet maintains or improves prediction accuracy despite the induced sparsity, achieving SOTA results on QH9 and PubchemQH datasets. The ablation study supports the claim that significant sparsity (up to 70% for PubChemQH) can be introduced without substantial accuracy loss. \n\nThe evidence provided appears generally supportive:\n* Experimental results on benchmark datasets (QH9, PubchemQH, MD17) are presented, comparing SPHNet against baselines like QHNet and WANet.\n* Quantitative results are reported, claiming up to 7x speedup and 75% memory reduction.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria are appropriate for the problem of accelerating DFT Hamiltonian prediction using physics-informed machine learning.\n\n* The core method involves introducing adaptive sparsity into an SE(3) equivariant GNN architecture. Targeting the tensor product operations, known computational bottlenecks in such networks, with sparsity is a reasonable strategy for improving efficiency. The specific gate mechanisms (Sparse Pair, Sparse TP) and the sparsity scheduler are novel contributions designed to implement this strategy effectively.\n* Evaluation Criteria: The evaluation uses standard metrics:\n    * MAE of the predicted Hamiltonian elements, observables and similarity of coeff matrix elements, compared to DFT calculations.\n    * Inference speed (Samples/Sec) and GPU memory usage for assessing scalability, which is a primary goal of the paper.\n    * Established datasets like QH9, PubchemQH, and MD17 are used. The larger basis set size of Def2-TZVP is important for showing scalability.", "theoretical_claims": "The paper focuses primarily on algorithmic innovation and empirical validation rather than presenting novel theoretical claims or proofs within ML or DFT. The effectiveness of the Three-phase Sparsity Scheduler seems justified empirically.", "experimental_designs_or_analyses": "As stated in methods and evaluation criteria, the experimental design and analyses are reasonable, with good baselines, datasets, metrics and ablation studies.", "relation_to_broader_scientific_literature": "This paper fits within the active research area of applying geometric deep learning to directly predict the converged Hamiltonian under DFT. It improves the efficiency of SE(3) equavariant GNNs by reducing the tensor product operations needed by improving network sparsity. The effectiveness of this technique is shown on DFT Hamiltonian prediction given a molecule graph.", "essential_references_not_discussed": "From the perspective of PIML for DFT, the paper covers the key related areas reasonably well, citing major works in SE(3) equivariance, Hamiltonian prediction, and general network sparsification.", "other_strengths_and_weaknesses": "* The authors show that by using the TSS induced sparsity, the inference speed and memory usage of SE(3) network based Hamiltonian prediction is greatly improved while maintaining similar accuracy as baseline models.\n*  The reported efficiency improvements (speed and memory) are substantial (up to 7x speedup, 75% memory reduction) while maintaining competitive or SOTA accuracy, demonstrating practical value. This is especially important for the PubchemQH results.\n* The paper is generally well-written and clearly structured. The methodology is explained with helpful diagrams. The experiments are comprehensive and well-documented in the main text and appendices.\n\n\nI am a bit confused about the notation used in equation (1). Shouldnt H and S be basis size dependent ($n \\times n$), $\\epsilon$ be orbital set size dependent ($n_o \\times n_o$) and C be ($n \\time n_o$)? I apologize if I misunderstood the notation.", "other_comments_or_suggestions": "* SPHNet is a somewhat famous model used for 3D point cloud analysis (https://arxiv.org/abs/1906.11555). I think the domains are far enough that it can't be confused, but consider changing the name of the model.\n* What do \"higher coefficient rates\" mean in this context? It is used to claim even more performance gains )Page 7, above Table 3).\n* Discuss the computational overhead of the sparsity gates and scheduler themselves, although this is likely small compared to the savings from reduced tensor products.\n* Minor: \n    * Use the draft/review latex environment when submitting for review so that the line numbers are visible.\n    * Page 6: Should be \"The GTO orbital basis is used for MD17 and QH9  is used for ...\"\n    * Page 6: Note instead of Noted\n    * Maintain consistent spacing for citations", "questions_for_authors": "1. the selection probability for pairs *increases* with distance in the Sparse Pair Gate is counter-intuitive compared to typical distance cutoffs and goes against the nearsightedness principle (E. Prodan,& W. Kohn,  Nearsightedness of electronic matter, Proc. Natl. Acad. Sci. U.S.A. 102 (33) 11635-11638, https://doi.org/10.1073/pnas.0505436102 (2005)). Can you explain this further, and does this hold across different systems/basis sets?\n2. What is the computational overhead introduced by the Sparse Pair Gate, Sparse TP Gate, and the adaptive phase of the scheduler relative to a dense model operating at the same FLOP count?", "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": "niooCqRIrq", "reviewer": "Reviewer_hk4X", "summary": "This paper tackles the Hamiltonian prediction task. It proposes to learn a mask to select important pairs in pair-wise interactions in both node interactions and non-diagonal pair construction blocks. Moreover, the paper also use similar techniques to select important paths in the tensor product during pair construction. The proposed SPHNet can reduce the computational cost. The experiments are conducted on QH9, PubChemQH and MD17 datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your valuable comments and suggestions. Below are our detailed responses.\n\n**Claims And Evidence** \n\n1. Thank you for your question. As you noted, SPHNet’s lightweight design allows it to run 1.73× faster than QHNet on the PubChemQH dataset. However, the primary acceleration comes from sparse gates. To isolate their effect, we included ablation study in Appendix B.2. Beside, to further address your concerns, we test these gates on QHNet. Please refer to our answer to the 4th reviewer EBLX for the details due to the length limitation of the response.\n\n2. The computational cost of key sparsity has already been included in Figure 3, where the number above each “$\\times$” symbol represents the training speed at that specific sparsity level. In Appendix Figure 8, we provide a more detailed visualization of how time scales with sparsity. Additionally, we have listed the complete results for training speed and GPU memory usage at each sparsity level on the PubChemQH dataset in the [table](https://imgur.com/a/tM8STT4) below. The results indicate that as sparsity increases, both time and memory costs decrease in an approximately linear manner.  \n\n**Experimental Designs Or Analyses:** \n\n1. We agree with your opinion that the sparse gates cause suboptimal results. As we had analysis in Section 5.3, since the water molecule is very small (only 3 atoms) compared to other molecules, there is not much room for the atom pairs and TP combination reductions. Therefore, spare gates may remove the necessary interaction combinations within the system and cause poor results. \n\n    However, we would like to declare that the sparse gates are not designed for these extremely small molecules, and we are more concerned about their performance on larger molecules, which turn out to be very efficient on the large molecule dataset PubChemQH.\n\n**Other Strengths And Weaknesses** \n\n1. Sorry for the confusion caused by the word \"probability.\" A clearer term would be \"percentage.\" Specifically, in the second phase, TOP(·) selects elements whose weights are within the top $( 1 - k )$ percent, with no randomness involved—selection is purely based on learnable weights. We have revised the manuscript for clarity. \n\n    Regarding path selection, the learnable weights continue to update in phase two, as they participate in subsequent operations (Equations 4 and 5). Only selected paths' weights are updated, while TOP(·)  consistently selects the highest-weighted elements. As weights evolve, the selected paths adjust accordingly, allowing the sparse gate to gradually learn the optimal path set.\n\n2. Thank you for the question. The $w_{ij}$ is a $R^{k}$ vector, where $k$ is the number of elements in the complete set $ U_c= \\{(\\ell_1, \\ell_2, \\ell_3) \\mid \\ell_3 \\in [|\\ell_1 - \\ell_2|, \\ell_1 + \\ell_2] \\} $ in the tensor product. And the $w_{ij}^{\\ell_1, \\ell_2, \\ell_3}$ is a single number from the $w_{ij}$ vector. Note that in Equation 12 we selected the $ w^{\\ell_1, \\ell_2, \\ell_3}_{ij} $ that within the $U^{TSS}_p$ set.\n\n3. We used a fixed scheduler step $( t = 3 )$ across all datasets and experiments, as results remained stable, indicating minimal impact on performance. We recommend setting $( t = 3 )$, but users can adjust it with minimal tuning cost.  \n\n    For sparsity rate, Section 5.4 shows performance remains stable within a reasonable range, with degradation occurring only beyond a certain upper limit, which depends on molecular size. Users can select this parameter based on our experimental results for their specific molecular sizes.\n\n**Other Comments Or Suggestions:** \n\n1. Thank you for pointing that out—the more precise shape should be $ n \\times n_0 $, where $ n_0 $ corresponds to the number of occupied orbitals. The reason we sometimes write $ C $ as $ n \\times n $ is that the KS equation is typically solved for all eigenstates, including virtual orbitals. We will revised this and running title in the revised version.\n\n**Questions For Authors** \n\n1. Please see previous responses.\n\n2. Thank you for raising this question. We addressed this in Section 4.1. To ensure that the combination selection is as unbiased as possible, we initialized the learnable matrix $ W $ (the sparsity weights) as an all-one vector. This means that, at the beginning, all combinations are considered to have the same importance.\n\n3. Thank you very much for pointing this out. Here, the result of inner product is still a irrep with different order of feature, and the superscript stands for the order of irrep feature we used (actually, all the superscripts in our paper stand for the order of irreps feature). But we just found that there is a typo in the writing of the formula. The correct formula should be $<x_i,x_j>^{1:}$, and the superscript ${1:}$ stands for the irrep feature larger than order zero. We are very sorry about the mistake and have fixed it in the manuscript.", "claims_and_evidence": "- The paper compares SPHNet and QHNet (as well as WANet) for speedup. However, the architecture of SPHNet and QHNet are not the same. As a result, the speed-up ratio may not entirely come from the sparse gates.\n\n- Figure 3 shows the effect of sparsity rate on prediction accuracy, it would be interesting to see the effect of sparsity rate v.s. computational cost for training and testing.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The equations mainly server to describe the method.", "experimental_designs_or_analyses": "- The result for water molecule in Table 3 does not seem to be good. Is it because of using the sparse gates?", "relation_to_broader_scientific_literature": "They are adequately discussed.", "essential_references_not_discussed": "Not I am aware of.", "other_strengths_and_weaknesses": "- In section 4.1: \"TOP(·) is used to select the elements with the highest weights from the set with a given probability 1 − k\", what does it mean? Is there still randomness? Can the path selection still get updated during training when using TOP?\n\n- In equation 12, there is superscripts $\\ell_1, \\ell_2, \\ell_3$ for the weights $w_{ij}$. However in equation 9 there is not. What is the shape of $w_{ij}$?\n\n- The proposed sparsity selection introduces additional hyper parameters, including the scheduler step $t$ and the sparsity rate. The sparsity rate may negatively affect the performance if not chosen adequently.", "other_comments_or_suggestions": "- Running title is not formatted.\n\n- Section 3: $C \\in \\mathbb{R}^{n\\times n}$, is the dimension for coefficient correct?", "questions_for_authors": "- Can the choices of retained paths be updated during the second phase of sparsity scheduler?\n\n- How are the sparsity weights initialized?\n\n- In equation 6, what does the superscript 0 of the inner product mean?", "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": "1kZqsAOwdK", "reviewer": "Reviewer_28vQ", "summary": "In this paper, the author proposes a new efficient equivariant operation based on Tensor Product (TP), named Sparse tensor product gate, to improve the efficiency of equivariant networks for Hamiltonian matrix prediction task. From the experiment, the proposed model achieves SOTA performance on QH9 and PubCHemQH, while improving the efficiency about 3-7 times faster than previous method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable comments and suggestions. We address each of the comments individually below.\n\n**Weakness：**\n\nYour assumption regarding the effect of the sparse gate is correct. As demonstrated in the ablation study presented in Table 6 of Appendix B.2, both types of sparse gates significantly improve speed with only a slight loss in accuracy. Therefore, compared to QHNet, the performance improvement in accuracy of our method primarily stems from SPHNet’s architectural design. Given that SPHNet has already achieved strong performance, the minor accuracy loss due to sparsity is entirely acceptable. For more details on the ablation study of different sparse gate, please refer to Appendix B.2 of the paper. Additionally, further experiments on applying the sparse gate to QHNet can be found in our response to the 4th Reviewer EBLX.\n\nBesides, to further address your concerns, we conducted an ablation study to evaluate the effect of different modules in the SPHNet architecture. Specifically, the standard SPHNet model has 4 Vectorial Node Interaction blocks, 2 Spherical Node Interaction blocks, and 2 Pair Construction blocks. We removed all the sparse gates and reduced the number of these three kinds of modules to 1, respectively, and observed the model performance. As shown in the table below, we found that both Vectorial Node Interaction block and Spherical Node Interaction block significantly affect the model performance, indicating the design of architectures with progressively increased irreps orders has an important positive impact on the models. Interestingly, we found that remove one Pair Construction block would not strongly affect the model accuracy, suggesting that there is actually room to further speed up the model. We will explore this further in our future work.\n\n**Table:** The effect of different modules on SPHNet model on PubChemQH dataset.\n\n| Model  | Sparse Pair Gate | Sparse TP Gate | Vectorial Node Interaction block | Spherical Node Interaction block | Pair Construction block | *H* [$10^{-6}E_h$] ↓ |\n|--------|------------------|----------------|----------------------------------|----------------------------------|-------------------------|----------------------|\n| SPHNet | ✗                | ✗              | 4                                | 2                                | 2                       | 86.35                |\n| SPHNet | ✗                | ✗              | 1                                | 2                                | 2                       | 96.01                |\n| SPHNet | ✗                | ✗              | 4                                | 1                                | 2                       | 97.35                |\n| SPHNet | ✗                | ✗              | 4                                | 2                                | 1                       | 89.17                |\n\nHowever, since our model has very different architecture and components from other models, and the specific module in the model is not a one-to-one correspondence, it's hard to make a substitution of the specific module and carry on the ablation study with the other model like QHNet. We are sorry that we are not able to carry on the ablation study with a different model and explain why the SPHNet has such performance improvements compared to others. We'd like to give a try in the future to explore more about this.\n\n**Questions For Authors 1:**\n\nThank you very much for your interest! We are very glad to share our data with the whole community. But our organization require a review process for open sourcing data, which might take some time. We are already actively promoting this data open source process and hope to be able to share our data soon.\n    \n**Questions For Authors 2:**\n\nThank you for your question. We have actually evaluated the relationship between model performance and sparsity ratio in Section 5.4. As shown in Figure 3, for all three datasets, the Hamiltonian MAE remained stable within a certain sparsity range. However, when the sparsity rate reached a particular threshold, we observed a significant increase in the Hamiltonian MAE, which we interpret as the upper limit of sparsity. This suggests that a suitable range of sparsity has little impact on model accuracy while significantly improving computational efficiency. We provide a detailed analysis of this in Section 5.4.", "claims_and_evidence": "The experiments are valid and well support the claim of this paper.", "methods_and_evaluation_criteria": "Sound evaluations. A comprehensive benchmark on existing datasets including QH9 and PubChemQH with reasonable metrics including the MAE on Hamiltonian matrix, MAE on eigen energies, and cosine similarity on the electronic wavefunction. The proposed method greatly improves the model efficiency, which is important for the training and inference of Hamiltonian matrix prediction task.", "theoretical_claims": "N/A.", "experimental_designs_or_analyses": "Sound evaluations and experiments.", "relation_to_broader_scientific_literature": "N/A.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "Strength:\n\n1. The proposed technique is interesting and valid. From the experiments, the sparse adaptive TP can greatly improve the efficiency while the performance will not drop a lot through the ablation study in Appendix B.1. This well supports the main motivation of this paper.\n\n2. The experimental results are strong, achieving the SOTA performance on QH9 and PubChemQH. \n\n3. The writing and organization of this paper is clear. \n\nWeakness:\n1. From the experiments, although the performance improves a lot, the reasons for such improvement have not been well discussed. Since the sparse adaptive TP mainly aims to improve efficiency, I assume the usage of such operation will not bring performance improvement. Therefore, more ablation studies should be included to discuss this.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "Questions:\n\n1. Would you mind sharing the source of the PubChemQH datasets? It seems they are not publicly available yet.\n\n2. I find the time/sec is linearly decreased with the increasing of sparsity rate. How about the relationship of model performance with sparsity ratio?", "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": "lOZMomqCJJ", "reviewer": "Reviewer_EBLX", "summary": "The paper presents a significant advancement in SE(3) equivariant neural networks by introducing a scalable and efficient approach for Hamiltonian prediction. Through innovative sparse gating mechanisms and an adaptive training scheduler, SPHNet achieves remarkable computational savings without sacrificing accuracy.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank all your valuable comments very much, and would like to discuss more on the 4th question of ‘Other Strengths And Weaknesses’: \n\nWe had conducted a series of ablation study to examine the contribution of each module, and the results had been included in the Appendix B.2. As shown in Appendix Table 6 (we also listed the results below), we removed the Sparse Pair gate, the Sparse TP gate, and the Vectorial Node Interaction block from the SPHNet model, respectively, to see their impact on the model efficiency. We found that the two kinds of sparse gates can improve the speedup ratio from 1.73x to 7.09x, demonstrating the effect of sparse gates on the model acceleration. Also, each of the Sparse pair gate and Sparse TP gate alone can provide significant acceleration, improving the speedup ratio from 3.98x to 7.09x and 2.73x to 7.09x. You can find the detailed analysis in Appendix B.2. \n\nAdditionally, we add a similar study by adding the Sparse pair gate and Sparse TP gate to the QHNet model here. Specifically, we applied the Sparse Pair gate on the second Non-diagonal pair block and applied the Sparse TP gate on the Node-wise interaction blocks and both Diagonal and Non-diagonal pair block. As shown in the table below, the sparse gate can effectively accelerate the training speed of the QHNet model and save computational consumption. We will also add this part of the experiment to the Appendix in our future version also.\n\n**Table:** The effect of sparse gates on SPHNet/QHNet model on PubChemQH dataset.\n\n| Model | Sparse Pair Gate | Sparse TP Gate | Vectorial Node Interaction block | Spherical Node Interaction block | *H* [$10^{-6}E_h$] | Memory [GB/Sample] ↓ | Speed [Sample/Sec] ↑ | Speedup Ratio ↑ |\n|-------|--------------------|------------------|-------------------------------------|------------------------------------|-----------------------------|-------------------------|-------------------------|--------------------|\n| SPHNet | ✓ | ✓ | 4 | 2 |  97.31  | 5.62 | 3.12 | 7.09x |\n| QHNet | ✗ | ✗ | 0 | 5 |  123.74  | 22.50 | 0.44 | 1.00x |\n| SPHNet | ✗ | ✓ | 4 | 2 | 94.31  | 8.04 | 1.75 | 3.98x |\n| SPHNet | ✓ | ✗ | 4 | 2 | 87.70  | 6.98 | 1.20 | 2.73x |\n| SPHNet | ✗ | ✗ | 4 | 2 | 86.35  | 10.91 | 0.76 | 1.73x |\n| SPHNet | ✓ | ✓ | 0 | 5 | 97.08  | 8.47 | 1.08 | 2.45x |\n| QHNet | ✗ | ✓ | 0 | 5 | 128.16  | 12.68 | 0.90 | 2.04x |\n| QHNet | ✓ | ✗ | 0 | 5 | 126.27  | 10.07 | 0.73 | 1.66x |\n| QHNet | ✓ | ✓ | 0 | 5 | 128.89  | 8.46 | 1.45 | 3.30x |\n\n\nBesides, we conduct an extra ablation study to evaluate the effect of different modules in the SPHNet architecture. Specifically, the standard SPHNet model has 4 Vectorial Node Interaction blocks, 2 Spherical Node Interaction blocks, and 2 Pair Construction blocks. We removed all the sparse gates and reduced the number of these three kinds of modules to 1, respectively, and observed the model performance. As shown in the table below, we found that both Vectorial Node Interaction block and Spherical Node Interaction block significantly affect the model performance, indicating the design of architectures with progressively increased irreps orders has an important positive impact on the models. Interestingly, we found that remove one Pair Construction block would not strongly affect the model accuracy, suggesting that there is actually room to further speed up the model. We will explore this further in our future work.\n\n\n**Table:** The effect of different modules on SPHNet model on PubChemQH dataset.\n\n| Model  | Sparse Pair Gate | Sparse TP Gate | Vectorial Node Interaction block | Spherical Node Interaction block | Pair Construction block | *H* [$10^{-6}E_h$] ↓ |\n|--------|------------------|----------------|----------------------------------|----------------------------------|-------------------------|----------------------|\n| SPHNet | ✗                | ✗              | 4                                | 2                                | 2                       | 86.35                |\n| SPHNet | ✗                | ✗              | 1                                | 2                                | 2                       | 96.01                |\n| SPHNet | ✗                | ✗              | 4                                | 1                                | 2                       | 97.35                |\n| SPHNet | ✗                | ✗              | 4                                | 2                                | 1                       | 89.17                |", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Experiments follow previous baseline methods.", "relation_to_broader_scientific_literature": "The Sparse TP and Pair Gates could be extended to other SE(3)-equivariant networks, benefiting a wide range of applications in computational chemistry. For example, molecular energy and force field predictions in addition to the Hamiltonian matrix prediction.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "1. The paper addresses a critical bottleneck in SE(3) equivariant graph neural networks—high computational cost due to tensor product (TP) operations.\n\n2. The Sparse Pair Gate filters out unimportant node pairs, reducing computational overhead. The Sparse TP Gate prunes insignificant interactions across different tensor product orders, improving efficiency while maintaining performance.\n\n3. The proposed Three-phase Sparsity Scheduler enables stable training and convergence by progressively optimizing sparse representations, ensuring the balance between efficiency and precision.\n\n4. Experimental results are good. But I'm curious about the effect of each module on the model performance. It would be better if authors can provide more experiments studying the contribution of each module.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: This paper proposes SPHNet, an efficient SE(3) equivariant graph neural network for Hamiltonian matrix prediction. By introducing sparse gating mechanisms and an adaptive sparsity scheduler, SPHNet effectively reduces computational costs while maintaining high accuracy. Experiments on two common datasets demonstrate that SPHNet achieves state-of-the-art performance with significant speed and memory improvements.

The reviewers were all recommending acceptance of the paper.
Decision: Accept (poster)