SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haiyang Yu*
haiyang@amu.edu
Meng Liu*
mengliu@tamu.edu
Youzhi Luo
yzluo@amu.edu
Alex Strasser
alexstrasser16410@tamu.edu
Xiaofeng Qian
feng@tamu.edu
Xiaoning Qian†
xqian@ece.tamu.edu
Shuiwang Ji†
sji@tamu.edu
Supervised machine learning approaches have been increasingly used in accelerating electronic structure prediction as surrogates of first-principle computational methods, such as density functional theory (DFT). While numerous quantum chemistry datasets focus on chemical properties and atomic forces, the ability to achieve accurate and efficient prediction of the Hamiltonian matrix is highly desired, as it is the most important and fundamental physical quantity that determines the quantum states of physical systems and chemical properties. In this work, we generate a new Quantum Hamiltonian dataset, named as QH9, to provide precise Hamiltonian matrices for 999 molecular dynamics trajectories and 130,831 stable molecular geometries, based on the QM9 dataset. By designing benchmark tasks with various molecules, we show that current machine learning models have the capacity to predict Hamiltonian matrices for arbitrary molecules. Both the QH9 dataset and the baseline models are provided to the community through an open-source benchmark, which can be highly valuable for developing machine learning methods and accelerating molecular and materials design for scientific and technological applications. Our benchmark is publicly available at https://github.com/divelab/AIRS/tree/main/OpenDFT/QHBench.
Machine learning methods have shown great potential in accelerating computations in quantum chemistry tasks [Zhang et al., 2023, Huang et al., 2023, Kirkpatrick et al., 2021]. For example, a variety of invariant geometric deep learning methods have been developed to encode pairwise distances and bond angles in molecular and materials systems [Schütt et al., 2018, Gasteiger et al., 2020, 2021, Liu et al., 2022, Wang et al., 2022, Lin et al., 2023, Yan et al., 2022, Liao et al., 2023,
Equal senior contribution
Figure 1: The target and content of the propose QH9 dataset and benchmark. Quantum tensor networks are built for predicting the target Hamiltonian matrix, facilitating the optimization loop within DFT by providing a precise approximation. Within this benchmark, stable and dynamic datasets are generated for training powerful quantum tensor networks and comprehensive evaluation metrics are proposed to measure the prediction quality.
Schütt et al., 2021, Batzner et al., 2022, Liu et al., 2021] to accelerate the prediction of their chemical properties as data-driven surrogate approximations. To enhance the prediction of vectorial properties, such as force fields, equivariant deep learning methods have been developed to capture permutation, translation, and rotation equivariance for equivariant property prediction [Satorras et al., 2021, Schütt et al., 2021, Thölke and Fabritiis, 2022, Thomas et al., 2018, Batzner et al., 2022, Fuchs et al., 2020, Liao and Smidt, 2023, Anderson et al., 2019, Brandstetter et al., 2022, Batatia et al., 2022]. To support and facilitate the development of machine learning methods on quantum chemistry property prediction, many datasets have been generated to benchmark the respective tasks on molecular property prediction [Blum and Reymond, 2009, Ruddigkeit et al., 2012, Ramakrishnan et al., 2014, Wang et al., 2009, Nakata and Shimazaki, 2017], catalyst prediction [Chanussot et al., 2021, Tran et al., 2023], and force field prediction [Chmiela et al., 2017, 2023].
In addition to these quantum chemistry prediction tasks, the quantum Hamiltonian is another significant and fundamental physical property that determines the quantum states and various materials properties [Marzari and Vanderbilt, 1997, Souza et al., 2001, Qian et al., 2010, Marzari et al., 2012, Bai et al., 2022]. The quantum Hamiltonian can be calculated using Density Functional Theory (DFT) [Hohenberg and Kohn, 1964, Kohn and Sham, 1965] with a time complexity of  $O(n^{3}T)$ , where  $n$  represents the number of electrons and  $T$  denotes the number of optimization steps required to achieve convergence. Given the high computational complexity of the DFT algorithms, accelerating such calculations for novel molecular and materials systems becomes a desirable but challenging task. To tackle this challenge, machine learning methods, such as quantum tensor networks [Li et al., 2022, Gong et al., 2023, Schütt et al., 2019, Yu et al., 2023, Unke et al., 2021], provide a highly promising approach for accelerating the DFT algorithms. These networks directly predict the final Hamiltonian matrix given the input 3D geometries, resulting in significant acceleration of calculations by orders of magnitude.
Unlike invariant chemical properties, Hamiltonian matrices obey intrinsic block-by-block matrix equivariance. This equivariance can be represented by the rotation Wigner D-Matrix, which may contain higher order rotations beyond 3D space. In order to make physically meaningful predictions, it is important to design quantum tensor network architectures that preserve this equivariance property. To perform systematic and in-depth study of this new task, there is a clear need to generate large-scale quantum tensor datasets and benchmarks. Currently, existing quantum Hamiltonian datasets
include the MD17 [Schütt et al., 2019, Gastegger et al., 2020] and mixed MD17 [Yu et al., 2023] datasets, which consist of data for a single and four molecules, respectively. Recent public dataset NablaDFT [Khrabrov et al., 2022] contains Hamiltonian matrices for molecular conformations.
To provide a realistic dataset and comprehensive evaluation, we generate a new quantum tensor dataset named QH9. This dataset contains Hamiltonian matrices for 130,831 stable molecular geometries and 999 molecular dynamic trajectories. In order to provide comprehensive studies for quantum tensor networks, we have designed four specific tasks. The first two tasks, QH9-stable-id and QH9-stable-ood, aim to explore the performance of the networks in both in-distribution and out-of-distribution scenarios, specifically focusing on stable molecular geometries. The QH9-dynamic-geo task follows the setting of the mixed MD17, containing the same molecule with different geometries in the training, validation, and test. On the other hand, the QH9-dynamic-mol task splits the trajectories based on different molecules. Finally, we evaluate the transferability of the trained models on molecules with larger sizes, thereby testing the models' ability to generalize beyond the training dataset. To demonstrate the quality of the predicted Hamilton matrix, we use four metrics. These metrics are based on the Mean Absolute Error (MAE) of the predicted Hamiltonian matrix  $H$ , as well as the derived properties such as orbital energies  $\epsilon$  and electronic wavefunction  $\psi$ . Furthermore, to evaluate the quality of the predicted Hamiltonian in accelerating DFT calculations, we calculate the DFT optimization ratio by taking the model predictions as DFT initialization. The target and content of the proposed QH9 dataset and benchmark are demonstrated in Figure 1.
2.1 Density Functional Theory (DFT)
Modeling the quantum states of physical systems is a central topic in computational quantum physics and chemistry. It aims to solve the Schrödinger equation [Schrödinger, 1926], which describes the electronic states shown as
$$
\hat {H} \Psi \left(\boldsymbol {r} _ {1}, \dots , \boldsymbol {r} _ {n}\right) = E \Psi \left(\boldsymbol {r} _ {1}, \dots , \boldsymbol {r} _ {n}\right), \tag {1}
$$
where  $\Psi (r_1,\dots ,r_n)$  is the  $n$ -electronic wavefunctions and  $\pmb{r}$  is the 3D coordinates. Electronic eigenvalues and wavefunctions play an important role in calculating numerous crucial physical properties, including the Highest Occupied Molecular Orbital (HOMO), the Lowest Unoccupied Molecular Orbital (LUMO), charge density and many others. However, due to the exponentially expanding input Hilbert space with the number of electrons, the computational cost to directly calculate many-electronic wavefunctions is extremely high. Therefore, various methods are proposed to approximate the solutions, such as the Hartree-Fock (HF) method [Szabo and Ostlund, 2012] that approximates the wavefunction itself, or density functional theory [Hohenberg and Kohn, 1964] that approximates the electron density. While the HF method scales with the number of electrons  $n$  as  $O(n^{4}T)$ , DFT scales with  $O(n^{3}T)$  and therefore DFT is better suited for large-scale systems. DFT is based on the key discovery that the total energy and thus all ground-state properties of a system are uniquely determined by the ground-state electron density [Kohn and Sham, 1965].
Both of these approaches divide an  $n$ -electron system into a set of  $n$  non-interacting one-electron wavefunctions  $\psi_{i}(\boldsymbol{r}_{i})$ , also called molecular orbitals in molecular systems. These one-electron orbitals can then be approximated by a linear combination of basis functions  $\phi_{j}(\boldsymbol{r})$  as  $\psi_{i}(\boldsymbol{r}) = \sum_{j} C_{ij} \phi_{j}(\boldsymbol{r})$ . The basis functions can be represented in analytical forms, such as Slater-type orbitals (STOs), Gaussian-type orbitals (GTOs), or plane waves, under numerical approximations for obtaining the coefficients matrix  $\mathbf{C}$ . With these approximations, the original Schrödinger Equation (1) for electrons can be transformed into a matrix form as
$$
\boldsymbol {H} \boldsymbol {C} _ {i} = \boldsymbol {\epsilon} _ {i} \boldsymbol {S} \boldsymbol {C} _ {i}, \tag {2}
$$
where  $\pmb{H}$  is the Hamiltonian matrix,  $S$  is the overlap matrix, and  $\epsilon_{i}$  is the energy for the  $i$ -th orbital. The Hamiltonian matrix can be decomposed into the sum
$$
\boldsymbol {H} = \boldsymbol {H} _ {e N} + \boldsymbol {H} _ {e e} + \boldsymbol {H} _ {X C}, \tag {3}
$$
which describes electron-nucleus interactions  $(\pmb{H}_{eN})$ , electron-electron interactions  $(\pmb{H}_{ee}$ , including kinetic energy and electron-electron Coulomb repulsion energy), and exchange-correlation energy  $(\pmb{H}_{XC})$ . These matrices take the electron density  $\rho(\pmb{r})$  as an input to evaluate the Hamiltonian matrix.
The exchange-correlation energy functional used in this paper was B3LYP [Lee et al., 1988, Becke, 1993], which is a hybrid functional that includes both the exchange energy from the HF method as well as a correlation potential. Thus, the complexity of using B3LYP in DFT is  $O(n^{4}T)$ , which is the same as HF. We implement the GTO basis set Def2SVP [Weigend and Ahlrichs, 2005] that incorporates aspects of DFT, namely, an exchange-correlation potential, in order to more accurately capture electron-electron interactions compared to the HF method, which uses a mean field approximation of electron density.
Equation (2) is satisfied for the final Hamiltonian matrix and its coefficient matrix once self-consistency is achieved using direct inversion in the iterative subspace (DIIS) [Pulay, 1980, 1982]. The equation is solved iteratively by building and solving the Hamiltonian and coefficient matrices, constructing an error vector based on a linear combination of energy differences in the previous steps, then diagonalizing and recalculating  $H$  until the error vector is below a convergence threshold. In our QH9 datasets, we provide the Hamiltonian matrix  $H$  used to train quantum tensor networks for directly predicting the Hamiltonian matrix.
2.2 Group Equivariance and Equivariant Matrices
In many quantum chemistry problems, the molecular property to be predicted (e.g., energy and force) is internally invariant or equivariant to transformations in SE(3) group, including rotations and translations. Formally, for an  $n$ -atom molecule whose 3D atom coordinates are  $\boldsymbol{r}_1,\dots,\boldsymbol{r}_n$ , any transformation in SE(3) group can be described as changing the 3D atom coordinates to  $\boldsymbol{R}\boldsymbol{r}_1 + \boldsymbol{t},\dots,\boldsymbol{R}\boldsymbol{r}_n + \boldsymbol{t}$ . Here, the translation vector  $\boldsymbol{t} \in \mathbb{R}^3$  is an arbitrary 3D vector, and the rotation matrix  $\boldsymbol{R} \in \mathbb{R}^{3\times 3}$  satisfies that  $\boldsymbol{R}^T\boldsymbol{R} = \boldsymbol{I}$ ,  $|\boldsymbol{R}| = 1$ . Let  $\boldsymbol{f}(\cdot)$  map the 3D atom coordinates to an  $(2\ell +1)$ -dimensional prediction target vector, we say  $\boldsymbol{f}$  is order- $\ell$  SE(3)-equivariant if
$$
\boldsymbol {f} \left(\boldsymbol {R} \boldsymbol {r} _ {1} + \boldsymbol {t}, \dots , \boldsymbol {R} \boldsymbol {r} _ {n} + \boldsymbol {t}\right) = D ^ {\ell} (\boldsymbol {R}) \boldsymbol {f} \left(\boldsymbol {r} _ {1}, \dots , \boldsymbol {r} _ {n}\right) \tag {4}
$$
holds for any rotation matrix  $\pmb{R}$  and translation vector  $\pmb{t}$ , where  $D^{\ell}(\pmb{R}) \in \mathbb{C}^{(2\ell + 1) \times (2\ell + 1)}$  is the order- $\ell$  Wigner-D matrix of  $\pmb{R}$  (please refer to Section A.3 of Brandstetter et al. [2022] and Section A.2.1 of Poulenard et al. [2022] for more background information about the Wigner-D matrix). To accurately predict SE(3)-equivariant properties, an effective approach is to develop neural network models that are designed to maintain the same equivariance relations between inputs and outputs as in Equation (4). Recently, many studies have proposed SE(3)-equivariant neural network architectures by using SE(3)-invariant feature encoding [Schütt et al., 2018, Gasteiger et al., 2020, 2021, Liu et al., 2022], tensor product operations [Thomas et al., 2018, Brandstetter et al., 2022, Liao and Smidt, 2023], or atomic cluster expansion framework [Batatia et al., 2022, Drautz, 2019, Dusson et al., 2022, Kovács et al., 2021, Musaelian et al., 2023].
Different from vector-like molecular properties, the Hamiltonian matrix  $\pmb{H}$  has a much more complicated SE(3) equivariance pattern that is associated with the intrinsic angular momentum of the atomic orbital pairs. In computational quantum chemistry algorithms such as DFT, the Hamiltonian matrix  $\pmb{H}$  obtained from DFT calculations can be used to represent the interactions between these atomic orbitals, and the block  $H_{ij}$  in Hamiltonian matrix represents the interactions between the atomic orbitals  $i$  in atom  $a_i$  with angular quantum number  $\ell_i$  and atomic orbitals  $j$  in atom  $a_j$  with angular quantum number  $\ell_j$ , and the shape of this block  $H_{ij}$  is  $(2\ell_i + 1) \times (2\ell_j + 1)$ . Usually, the atomic orbitals are arranged sequentially for the orbitals in the same atom and with the same angular quantum number. For example,  $H_{ij}$  can be located within the  $s_i$ -th to  $(s_i + 2\ell_i)$ -th row, and the  $s_j$ -th to  $(s_j + 2\ell_j)$ -th column of Hamiltonian matrix  $\pmb{H}$ . Specifically, its SE(3) equivariance can be described as
$$
\boldsymbol {H} _ {i j} (\boldsymbol {\rho} (\boldsymbol {R} \boldsymbol {r} + \boldsymbol {t})) = D ^ {\ell_ {i}} (\boldsymbol {R}) \boldsymbol {H} _ {i j} (\boldsymbol {\rho} (\boldsymbol {r})) D ^ {\ell_ {j}} (\boldsymbol {R} ^ {T}), \tag {5}
$$
where  $\rho(r)$  is the electronic density at position in  $r$  and Hamiltonian matrix  $\pmb{H}$  is a function of the electronic density  $\rho(r)$  in the DFT algorithm. In other words, the SE(3) equivariance of different submatrices in  $\pmb{H}$  has different mathematical forms, which is much more complicated than the SE(3) equivariance of vector-like molecular properties. Hence, it is much more challenging to develop SE(3)-equivariant neural network architectures for the prediction of Hamiltonian matrices. Nowadays, only a few studies [Li et al., 2022, Gong et al., 2023, Yu et al., 2023, Unke et al., 2021] have made initial exploration in this direction.
2.3 Datasets for Quantum Chemistry
To facilitate the usage of machine learning models to predict chemistry properties and accelerate simulations, numerous quantum chemistry datasets have been built to provide extensive and faithful data. Here, we introduce several existing datasets that have been constructed for different tasks respectively, including molecular property prediction, catalyst modeling, molecular force field prediction, and molecular Hamiltonian matrix prediction. For molecular property prediction, the QM7 [Blum and Reymond, 2009] dataset was initially constructed using 7,165 molecules from the organic molecule database GDB-13 [Blum and Reymond, 2009], with each selected molecule having no more than 7 heavy atoms. The primary purpose of creating the QM7 dataset is to provide atomization energies as the target molecular property. Then QM9 [Ramakrishnan et al., 2014, Schutt et al., 2018] was built based on GDB-17 [Ruddigkeit et al., 2012] to provide 134k stable small organic molecules with no more than 9 heavy atoms in each molecule. Moreover, it provides 13 different important quantum chemistry properties, including HOMO and LUMO energies. Based on the molecules from PubQChem [Wang et al., 2009, 2017, Kim et al., 2019, 2021, 2023], PubQChemQC [Nakata and Shimazaki, 2017] provides 3M ground-state molecular structures as well as the HOMO-LUMO gap and excitation energies for 2M molecules. In addition to the molecular property datasets, OC20 [Chanussot et al., 2021] and OC22 [Tran et al., 2023] were developed to provide the data of interactions of catalysts on material surfaces. They provide the geometries of the initial structures to predict the final structures or energies as well as the relaxation trajectories with energy and atomic forces. For the molecular force field prediction datasets, MD17 [Chmiela et al., 2017] and MD22 [Chmiela et al., 2023] contain atomic forces for molecular and supramolecular trajectories respectively as valuable datasets to develop machine learning methods. The last category is the Hamiltonian matrices datasets. MD17 [Schutt et al., 2019, Gastegger et al., 2020] provides the Hamiltonian matrices for single molecular dynamic trajectories to study the Hamiltonian matrices for molecules with various geometries. Building upon this dataset, mixed MD17 [Yu et al., 2023] combines four molecular trajectories in the MD17 to study Hamiltonian matrix prediction tasks with multiple molecules. NablaDFT [Khrabrov et al., 2022] has million of Hamiltonian matrices for molecular conformers and provides the MAE on predicted Hamiltonian matrices with model trained with in-distribution data split. Alongside the increasing interest in Hamiltonian matrix prediction, there is a growing need for datasets that include Hamiltonian matrices with a greater number of molecules and benchmark with comprehensive evaluation to facilitate the subsequent studies.
3 Datasets, Tasks, Methods, and Metrics
Dataset Generation. For the QH9 dataset, we use open-source software PySCF [Sun et al., 2018, 2020] to conduct computational quantum chemistry calculations. In the QH9, there are two sub datasets. The first one is the QH9-stable dataset containing Hamiltonian matrices for 130,831 molecules with their geometries. We obtain the molecules and their geometries in QH9-stable from the QM9 version with 130,831, which is widely used in molecular property prediction tasks in the literature [Schutt et al., 2018, Gasteiger et al., 2020, Liu et al., 2022, Wang et al., 2022]. Note that we cover all the molecules in this QM9 version. The second one is the QH9-dynamic dataset, it has molecular trajectories for 999 molecules and each trajectory contains 100 geometries. To obtain the accurate Hamiltonian matrices for this dataset, we set the hyper-parameters of the DFT algorithms to a tight level. Specifically, we set the grid density level to 3 to calculate accurate electronic density, and the SCF convergence condition is set to SCF tolerance of  $10^{-13}$  and gradient threshold of  $3.16 \times 10^{-5}$  to ensure the final states achieve tight convergence. For the density functional, we select the B3LYP exchange-correlation functional to conduct DFT calculations, and GTO orbital basis Def2SVP is selected to approximate the electronic wavefunctions. To accelerate and achieve the convergence of SCF algorithm, we use DIIS algorithm with consideration of the 8 previous steps. For the QH9-dynamic dataset, molecular dynamics simulations are conducted under the microcanonical ensemble, where the number of particles, volume, and energy remain constant (NVE). The temperature is set to 300K, and the time step for recording the molecular trajectory is set to 0.12 fs with 1,000 total steps. In the dataset collection, we sampled 100 time steps for each trajectory with a timestep of 1.2 fs. The generated QH9 dataset is available at Zenodo (https://zenodo.org/records/8274793) and GitHub (https://github.com/divelab/AIRS/tree/main/OpenDFT/QHBench/QH9).
(a) The histogram of the molecule size for QH9-stable dataset.
(b) Number of heavy atom percentage in QH9-stable dataset.
(c) The histogram of the molecule size for QH9-dynamic.
Figure 2: The dataset statistics on QH9-stable and QH9-dynamic, including molecule size distribution and percentage of molecules with different number of heavy atoms.
Table 1: The statistics of our defined four tasks.
Task	# Total geometries	# Molecules	# Training/Validation/testing geometries
QH9-stable-id	130,831	130,831	104,664/13,083/13,084
QH9-stable-ood	130,831	130,831	104,001/17,495/9,335
QH9-dynamic-geo	99,900	999	79,920/9,990/9,990
QH9-dynamic-mol	99,900	999	79,900/9,900/10,100
Dataset Statistics. The statistical data, including the number of molecules and geometries for QH9-stable and QH9-dynamic, is presented in Table 1. These molecules consist of no more than 9 heavy atoms and are composed of four specific heavy atoms: carbon (C), nitrogen (N), oxygen (O), and fluorine (F). The distribution of molecule size for QH9-stable and QH9-dynamic is shown in Figure 2a and Figure 2c. Meanwhile, the percentage of molecules in QH9-stable with different the number of heavy atoms is shown in 2b.
To comprehensively evaluate the quantum Hamiltonian prediction performance, we define the following tasks based on the obtained stable and dynamic geometries in the QH9 dataset.
QH9-stable-id. We first randomly divide the obtained stable geometries in QH9 into three subsets, including  $80\%$  for training,  $10\%$  for validation, and  $10\%$  for testing. This serves as the basic evaluation task for predicting quantum Hamiltonian matrices.
QH9-stable-ood. We further split the stable geometries in QH9 by molecular size based on the number of constituting atoms. The training set consists of molecules with 3 to 20 atoms, maintaining a similar number of training samples as in the QH9-stable-id split. The validation set includes molecules with 21 to 22 atoms, while the testing set has molecules with 23 to 29 atoms. This task allows for an evaluation of the model's generalization ability under an out-of-distribution training setup.
QH9-dynamic-geo. For this split and the following QH9-dynamic-mol split, there are 999 molecular dynamics trajectories, while each trajectory includes 100 geometries. In QH9-dynamic-geo, the split is performed geometry-wise. Specifically, for each molecule, 100 geometries are randomly divided into 80 for training, 10 for validation, and 10 for testing. Here, the molecules in the test set are visible during training but the geometric structures are different from training structures.
QH9-dynamic-mol. In this QH9-dynamic-mol split, the 999 molecules are divided into training, validation, and testing subsets in a ratio of  $0.8 / 0.1 / 0.1$ . Importantly, different from the above QH9-dynamic-geo setup, all 100 geometries corresponding to a specific molecule are grouped together and assigned to the same subset. This setup introduces a more challenging task than QH9-dynamic-geo since the geometries in the testing set correspond to different molecules as those in training.
3.3 Methods