SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Generating QM1B with  $\mathbf{PySCF}_{\mathbf{IPU}}$
Alexander Mathiasen<sup>1</sup>  Hatem Helal<sup>1</sup>  Kerstin Klaser<sup>1</sup>  Paul Balanca<sup>1</sup>  Josef Dean<sup>1</sup>  Carlo Luschi<sup>1</sup>  Dominique Beaini<sup>2,3,4</sup>  Andrew Fitzgibbon<sup>1</sup>  Dominic Masters<sup>1</sup>
<sup>1</sup>Graphcore <sup>2</sup>Mila - Québec AI Institute <sup>3</sup>Valence <sup>4</sup>Université de Montréal
The emergence of foundation models in Computer Vision and Natural Language Processing have resulted in immense progress on downstream tasks. This progress was enabled by datasets with billions of training examples. Similar benefits are yet to be unlocked for quantum chemistry, where the potential of deep learning is constrained by comparatively small datasets with 100k to 20M training examples. These datasets are limited in size because the labels are computed using the accurate (but computationally demanding) predictions of Density Functional Theory (DFT). Notably, prior DFT datasets were created using CPU supercomputers without leveraging hardware acceleration. In this paper, we take a first step towards utilising hardware accelerators by introducing the data generator  $\mathrm{PySCF_{IPU}}$  using Intelligence Processing Units (IPUs). This allowed us to create the dataset QM1B with one billion training examples containing 9-11 heavy atoms. We demonstrate that a simple baseline neural network (SchNet 9M) improves its performance by simply increasing the amount of training data without additional inductive biases. To encourage future researchers to use QM1B responsibly, we highlight several limitations of QM1B and emphasise the low-resolution of our DFT options, which also serves as motivation for even larger, more accurate datasets. Code and dataset.
Artificial Intelligence research is quickly moving towards building systems that are able to perform generally applicable functions using foundation models that are pretrained on billions of training examples. Foundation models have transformed natural language processing (NLP) [1] and computer vision (CV) [2], but have not yet been demonstrated in the important field of molecular machine learning. Addressing this has the potential to accelerate the design of new medicines and materials.
A promising approach within molecular machine learning is to train neural networks (NN) to approximate the predictions of quantum chemistry, yielding 1000x faster predictions [3, 4] with errors approaching that of experimental measurement. However, current molecular datasets are limited to 100k-20M training examples, that are seemingly insufficient to train foundation models. Due to the high cost and time consuming nature of constructing chemical datasets directly from experimental measurement, prior datasets rely on computational methods to generate labels. The computational method routinely used to generate machine learning datasets is Density Functional Theory (DFT). Although DFT strikes a favourable balance between accuracy and computational cost relative to other molecular simulation methods, its computational cost has nevertheless limited prior  $\text{dists}^1$  to fewer than 20M training examples. Notably, prior DFT datasets such as QM9 [6, 7], ANI-1 [8] or PCQ [9] were created by running DFT on CPUs. To the best of our knowledge no prior work utilises hardware acceleration to create larger DFT datasets for deep learning.
$^{1}$ Gaussian basis DFT datasets. This excludes semi-empirical methods [5] and plane-wave DFT [4].
Figure 1: We took  $1.09\mathrm{M}$  molecules with  $\geq 9$  atoms from GDB11 [11, 12], and used  $\mathrm{PySCF_{IPU}}$  to compute molecular properties (e.g. HOMO-LUMO (HL) gap) for a set of up to 1000 conformers. Three such sets are visualised here, corresponding to low, modal, and high variance of HL gap. Training a SchNet 9M model to predict HL gap shows improvement as the number of training samples approaches 500M. All training runs see  $10^{9}$  examples, sampling with replacement. We conjecture that billions of samples will enable the training of molecular foundation models.
We introduce  $\mathrm{PySCF_{IPU}}$ , a DFT data generator which utilises Intelligence Processing Units (IPUs [10]) to accelerate molecular property dataset generation. We used  $\mathrm{PySCF_{IPU}}$  to generate QM1B with one billion DFT training examples within 40000 IPU hours, substantially less than the two years it took to create PCQ on CPU supercomputers. The larger scale was enabled by several changes in the dataset generation which we outline in Section 4. The most important change, from a machine learning perspective, is that we increased the number of training examples by reducing the DFT accuracy. We advise that QM1B is only used for research purposes and discourage applications in deployed systems without carefully investigating the consequences of the reduced DFT accuracy. It remains to be found how pretraining on 1B low-resolution DFT training examples may bias a NN which is subsequently fine-tuned on downstream tasks. We encourage future research to investigate how DFT options bias subsequent neural networks by jointly iterating on dataset creation and model training. To enable such research we do not only release QM1B but also the entire software stack used to generate QM1B.
To investigate whether more training data improves neural networks, we trained a baseline SchNet with 9M parameters on differently sized subsets of QM1B. We found the validation Mean Absolute Error (MAE) improved from  $285\mathrm{meV}$  to  $32\mathrm{meV}$ , see Figure 1. Notably, the training and validation MAE overlap at the larger scales, indicative of underfitting, paving the way for future work to train large foundation molecular neural networks.
After pre-training SchNet on QM1B we fine-tuned SchNet on QM9. Compared to training SchNet on QM9 from scratch, we found this to improve validation MAE from  $54.13\mathrm{meV}$  to  $30.2\mathrm{meV}$ .
The contributions of this work can be summarised as follows:
1. We present  $\mathrm{PySCF_{IPU}}$ , a new hardware accelerated DFT data generator for deep learning.
2. We release the dataset QM1B, a low-resolution DFT dataset with  $10^{9}$  training examples.
3. We train a simple NN baseline which improves performance when trained on more data, validating the scale of QM1B.
The core contributions of this work touch on three mature fields that have typically been studied in isolation. We provide a brief history of DFT software packages and follow this with their application to generating large quantum chemical datasets. Finally, we summarise efforts to approximate DFT with machine learning models.
2.1 DFT Libraries
A suitable approximation to quantum mechanics can provide a universal framework for predicting properties of molecules and materials [13]. The most widely applied approximation to quantum theory is Kohn-Sham DFT [14] which is implemented in a number of widely available DFT packages. These packages have evolved over decades and have successfully adapted to the constantly shifting landscape of high-performance computing. Notably, a substantial fraction of nation-level supercomputers are
dedicated to running DFT simulations, e.g.,  $30\%$  of the US National Energy Research Scientific Computing Center (NERSC) and up to  $40\%$  of the UK ARCHER2 supercomputer [15, 16]. We highlight three advances in scientific computing that have accelerated progress in machine learning research that have yet to be widely adopted by DFT software libraries:
1. Ubiquitous hardware acceleration using low-precision numerical formats.
2. Tensor computations in a high-level API like NumPy [17].
3. Automatic differentiation [18].
Crucially, the core algorithm development of many widely used DFT packages predates the rise of NumPy and the Python ecosystem. This also coincides with the increased availability of specialised hardware accelerators for scientific computing. For this reason, most DFT libraries are still developed in languages like FORTRAN and C, without support for automatic differentiation. Hardware acceleration in this low-level programming environment is limited, with marginal performance gains reported [19]. One notable exception is the commercial TeraChem library [20] which was engineered specifically for GPUs.
Table 1: Comparison of different Gaussian basis set DFT libraries frequently deployed on high performance computing clusters<sup>2</sup>. Despite the fact that there are many DFT libraries that support hardware acceleration, we know of none that have been used to generate machine learning datasets.
Library	Autodiff	Hardware	Precision	License	Language
Gaussian9 [21]	no	CPU/GPU	float64	Commercial	-
TeraChem [20]	no	CPU/GPU	mixed float32/64	Commercial	C++
Psi4 [22]	no	CPU/GPU	float64	Open Source	C++
GAMESS [23]	no	CPU/GPU	float64	Open Source	C, Fortran77
PySCF [24]	no	CPU	float64	Open Source	NumPy, C
PySCFIPU (ours)	yes*	IPU	float32	Open Source	Jax, C
2.2 DFT Datasets
The authors of previous DFT datasets made different choices, e.g., basis set and number of conformers. We list several examples in Table 2 limited to Gaussian basis DFT. It is not clear how these choices impact subsequent deep learning models. We hope  $\mathrm{PySCF_{IPU}}$  will enable future researchers to jointly iterate on dataset creation and model training. While prior datasets use DFT libraries with some support for hardware acceleration, we found all dataset papers relied on CPUs.
Table 2: Comparison of different DFT datasets.
Dataset	Graphs	Conformers	Basis Set	XC	Atoms	Time	Library	Hardware
QM9 [7]	133.9k	133.9k	6-31G(2df,p)	B3LYP	≤ 9	-	?	CPU
PCQ [9]	3.38M	3.38M	6-31G*	B3LYP	≤ 20	~2 years [25]	GAMESS	CPU
SPICE [26]	19238	1.13M	def2-TZVPPD	ωB97M-D3(BJ)	2-96	-	PSI4	CPU
ANI-1 [27]	57462	20M	6-31G(d)	ωB97x	≤ 8	-	Gaussian9	CPU
QM1B	1.09M	1B	STO-3G	B3LYP	9-11	~5 days4	PySCFIPU	IPU
2.3 Neural Net Approximations to DFT
Similar to CV and NLP, molecular machine learning has seen substantial effort towards optimising neural networks to better approximate DFT. Broadly speaking, these approaches include denoising and auxiliary tasks [28, 29, 30, 31, 32], pretraining/finetuning [30, 33, 34, 35, 36], positional encodings [37, 38, 39, 40, 41] and architectural improvements [42, 43, 44]. Some work concentrated on performing equilibrium based predictions with 2D molecular representations bypassing expensive DFT-based structure optimisations [45, 46, 47, 48, 49]. [4] investigated how NNs improve their DFT
2This excludes e.g. the unpublished github repositorygpu4pyscf.
This excludes plane-wave DFT as used in [4] and semi-empirical methods as used in [25].
Estimated as 40000 IPU hours using 320 IPUs.
approximation while increasing training examples. They concluded that current methods would require orders of magnitude more training examples to solve their catalysis task. For demonstration purposes, we chose to scale one of the simplest models, SchNet [50], to investigate how far performance can be improved by simply scaling training data without incorporating additional architectural innovations. In Table 3 we compare SchNet trained on QM9 against a selection of other methods that were trained from scratch (i.e. no pretraining/finetuning). We point out that the parameter count is orders of magnitude lower than foundation models in CV and NLP. Furthermore, the machine learning errors around  $20 - 70\mathrm{meV}$  are orders of magnitude larger than the convergence thresholds commonly employed in DFT. This motivated us to challenge the common assumption that DFT requires double precision (float64) for the special case of generating training data for machine learning.
Table 3: Results of previous Neural Networks trained from scratch on the DFT dataset QM9.
Method	# Params	MAE QM9 HL gap (meV)	MAE QM9 energy (meV)
SchNet [50]	432k	63	14.6
PaiNN [51]	589k	45.7	5.85
DimeNet++ [52]	1.8M	32.6	6.32
GNS + Noisy Nodes [31]	-	28.6	7.3
3 PySCFIPU
The Python-based Simulations of Chemistry Framework (PySCF) implements multiple Self-Consistent Field methods (SCF) including DFT. Importantly, relative to many other DFT libraries, PySCF is free to use, open-source and implemented mainly in Python and NumPy with a few of the computational demanding operations written in C. Furthermore, PySCF supports a plethora of DFT options, enabling it to reproduce quantum chemical properties as computed in machine learning datasets like QM9 or PCQ. All of the above made PySCF a good starting point to port for IPUs.
Intelligence Processing Units (IPUs) are a type of accelerated hardware designed specifically for machine learning workloads. We found IPUs allowed us to speed up DFT data generation due to two reasons:
1. IPUs have 940MB on-chip memory with 12-65TB/s bandwidth, enough to perform small DFT computations without relying on off-chip RAM with  $< 3\mathrm{TB / s}$  bandwidth.
2. IPUs support Multiple Instruction Multiple Data (MIMD) parallelism which simplifies parallel computation of the notoriously difficult computation of Electron Repulsion Integrals.
3.1 DFT Primer.
Our data generator  $\mathrm{PySCF_{IPU}}$  implements DFT. To allow machine learning researchers to reason about our data generator, this section describes what Gaussian basis set DFT implementations like  $\mathrm{PySCF_{IPU}}$  do. We deliberately limit ourselves to explaining what DFT do and exclude explaining the underlying physics, e.g., under which assumptions DFT can be derived from the Schrödinger equation (see instead [53] or [54, Appendix B]).
The inputs to a DFT computation are atomic positions and atomic numbers. These inputs are used to compute a matrix  $\mathbf{V}$  which represents the system of interest. One then solves the following matrix equation for  $(\mathbf{C},\epsilon)$ .
$$
\left[ \mathbf {V} + \mathbf {T} + J (\rho (\mathbf {C})) + K (\rho (\mathbf {C})) + V _ {X C} (\rho (\mathbf {C})) \right] \mathbf {C} = \mathbf {S} \mathbf {C} \boldsymbol {\epsilon}
$$
$$
\rho (\mathbf {C}) = \mathbf {C D C} ^ {T} = \boldsymbol {\rho} \quad D _ {i j} = \left\{ \begin{array}{l l} 2 & \text {i f} i = j \leq N _ {\text {e l e c t r o n s / 2}} \\ 0 & \text {e l s e} \end{array} \right.
$$
$$
J (\pmb {\rho}) _ {k l} = \sum_ {i j} I _ {i j k l} ^ {2 e} \rho_ {i j} K (\pmb {\rho}) _ {i l} = - \frac {1}{2} \sum_ {j k} I _ {i j k l} ^ {2 e} \rho_ {j k}
$$
Here  $\mathbf{I}^{2e}\in \mathbb{R}^{N\times N\times N\times N}$  while  $\{\pmb {\rho},\mathbf{V},\mathbf{T},\mathbf{C},\mathbf{S},\pmb {\epsilon}\}$  are in  $\mathbb{R}^{N\times N}$ ,  $(\mathbf{D},\epsilon)$  are diagonal matrices and  $\{\rho ,J,K,V_{XC}\}$  are functions with  $(N,N)$  matrices as input and output. Given a physical system,
Under the Apache-2.0 License under which we plan to open-source  $\mathrm{PySCF_{IPU}}$ .
the DFT algorithm proceeds in iterations. At iteration 0, we precompute  $(\mathbf{V}, \mathbf{T}, \mathbf{S}, \mathbf{I}^{2e})$  and initialise  $\mathbf{C}_1$ . At iteration  $i$ , we compute  $(\pmb{\rho}_i = \rho(\mathbf{C}_i), V_{XC}(\pmb{\rho}_i), J(\pmb{\rho}_i), K(\pmb{\rho}_i))$ , and then solve the above equation for  $\mathbf{C}_{i+1}$  (which corresponds to a generalised eigenproblem). If the iterations converge to self-consistency, one attains a solution  $(\mathbf{C}, \epsilon)$  from which several molecular properties can be computed, e.g.
$$
\mathrm {H L} \quad \text {g a p} := \epsilon_ {\text {H O M O}} - \epsilon_ {\text {L U M O}}
$$
DFT has two well-known trade-offs between accuracy and compute time:
1.  $V_{XC}$  attempts to correct the approximation error between DFT and the Schrödinger equation; hundreds of corrections exist such as LDA [13], B3LYP [55, 56] and  $\omega \mathrm{B97x}$  [57].
2. C represents molecular orbitals as a linear combination of Gaussian functions; many such Gaussian basis sets have been extensively studied [58].
Prior work trained neural networks to improve the inherent errors from DFT for  $V_{XC}$  [59, 60] and the representation of molecular orbitals [54].
3.2 PySCF IPU Implementation
To utilise IPUs we ported PySCF from NumPy to JAX [61], which can target IPUs by using IPU TensorFlow XLA backend [62]. This left two remaining parts in C:
1. LIBXC [64] which computes  $V_{XC}$
2. LIBCINT [65] which computes  $(\mathbf{T},\mathbf{V},\mathbf{S},\mathbf{I}^{2e})$
Libxc. We only support the B3LYP functional. The energy computation of B3LYP was implemented in JAX, which allowed us to use JAX autograd to compute the B3LYP potential. We use float32 instead of float64. JAX-XC [66] recently machine translated the libxc Maple files to JAX. This may allow us to support more functionals by extending JAX-XC to float32.
Libcint. We implemented the computation of  $\mathbf{I}^{2e}$  by porting the libcint implementation of the Rys Quadrature algorithm from C to the IPU. The code implements a function INT2E which is called many times with different inputs. Each call can be done in parallel. However, the different calls perform different computations making it tricky to parallelise using Single Instruction Multiple Data (SIMD). In contrast, using the IPUs MIMD parallelism all 8832 IPU threads can independently compute one call to INT2E in parallel.  $\mathbf{I}^{2e}$  satisfies an 8x symmetry which, when exploited, (1) reduces bytes needed to store  $\mathbf{I}^{2e}$  by 8x and (2) gives an 8x reduction in FLOPs needed to compute  $K(\pmb {\rho},\mathbf{I}^{2e})$  and  $J(\pmb {\rho},\mathbf{I}^{2e})$  (usually computed using np.einsum). To utilise the 8x symmetry we implemented a custom einsum algorithm. The memory usage of our current implementation can be improved as outlined in Section 5.
Numerical Error of  $\mathrm{PySCF_{IPU}}$ . DFT libraries usually use double precision (float64), see Table 1. In contrast, NNs are usually trained in float32 to decrease compute time, e.g., some hardware accelerators can perform matrix multiplications more than  $10\mathrm{x}$  faster in float32 than float64. For the particular case of libcint, we found float32 to be around 6x faster than float64, likely due to IPUs ability to hardware accelerate float32 instead of simulating float64 in software (similar to the 10x reported by [67]). While float32 has yet to be thoroughly demonstrated in DFT libraries, our goal is to generate data to train NNs, therefore, our requirements for numerical precision is less stringent compared to mainstream DFT libraries, our errors just have to be below that of NNs (20-70meV for HL gap). To investigate our numerical errors, we recomputed DFT with PySCF in float64 for 10k SMILES strings from QM1B. This allowed us to compute a the mean absolute error (MAE) of  $\mathrm{PySCF_{IPU}}$  relative to PySCF. For converged<sup>8</sup> molecules our MAE was 6meV for energies and 0.2meV for HL gap, see Figure 2. Our numerical error for HL gap is thus 100x smaller than the errors achieved by neural networks (but similar for energy), see Table 3.
[63] cleverly extends PySCF with AD by using JAX to pair all C calls with gradient C calls (only for CPUs).
$^{7}\mathbf{I}_{ijkl}^{2e} = \mathbf{I}_{ijlk}^{2e} = \mathbf{I}_{jilk}^{2e} = \mathbf{I}_{jilk}^{2e} = \mathbf{I}_{klij}^{2e} = \mathbf{I}_{klj}^{2e} = \mathbf{I}_{lki j}^{2e} = \mathbf{I}_{lkji}^{2e}$
${}^{8}$  Converged is defined as np.std(energies[-5:])<0.01 . Around 99% converged. In the future the remaining 1% could be added using PySCF for a negligible cost.
Figure 2: Numerical error of  $\mathrm{PySCF_{IPU}}$  in float32 compared to PySCF in float64. We suspect the difference between HL gap and energy is caused by energies being  $100\mathrm{x}$  larger than HL gaps.
Timing  $\mathbf{PySCF}_{\mathbf{IPU}}$ . Our logs recorded 38729 IPU hours generating QM1B. We estimate less than  $4\%$  was spent compiling. We completed the data generation within 5 days, by utilizing a varying number of 256-384 IPUs split over 16-24 POD16s. Each IPU POD16 has two physical EPYC CPUs with 240 vCPUs. On a POD16, for the example molecule FC=C1C2CN2N=C1C=O, we can do 82 DFTs/sec with PySCF and 228 DFTs/sec with  $\mathbf{PySCF}_{\mathbf{IPU}}$ , see Figure 3 for a profile of  $\mathbf{PySCF}_{\mathbf{IPU}}$ .
Figure 3: Profile of our DFT computation for "Fc1nocnc  $(= 0)$  c1  $= 0$
3.3 Limitations
$\mathrm{PySCF}_{\mathrm{IPU}}$  is an ongoing research project and thus has several limitations:
1. The 940MB limits  $\mathrm{PySCF_{IPU}}$  to  $\leq 12$  heavy atoms with a less accurate electron representation (STO-3G).
2. Larger numerical errors due to float32 instead of float64.
3. The first simulation of a molecule with  $N$  atomic orbitals takes approximately five additional minutes due to the ahead-of-time compilation model used for IPUs (for QM1B we estimate spending less than  $4\%$  time compiling).
Our current implementation supports restricted Kohn-Sham, B3LYP functional [55, 56], no interatomic force and the basis sets  $\{\mathrm{STO3G}, 6-31\mathrm{G}\}$ .
4 QM1B
We created QM1B by taking 1.09M SMILES strings from the Generated Data Bank (GDB)<sup>9</sup> with 9-11 heavy atoms (GDB11). Hydrogen atoms were added by using RDKIT resulting in 0 to 11 Hydrogen atoms. We subsequently generated up to 1000 conformers for each molecule using the ETKDG algorithm [68] as implemented in RDKit [69]. This led to a total of 305.8M, 568,7M and 205.4M conformers for 9,10,11 heavy atoms respectively. We then computed the following properties on the resulting 1B conformers using  $\mathrm{PySCF_{IPU}}$ :
1. Energy
2. Highest Occupied Molecular Orbital (HOMO) Energy
3. Lowest Occupied Molecular Orbital (LUMO) Energy
Conformers. Datasets like QM9 and PCQ use DFT to compute atom positions for which each molecule is in equilibrium, i.e., the forces acting on all atoms are zero:
$$
\mathbf {F} = - \frac {\partial \text {D F T} _ {-} \text {e n e r g y} (\text {a t o m} _ {-} \text {p o s i t i o n s})}{\partial \text {a t o m} _ {-} \text {p o s i t i o n s}} = \mathbf {0} \in \mathbb {R} ^ {\text {n u m} _ {-} \text {a t o m s} \times 3}.
$$
Equilibrium atom positions are computed through structure optimization, typically done by minimizing energy with respect to atom positions using standard optimisation methods such as conjugate gradients or LM-BFGS. This takes  $\approx 50$  iterations of LM-BFGS where each iteration use DFT to
$^9$ GDB is released under the Creative Commons Attribution 4.0 International. We received written consent from the author of GDB use it for QM1B.
compute F. One can thus get  $50\mathrm{x}$  more examples of DFT by evaluating DFT on atom positions that are not structure optimised (similar in spirit to ANI-1 [8]). By utilising the ETKDG algorithm to generate conformers we get  $50\mathrm{x}$  more data, while relying on the ETKDG algorithm to explore the conformational space. We investigate the diversity of conformers generated using this method in Figure 4(a), which depicts the distribution of the standard deviation of each molecule's HL Gap over all conformers. Figure 4(b) shows the distribution of conformers of a few molecules from the extremes of the dataset, and visualises several different conformers for each molecule super-imposed on top of each other.
Figure 4: (a) If we compute the standard deviation of HL gaps over all conformers for a single SMILES string, then this figure plots the distribution of those standard deviations over all SMILES strings in the dataset. (b) The distribution of HL gaps across different conformers for 3 example molecules. These examples are chosen from the left tail, the peak, and the right tail of the distribution of all HL gap variances seen in (a).