SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To predict the quantum Hamiltonian matrix, several quantum tensor networks have been proposed [Li et al., 2023, Gong et al., 2023]. SchNOrb [Schütt et al., 2019] uses pairwise distance and direction as the input geometric information to predict the final Hamiltonian matrix. However, SchNOrb lacks the ability to ensure matrix equivariance and relies on data augmentation techniques to encourage equivariance. Another network, DeepH [Li et al., 2022], uses invariant local coordinate systems and a global coordinate system to handle the equivariance challenge. It uses the geometric features within these invariant local coordinate systems to predict invariant Hamiltonian matrix blocks. Next, as a post-processing step, DeepH applies a rotation using Wigner D-Matrix to transform the Hamiltonian matrix blocks from local coordinate systems back to the global coordinate system. Currently, DeepH is applied on predicting Hamiltonian matrices for materials. PhiSNet [Unke et al., 2021] uses an equivariant model architecture that inherently guarantees matrix equivariance. However, current implementation of PhiSNet is limited to supporting single molecule. This limitation arises from the design of the matrix prediction module in PhiSNet, which is designed to predict matrices for the same molecules with fixed matrix size. Therefore, equivariant quantum tensor network QHNet [Yu et al., 2023] is selected as the main baseline method in the QH9 benchmark currently. QHNet has an extendable expansion module that is built upon intermediate full orbital matrices, enabling its capability to effectively handle different molecules. This flexibility allows QHNet to accommodate various molecules in the QH9 benchmark.
3.4 Metrics
To evaluate the quality of the predicted Hamiltonian matrix, we adopt several metrics that are used to measure both approximation accuracy and computational efficiency.
MAE on Hamiltonian matrix H. This metric calculates the Mean Absolute Error (MAE) between the predicted Hamiltonian matrix and the ground-truth labels from DFT calculation. Each Hamiltonian matrix consists of diagonal blocks and non-diagonal blocks, representing the interactions within individual atoms and the interactions between pairs of atoms, respectively. When the atom pair is distant, the values in the Hamiltonian matrix blocks are typically close to zero. Consequently, as the molecules increase in size, the proportion of distant atom pairs also increases, causing the overall mean value of the Hamiltonian matrix to decrease. Hence, in the subsequent experiments, we compare the MAEs of the diagonal and non-diagonal blocks separately as well as the total MAE on the Hamiltonian matrix.
MAE on occupied orbital energies  $\epsilon$ . Orbital energy, which includes the Highest Occupied Molecular Orbital (HOMO) and Lowest Unoccupied Molecular Orbital (LUMO) energies, is a highly significant chemical property. It can be determined by diagonalizing the Hamiltonian matrix using Equation 2. Hence, this metric can serve as a measure to reflect the quality of the predicted Hamiltonian matrix in accurately deducing the desired property. Specifically, it calculates the MAE on all the occupied molecular orbital energies  $\epsilon$  derived from the predicted and the ground-truth Hamiltonian matrix.
Cosine similarity of orbital coefficients  $\psi$ . Electronic wavefunctions can describe the quantum states of molecular systems and are used to derive a range of chemical properties. In order to measure the similarity between the ground-truth wavefunctions and the predicted wavefunctions, we calculate the cosine similarity of the coefficients for the occupied molecular orbitals  $\psi$ . The corresponding coefficients  $\mathbf{C}$  are derived from the predicted and ground-truth Hamiltonian matrix shown in Equation 2.
Acceleration ratio. Besides the metrics assessing molecular properties, several acceleration ratios, such as achieved ratio and error-level ratio, are proposed to measure the quality of the predicted Hamiltonian matrix in accelerating DFT calculation. Specifically, the achieved ratio calculates the ratio of the number of optimization steps between initializing with the predicted Hamiltonian matrix and using initial guess methods like minao and 1e. When the Hamiltonian matrix is accurately predicted, the Self-Consistent Field (SCF) algorithm approaches convergence, resulting in a substantial reduction in the number of optimization steps. As a comparison, the optimal ratio calculates the single optimization step for each molecule divided by the total number of steps, serving as an illustrative benchmark of the ideal performance. Meanwhile, the error-level ratio calculates the ratio over the
Table 2: (Results updated by using the newest QHNet model.) The overall performance on the testing set on the defined four tasks. The unit for the Hamiltonian  $\mathbf{H}$  and eigenenergies  $\epsilon$  is Hartree denoted by  $E_{h}$ .
Dataset	Model	H [10-6Eh] ↓	ε [10-6Eh] ↓	ψ [10-2] ↑
diagonal	non-diagonal	all
QH9-stable-id	QHNet	111.21	73.68	76.31	798.51	95.85
QH9-stable-ood	QHNet	111.72	69.88	72.11	644.17	93.68
QH9-dynamic-geo	QHNet	149.62	92.88	96.85	834.47	94.45
QH9-dynamic-mol	QHNet	416.99	153.68	173.92	9719.58	79.15
number of steps required to reach a the same error level as model prediction during the DFT SCF loop compared to the total number of steps in the DFT process.
Setup. To assess how deep learning approaches perform on the proposed dataset, we conduct experiments on the four designed tasks, as described in Section 3.2. To be more specific, we evaluate the performance of QHNet [Yu et al., 2023], a recently proposed SE(3)-equivariant network specifically designed for efficient and accurate quantum Hamiltonian matrix prediction. QHNet is known for its effectiveness and efficiency in handling the task at hand, making it a suitable testing method for our benchmark evaluation. For quantitative evaluation, we use the metrics as introduced in Section 3.4. Our implementation is based on PyTorch [Paszke et al., 2019], PyG [Fey and Lenssen, 2019], and e3nn [Geiger et al., 2022]. We train models on either (1) a single 48GB Nvidia GeForce RTX A6000 GPU and Intel Xeon Silver 4214R CPU, or (2) a single Nvidia A100 GPU and Intel Xeon Gold 6258R CPU.
Following the model setup in QHNet, in all implemented models, we employ five node-wise interaction layers to aggregate messages from neighboring nodes to update the node irreducible representations. We train all models with a total training step of either 210,000 or 260,000 using a batch size of 32. To expedite the convergence of model training, following the QHNet setup, we implement a learning rate scheduler. The scheduler gradually increases the learning rate from 0 to a maximum value of  $5 \times 10^{-4}$  over the first 1,000 warm-up steps. Subsequently, the scheduler linearly reduces the learning rate, ensuring it reaches  $1 \times 10^{-7}$  by the final step.
Overall performance. We first evaluate the overall performance of the model on the four defined tasks by demonstrating its accuracy of the predicted Hamiltonian matrices on the testing set. As summarized in Table 2, the employed QHNet models can achieve a reasonably low MAE in predicting the Hamiltonian matrices on all proposed tasks. For reference, QHNet can achieve an MAE of  $83.12 \times 10^{-6} E_h$  on the mixed MD17 dataset, which has a similar setup to our QH9-dynamic-geo setup. In addition to MAE on Hamiltonian matrices, the trained models also achieve low errors on the predicted occupied orbital energies and orbital coefficients. This aligns with the prior reported work that QHNet is effective to predict the Hamiltonian matrices for multiple molecules [Yu et al., 2023]. Notably, compared to the existing Hamiltonian matrix datasets, such as MD17 [Chmiela et al., 2017] and mixed MD17 [Yu et al., 2023], our proposed tasks involve predicting Hamiltonian matrices for significantly more molecules. Overall, we anticipate that the proposed new datasets and corresponding tasks can serve as more challenging and realistic testbeds for future research in Hamiltonian matrix prediction.
Investigation on out-of-distribution generalization. Since we maintain a similar number of training samples for QH9-stable-id and QH9-stable-ood, it is feasible to compare the performance of these two settings to investigate the out-of-distribution challenge in predicting Hamiltonian matrices. It is worth noting that we cannot directly compare the performance on their respective test sets, as reported in Table 2, to demonstrate the out-of-distribution generalizability challenge. This is because the molecules in the QH9-stable-ood test set have a larger number of atoms on average than those in QH9-stable-id. As explained in Section 3.4, molecules with larger size typically have more distant atom pairs, thus leading to a lower overall mean value of the Hamiltonian matrix. Hence, numerical results on molecules with different sizes are not directly comparable.
Table 3: The performance of in-distribution (ID) training and out-of-distribution (OOD) training on the constructed evaluation set for the OOD investigation.
Training schema	Models	H [10-6Eh] ↓	ε [10-6Eh] ↓	ψ [10-2] ↑
diagonal	non-diagonal	all
ID	QHNet	84.19	56.01	57.53	442.78	95.26
OOD	QHNet	113.05	70.52	72.78	630.49	94.01
Table 4: The performance of DFT calculation acceleration. Both models, trained on the QH9-stable-id split and the QH9-stable-ood split respectively, are evaluated on a common set of 50 randomly chosen molecules from the intersection of their test sets, making their results directly comparable. Similarly, results from models trained on the QH9-dynamic-geo and QH9-dynamic-mol splits can be compared directly.
Training Dataset	DFT initialization	Metric	Ratio	Training Dataset	DFT initialization	Metric	Ratio
QH9-stable-id	le	Optimal ratio	0.057±0.004	le	Optimal ratio	0.057±0.004
Achieved ratio ↓	0.395±0.030	Achieved ratio ↓	0.400±0.030
Error-level ratio ↑	0.635±0.039	QH9-stable-ood	Error-level ratio ↑	0.620±0.037
Optimal ratio	0.102±0.005	Optimal ratio	0.102±0.005
minao	Achieved ratio ↓	0.706±0.031	minao	Achieved ratio ↓	0.715±0.033
Error-level ratio ↑	0.408±0.025	Error-level ratio ↑	0.406±0.021
QH9-dynamic-geo	le	Optimal ratio	0.056±0.006	le	Optimal ratio	0.056±0.006
Achieved ratio ↓	0.392±0.036	Achieved ratio ↓	0.512±0.138
Error-level ratio ↑	0.648±0.041	QH9-dynamic-mol	Error-level ratio ↑	0.622±0.048
Optimal ratio	0.098±0.008	Optimal ratio	0.098±0.008
minao	Achieved ratio ↓	0.679±0.041	minao	Achieved ratio ↓	0.882±0.217
Error-level ratio ↑	0.443±0.044	Error-level ratio ↑	0.406±0.066
To examine the presence of the out-of-distribution issue in the Hamiltonian prediction task, we adopt an alternative evaluation strategy. To be specific, we assess models that have been trained respectively on the QH9-stable-id and QH9-stable-ood training sets, employing the same set of samples for evaluation in each instance. Specifically, we use the intersecting set of the QH9-stable-id and QH9-stable-ood testing sets as our evaluation dataset. Clearly, the samples contained within this evaluation set are previously unseen during the training phase of either model, thereby maintaining the integrity of the assessment. The evaluation set contains 923 molecules with 23 to 29 atoms. Under this experimental setup, the primary challenge faced by the model trained on the QH9-stable-ood training set stems from the novelty of molecular sizes during the evaluation phase. On the other hand, the model trained on the QH9-stable-id training set benefits from having been exposed to such molecular sizes during training. We denote that these two models are trained under out-of-distribution (OOD) and in-distribution (ID) training schema respectively in Table 3. By comparing the performance on the identical evaluation set, it becomes apparent that the model employing the ID training schema outperforms its OOD-trained counterpart, across all metrics including Hamiltonian MAE and predicted orbital energies and coefficients. Such a performance gap demonstrates that the out-of-distribution issue in molecular size is actually a valid concern particularly when extending trained models to molecular sizes not encountered during training.
Geometry-wise vs. molecule-wise generalization. We further explore geometry-wise and molecule-wise generalizability by analyzing the difficulty differences between the QH9-dynamic-geo and QH9-dynamic-mol tasks. We consider the results in Table 2 for these two tasks to be comparable given that both models are trained with a similar number of geometry structures. We note that the model in the QH9-dynamic-geo task demonstrates numerically better test performance than the model in the QH9-dynamic-mol task. This is consistent with our intention when designing the tasks. Specifically, in QH9-dynamic-geo, although the geometric structures in the test set are different, the molecules themselves are not entirely novel to the model due to the exposure during the training phase. In comparison, the QH9-dynamic-mol task presents a more challenging and demanding scenario. In particular, the test set geometries in QH9-dynamic-mol correspond to entirely different molecules than those seen during training. This task requires the model to generalize from its learned patterns to the unseen molecular structures. To summarize, both tasks serve as valuable testbeds in evaluating the model's generalization ability, and our analysis shows that QH9-dynamic-mol task, which requires extrapolating to entirely new molecular structures, is notably more challenging and demanding.
Accelerating the DFT calculation. We further measure the quality of the predicted Hamiltonian matrix by evaluating its ability in accelerating the DFT calculation. As introduced in Section 3.4, we
compute the ratio of optimization steps required when initializing with the predicted Hamiltonian matrix as compared to classic initial guess methods such as minao and 1e. Note that minao diagonalizes the Fock matrix obtained from a minimal basis to get the guess orbitals, and 1e is one-electron guess which diagonalizes the core Hamiltonian to obtain the guess orbitals. In this experiment, following our data collection process, we use PySCF [Sun et al., 2018] to perform the DFT calculation with using B3LYP exchange-correlation functional and def2SVP basis set. We select DIIS as the SCF algorithm for the DFT calculation and set a grid density level of 3 to ensure an accurate DFT calculation. For each dataset, we compute the average optimization step ratio for 50 randomly selected molecules. As shown in Table 4, we provide several metrics to reflect the optimization ratio. The optimal ratio is the ratio of a single step to the DFT optimization steps. The achieve ratio is the number of optimization steps when initialized by model prediction to optimization steps by DFT initialization. The error-level ratio is the number of optimization steps that achieve similar MAE error with model prediction to the total DFT optimization steps. We can observe that when initializing from the predicted Hamiltonian matrices given by QHNet, it requires fewer optimization steps to reach the converged Hamiltonian matrix, which indicates that the predicted Hamiltonian matrix is close to the convergence condition. This set of experimental results demonstrates that machine learning approaches are helpful in accelerating the DFT calculation.
We are interested in accelerating computation of quantum Hamiltonian matrices, which fundamentally determine the quantum states of physical systems and chemical properties. While various invariant and equivariant deep learning methods have been developed recently, current quantum Hamiltonian datasets consist of Hamiltonian matrices of molecular dynamic trajectories for only a single and four molecules, respectively. To significantly expand the size and variety of such datasets, we generate a much larger dataset based on the QM9 molecules. Our dataset provides precise Hamiltonian matrices for 130,831 stable molecular geometries and 999 molecular dynamics trajectories with sampled 100 geometries in each trajectory. Extensive and carefully designed experiments are conducted to demonstrate the quality of our generated data.
This work was supported in part by National Science Foundation grants IIS-2006861, CCF-1553281, DMR-2119103, DMR-2103842, CMMI-2226908, and IIS-2212419 and by the donors of ACS Petroleum Research Fund under Grant 65502-ND10. Portions of this research were conducted with the advanced computing resources provided by Texas A&M High Performance Research Computing.
Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Efficient and equivariant graph networks for predicting quantum Hamiltonian. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang, Haiyang Yu, YuQing Xie, Xiang Fu, Alex Strasser, Shenglong Xu, Yi Liu, Yuanqi Du, Alexandra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung, Minkai Xu, Chaitanya K. Joshi, Simon V. Mathis, Kamyar Azizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon, Pietro Lio, Rose Yu, Stephan Gunnemann, Jure Leskovec, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423, 2023.
Table 5: Statistics of the datasets in Training/Validation/Testing datasets.
Statistic Type	QH9-stable-id	QH9-stable-ood	QH9-dynamic-geo	QH9-dynamic-mol
# of atoms	Mean	18.02 / 18.02 / 18.03	16.97 / 21.25 / 23.67	16.53 / 16.53 / 16.53	16.56 / 16.38 / 16.38
Median	18 / 18 / 18	17 / 21 / 23	17 / 17 / 17	17 / 17 / 17
Max	29 / 29 / 29	20 / 22 / 29	19 / 19 / 19	19 / 19 / 19
Min	3 / 6 / 4	3 / 21 / 23	10 / 10 / 10	10 / 10 / 10
# of electronics	Mean	65.89 / 65.90 / 65.86	64.98 / 68.83 / 70.57	64.68 / 64.68 / 64.68	64.71 / 64.44 / 64.68
Median	66 / 66 / 66	66 / 70 / 70	66 / 66 / 66	66 / 66 / 66
Max	74 / 74 / 74	74 / 74 / 74	74 / 74 / 74	74 / 70 / 70
Min	10 / 18 / 24	10 / 56 / 58	38 / 38 / 38	40 / 38 / 46
Size of Hamiltonian	Mean	141.59 / 141.62 / 141.56	138.99 / 149.85 / 155.09	137.96 / 137.96 / 137.96	138.04 / 137.34 / 137.88
Median	144 / 144 / 144	142 / 150 / 154	140 / 140 / 140	142 / 142 / 142
Max	166 / 166 / 166	148 / 152 / 166	146 / 146 / 146	146 / 146 / 146
Min	18 / 36 / 48	18 / 126 / 130	82 / 82 / 82	86 / 82 / 92
Table 6: Test performance for the QHNet trained on QH9-stable-id and QH9-stable-ood datasets on the intersection of test sets over various molecule size.
# of Atoms	# of Samples	Training schema	H [10-6Eh] ↓	ε [10-6Eh] ↓	ψ [10-2] ↑
diag	non-diag	all
23	620	ID	86.22	57.75	59.31	435.88	94.83
OOD	107.09	69.87	71.92	583.76	93.47
24	82	ID	92.25	58.53	60.31	353.93	96.92
OOD	125.20	75.29	77.93	365.96	96.74
25	180	ID	76.73	50.99	52.30	505.51	96.34
OOD	118.84	69.56	72.07	1135.47	94.38
26	4	ID	79.75	51.49	52.88	254.94	97.66
OOD	131.86	72.89	75.79	406.69	95.73
27	33	ID	69.33	46.10	47.21	264.35	95.27
OOD	146.63	73.43	76.90	491.40	91.97
28	0	ID	-	-	-	-	-
OOD	-	-	-	-	-
29	4	ID	71.42	45.87	47.00	243.20	94.55
OOD	223.41	90.25	96.15	772.10	90.61
Table 7: Test performance with various training set sizes.
Training sets	# of examples	Models	H [10-6Eh] ↓	ε [10-6Eh] ↓	ψ [10-2] ↑
diag	non-diag	all
QH9-Stable-id	104, 664	QHNet	111.21	73.68	76.31	798.51	95.85
QH9-Stable-id-50k	50,000	QHNet	116.23	74.26	77.20	1127.33	95.37
QH9-Stable-id-10k	10,000	QHNet	127.07	74.93	78.60	2471.46	96.18
A.1 Detailed dataset statistics
Here we provide a detailed statistics including the number of atoms, the number of electronics and the size of Hamiltonian matrix for training, validation and test sets in Table 5 for all the proposed datasets.
A.2 Test performance over molecule size
To evaluate the model performance and alleviate the influence of molecular size, we test the QHNet trained on QH9-stable-id and QH9-stable-ood datasets, and demonstrate the test performance on the molecule with various sizes in Table 6.
A.3 Test performance over various training sets
To investigate the model performance with various number of examples in training sets, we train the QHNet on datasets with various sizes. As shown in Table 7, we can observe that while the MAE on Hamiltonian matrix and cosine similarity on wavefunction  $\psi$  are in similar level with various training sets, the MAE on occupied eigen energies  $\epsilon$  becomes much worse when the training sets become smaller.