SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lower Resolution. For a fixed compute budget to generate data for machine learning, we can think of the DFT options as trading-off dataset size for data quality (DFT accuracy). Relative to prior datasets, we chose a less accurate DFT to generate more training examples. To compare the resulting inaccuracies, we compare the HL gap computed for the same molecule with different options. In particular, Figure 5 compares our STO-3G/B3LYP options against 6-31G*/B3LYP options (attempting to mimic the options from PCQ as closely as possible). The mean absolute error between the options are  $360\mathrm{meV}$ , much larger than chemical accuracy  $(43\mathrm{meV})$  or the errors attained by SchNet  $(63\mathrm{meV})$ . It remains an open problem how pretraining on 1B such low-resolution conformers may bias a Neural Network subsequently fine-tuned on a smaller but more accurate dataset like PCQ. As a preliminary investigation, we pre-trained SchNet on QM1B and then finetuned it on QM9. Compared to only training SchNet on QM9, fine-tuning improved performance from  $54.13\mathrm{meV}$  to  $30.2\mathrm{meV}$ .
Figure 5: Comparison of different DFT options for  $10\mathrm{k}$  molecules from QM1B. Notably, the disagreement between the DFT options are larger than chemical accuracy. It remains to be found how pretraining on 1B molecules with "lower resolution DFT" biases subsequent models fine-tuned on downstream tasks.
KDEplot of HL Energy Gap. To compare HL gap prediction tasks across different datasets we plotted the kernel density estimate (KDE) of the HL gaps. We plot these in Figure 6 for QM9, PCQ and QM1B. To further investigate the consequences of our lower resolution DFT, we also compared our DFT implementation with the optimised atom positions from QM9 against another version of QM9 with atom positions from ETKDG [68] as implemented in RDKit [69].
Basis Sets and Neural Network. While Figure 5 visualises the disagreement between STO3G and 6-31G*, it remains unclear how the disagreement impacts the training behaviour of neural networks. To investigate this we created a version of QM9 with the same DFT options as used for QM1B while using the same atom positions. This allowed us to train the same SchNet model on the same molecules, but with labels computed using different DFT options. Figure 7 shows that the resulting loss curves broadly overlap throughout training.
We also investigated the impact of using off-equilibrium molecules through the RDKit conformers, that is, we recomputed QM9 but using atom positions from RDKit conformers. We found the resulting loss curves were worse (larger MAE), suggesting off-equilibrium prediction may be a fundamentally harder task. We emphasise further experiments are needed to gain certainty, and point out that  $\mathrm{PySCF_{IPU}}$  may allow future research to investigate these types of molecular dataset biases.
Figure 6: Kernel Density Estimate (KDE) plot of the HL gaps in different datasets.
Figure 7: The SchNet 9M model trained on three different versions of QM9 to investigate how choices in dataset creation impact training.
Scaling SchNet Baseline on QM1B. We trained a 9M parameters SchNet model on varying size subsets of the full QM1B dataset. Each subsets consists of N molecule and conformer pairs sampled uniformly from QM1B. For the full dataset a single epoch of training is performed with the number of epochs for each smaller subset adjusted to keep the total number of iterations fixed. The results are presented in Figure 1 and report mean values for three independent seeds. They show that the validation MAE improves as more samples are added to the dataset, as expected. Interestingly we see that the training and validation performance plateaus after approximately 100M samples are in the dataset, this suggests that the model is underfitting, and increasing the model size would further improve performance.
Application Scenario: Pre-training on QM1B, fine-tuning on QM9. After pre-training the 9M SchNet on QM1B we fine-tuned it QM9. Compared to trained the 9M SchNet on QM9 from scratch, pretraining (and subsequent fine-tuning) improved validation performance from  $54.13\mathrm{meV}$  to  $30.2\mathrm{meV}$ . To avoid data leakage, we restricted fine-tuning to a subset of  $88\mathrm{k}$  molecules from QM9 for which their SMILES strings did not appear in the version of GDB used to create QM1B. For discussion of pre-training tasks see Appendix A.
5 Future Work: Larger Molecules and Higher Resolution DFT
The main limitation of our implementation is the fact that it runs out of memory when  $N > 70$ , where  $N$  is the number of atomic orbitals. This limited us to use  $\leq 12$  atoms in STO-3G with a small grid size (10-30k).
Current (Naive) Bottleneck. We can view  $\mathbf{I}^{2e}$  as a matrix  $M\in \mathbb{R}^{N^2,N^2}$  so  $K(\pmb {\rho}) = Mx$  where  $x = \mathrm{flat}(p)$ . For  $N = 70$  we spend 173MB computing  $v = Mx$ . We represent  $M$  with the 8x
symmetry as a list of matrices of size (num_integrals, integral_size).  $M$  is stored distributed over all the cores inside one IPU and thus never needs to move. We compute  $v_{i} = Mx$  by copying  $x$  for every thread allowing us to compute  $x_{i} = M_{\text{thread}_{i}}x$  where  $M_{\text{thread}_{i}}$  is the part of  $M$  stored on core  $i$ . The result is then  $v = \sum_{i} x_{i}$ . This uses  $N^{2}$  number_of Threads floats (or 173MB). The memory consumption can be reduced to 173MB/N=2.47MB by splitting  $x$  into  $N$  batches each with  $N$  floats. The current strategy was a stepping stone, the main advantage is ease of implementation. Notably, prior work disregarded this 8x symmetry at the cost of a 5-8x slowdown [63].
Recomputation. Instead of precomputing  $\mathbf{I}^{2e}$  we can recompute the entries of  $\mathbf{I}^{2e}$  whenever needed during the simultaneous computation of  $K(\pmb{\rho},\mathbf{I}^{2e})$  and  $J(\pmb{\rho},\mathbf{I}^{2e})$ . For the case of Figure 3(a) it took 15.3M out of 105M cycles to compute  $\mathbf{I}^{2e}$ . Recomputing  $\mathbf{I}^{2e}$  all 20 iterations would increase cycle count from 105M to an estimated 395.7M cycles (57.5ms to 216.8ms on a 1.825GHz Mk2 BOW). Finally, if we also recompute the evaluation of the atomic orbitals on the XC grid our memory consumption becomes  $N^2 + \text{grid\_size} \cdot 4 + B$ , where  $B$  is the memory used while (re)computing  $K(\pmb{\rho},\mathbf{I}^{2e})$  and  $J(\pmb{\rho},\mathbf{I}^{2e})$  (independent of  $N$ ).
Using Multiple IPUs. We exemplify our plan to parallelise over 15GB SRAM in 16 IPUs following Figure 3(a). The computation of  $\mathbf{I}^{2e}$  (Eletron Repulsion Integral) can be split over 16 IPUs in the same way we already utilise MIMD parallelism to split them over the 8832 threads. The computation of  $K$  and  $J$  can also be split over 16 IPUs with one REDUCE_SUM of size  $N^2$ . The remainder of the computation can be repeated independently on each IPU.
6 Discussion.
Prior work usually falls in one of the following three categories: (1) DFT libraries, (2) DFT datasets, and (3) neural networks trained on DFT datasets. This article touches on all three, demonstrating how choices made for DFT libraries and dataset generation (1,2) can impact neural network training (3). As molecular machine learning turns towards training large foundation models, we hope our work will help generate sufficiently large datasets with billions (or trillions) of molecules. Furthermore, we hope to allow researchers to generate datasets for downstream tasks used solely to fine-tune such foundation models.
7 Broader Impact.
Our work attempts to improve neural network approximations to Density Functional Theory (DFT) by generating larger datasets. Improving DFT holds the promise to accelerate the development of new medicines and materials, with positive implications for health outcomes and global warming. However, the same technology may someday also allow bad actors to design pathogens, viruses or weapons.
We thank Graphcore research for help and feedback through the entire project. In particular, a huge thanks to Charlie Blake for help with numerical analysis and beautiful visualizations. We also thank Ivan Franzoni, Hadrien Mary Prudencio Tossou, Therence Bois and Cristian Gabellini for valuable feedback. The sole funder of this project was Graphcore.
[1] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling Laws for Neural Language Models. arXiv preprint arXiv:2001.08361, 2020.
[2] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12104-12113, 2022.
[3] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural Message Passing for Quantum Chemistry. In International conference on machine learning. PMLR, 2017.
[4] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, et al. Open Catalyst 2020 (OC20) Dataset and Community Challenges. Acs Catalysis, 2021.
[5] Maho Nakata, Tomomi Shimazaki, Masatomo Hashimoto, and Toshiyuki Maeda. PubChemQC PM6: Data Sets of 221 Million Molecules with Optimized Molecular Geometries and Electronic Properties. Journal of Chemical Information and Modeling, 60(12):5891-5899, 2020.
[6] Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(11):2864–2875, 2012.
[7] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum Chemistry Structures and Properties of 134 Kilo Molecules. Scientific data, 2014.
[8] Justin S Smith, Alexandre Isayev, and Adrian E Roitberg. ANI-1: An Extensible Neural Network Potential With DFT Accuracy at Force Field Computational Cost. Chemical science, 2017.
[9] Maho Nakata and Tomomi Shimazaki. PubChemQC Project: A Large-Scale First-Principles Electronic Structure Database for Data-Driven Chemistry. Journal of Chemical Information and Modeling, 2017.
[10] Zhe Jia, Blake Tillman, Marco Maggioni, and Daniele Paolo Scarpazza. Dissecting the graphcore ipu architecture via microbenchmarking. arXiv preprint arXiv:1912.03413, 2019.
[11] Tobias Fink and Jean-Louis Reymond. Virtual Exploration of the Chemical Universe up to 11 Atoms of C, N, O, F: Assembly of 26.4 Million Structures (110.9 Million Stereoisomers) and Analysis for New Ring Systems, Stereochemistry, Physicochemical Properties, Compound Classes, and Drug Discovery. Journal of chemical information and modeling, 2007.
[12] Tobias Fink, Heinz Bruggesser, and Jean-Louis Reymond. Virtual exploration of the small-molecule chemical universe below 160 daltons. Angewandte Chemie International Edition, 2005.
[60] Li Li, Stephan Hoyer, Ryan Pederson, Ruoxi Sun, Ekin D Cubuk, Patrick Riley, Kieron Burke, et al. Kohn-Sham Equations as Regularizer: Building Prior Knowledge Into Machine-Learned Physics. Physical review letters, 126(3):036401, 2021.
[61] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
[62] Amit Sabne. XLA: Compiling Machine Learning for Peak Performance. 2020.
[63] Xing Zhang and Garnet Kin-Lic Chan. Differentiable Quantum Chemistry with PySCF for Molecules and Materials at the Mean-Field Level and Beyond. The Journal of Chemical Physics, 2022.
[64] Susi Lehtola, Conrad Steigemann, Micael J. T. Oliveira, and Miguel A. L. Marques. Recent developments in libxc — a comprehensive library of functionals for density functional theory. SoftwareX, 7:1-5, Jan 2018.
[65] Qiming Sun. Libcint: An efficient general integral library for gaussian basis functions. (arXiv:1412.0649), Dec 2014. arXiv:1412.0649 [physics].
[66] Kunhao Zheng and Min Lin. JAX-XC: Exchange Correlation Functionals Library in Jax. In Workshop on "Machine Learning for Materials" ICLR 2023.
[67] Ryan Pederson, John Kozlowski, Ruyi Song, Jackson Beall, Martin Ganahl, Markus Hauru, Adam GM Lewis, Yi Yao, Shrestha Basu Mallick, Volker Blum, et al. Large Scale Quantum Chemistry with Tensor Processing Units. Journal of Chemical Theory and Computation, 2022.
[68] Sereina Riniker and Gregory A Landrum. Better Informed Distance Geometry: Using What We Know To Improve Conformation Generation. Journal of chemical information and modeling.
[69] Greg Landrum et al. RDKit: A Software Suite For Cheminformatics, Computational Chemistry, and Predictive Modeling. Greg Landrum, 2013.
A Pre-training Quantum Foundation Models
While it remains an open question which quantum labels are best suited for pre-training, we believe there is a strong argument for pre-training on energy. DFT computes quantum properties by solving the optimization problem:
$$
\text {e n e r g y} ^ {*} = \min  _ {\text {e l e c t r o n} \text {d e n s i t y}} \text {E n e r g y} [ \text {e l e c t r o n} \text {d e n s i t y} ].
$$
All the quantum properties are then derived from the electron density. Pre-training the foundation model to predict energy thus asks the foundation model to perform the same task as DFT. It remains an open problem whether combining energy predicting with other tasks for pre-training would lead to better performance.[10]
Self-supervision. Self-supervised NLP bypass the need for labels  $y$  by predicting parts of the input  $x$  from  $x$ . This is useful because it circumvents creating  $y$  which require expensive human annotation. In Quantum Chemistry,  $x$  is a molecule represented by atom positions and atom types, while  $y$  is a quantum chemical property. We note two differences to NLP
1. In Quantum Chemistry  $y$  may be computed with DFT and require no human annotation.
2. If the input  $x$  arrises from classical chemistry simulations  $x$  contains no quantum information. If  $x$  arrises from quantum chemistry (e.g. PCQ), the compute required to get  $y$  is around 1% of the compute required to compute  $x$ .
<sup>10</sup>Note: Our SchNet experiments trained on gap between molecular orbital energies (HLgap).