SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We present our benchmarking results on OC20 S2EF in Table 2, which summarizes the performance of the five models with respect to the following four metrics: energy and force MAE (i.e., the mean absolute error between ground truth and predicted energies and forces); force cos (i.e., the cosine similarity between ground truth and predicted forces); and energy and forces within threshold (EFwT) - i.e., the percentage of systems whose predicted energies and forces are within a specified threshold from the ground truth [17]. Since force cos and EFwT are correlated with energy and force MAE, we focus on the latter throughout the paper but present all four for completeness.
Table 2: Evaluation of the performance of our five baseline models on the OC20 S2EF task. All models are trained on the OC20-2M dataset. Values represent the average across the four available validation sets. Results for individual validation datasets are provided in Appendix B.
Model	Inference Throughput	OC20 S2EF Validation
Samples / GPU sec. ↑	Energy MAE meV ↓	Force MAE meV/Å ↓	Force cos ↑	EFwT % ↑
SchNet	1100	1308	65.1	0.204	0
PaiNN-small	680	489	47.1	0.345	0.085
PaiNN	264	440	45.3	0.376	0.14
GemNet-OC-small	158	344	31.3	0.524	0.51
GemNet-OC	107	286	25.7	0.598	1.06
<sup>3</sup>https://github.com/Open-Catalyst-Project/ocp
Our results highlight the substantial trade-off between predictive accuracy and computational cost across the GNN architectures, and, therefore, the need for methods that can alleviate this limitation. We observe the same trend on the COLL dataset (see Appendix C).
Similarity analysis of baseline models. To make our analysis exhaustive, we set out to design experiments involving teacher and student architectures of a variable degree of architectural disparity. As a proxy of that, we derive similarity scores based on central kernel alignment (CKA) [39, 40, 35]. In particular, we calculate the pairwise CKA similarity between the node features of our trained SchNet, PaiNN and GemNet-OC models. The results of this analysis are summarized in Figure 2. Focusing on intramodel similarities first (plots on the diagonal), we observe that, while representations from different layers within PaiNN and SchNet have a generally high degree of similarity, GemNet-OC exhibits the opposite behavior, with features extracted at each layer being significantly different from those captured across the rest of the architecture. This is consistent with the architectures of these three models, with features in PaiNN and SchNet being iteratively updated by adding
Figure 2: Similarity analysis between the node features of SchNet, PaiNN and GemNet-OC using CKA (averaged over  $n = 987$  nodes).
displacement features computed at each layer, while those in GemNet-OC representing separate output features. When examining inter-model similarity instead, we notice that, generally speaking, node features in SchNet and PaiNN are similar, whereas those between SchNet and GemNet-OC, and PaiNN and GemNet-OC, diverge significantly as we move deeper into GemNet-OC.
Knowledge distillation results. Based on the aforementioned analyses, we define the following teacher-student pairs, covering the whole spectrum of architectural disparity: PaiNN to PaiNN-small (same architecture); PaiNN to SchNet (similar architectures); GemNet-OC to PaiNN (different architectures). We additionally explore KD from GemNet-OC to GemNet-OC-small (same architecture) on OC20. We train student models by utilizing an offline KD strategy [23], where we distill knowledge from the more competent, pre-trained teacher model to the simpler, more lightweight student model during the training of the latter. We augment the training of each student model with our KD protocols and evaluate the effect on predictive accuracy against the models trained without KD. If not mentioned otherwise, we utilize the following setup: we use MSE as a distillation loss  $\mathcal{L}_{\mathrm{feat}}$ ; a learned linear layer as a transformation function  $\mathcal{M}_s$  applied to the features of the student; and the identity transformation as  $\mathcal{M}_t$ . When distilling knowledge from/into GemNet-OC models, we use the aggregated node- and edge-level output features, which is reminiscent of the review-based setup proposed in [41]. For PaiNN and SchNet, we use the final node features.
The results of our experiments are summarized in Tables 3 and 4, presenting a comparative analysis of the predictive performance of different student models trained with and without the implementation of knowledge distillation on the OC20-2M and COLL datasets, respectively. Focusing on energy predictions first, we observe that, by utilizing KD, we achieve significant improvements in performance in virtually all teacher-student configurations. In particular, we manage to close the gap in performance between student and teacher models by  $\sim 60\%$  or more in six out of the seven configurations, reaching results as high as  $96.7\%$  (distilling PaiNN to PaiNN-small on the COLL dataset). Putting our results into context, we remark that our PaiNN model trained with  $n2n$  KD from GemNet-OC, for instance, provides more accurate energy predictions than more advanced models such as GemNet-dT [7, 8] which is substantially slower. The only setup where results are not as definite is when training SchNet on OC20-2M with KD from PaiNN, where we close  $10.8\%$ . However, it is noteworthy to highlight that this still corresponds to a significant absolute improvement (i.e., notice the big initial difference in performance between the two baseline models on this dataset).
We similarly observe an improvement in the accuracy of student models in force predictions in all teacher-student configurations. Although we observe force improvements (typically  $\sim 5 - 25\%$ ) that are generally not as pronounced as those achieved in energy prediction, we note that we reach
Table 3: Evaluation of the performance of our KD strategies across teacher-student architectures on the OC20 S2EF task. All models are trained on the OC20-2M dataset. Numbers in brackets represent the proportion of the gap between the student  $(S)$  and the teacher  $(T)$  that has been closed by the respective KD strategy (in%). Best results are given in bold. Values represent the average across the four available validation sets. Results for individual validation datasets are provided in Appendix B. Error bars for selected configurations can be found in Appendix F.
Model	OC20 S2EF Validation
Energy MAE meV ↓	Force MAE meV/Å ↓	Force cos ↑	EFwT % ↑
same	S: PaiNN-small	489	47.1	0.345	0.085
T: PaiNN	440	45.3	0.376	0.139
Vanilla (1)	515(-52.4%)	48.5(-81.0%)	0.269(-237%)	0.07(-28%)
Vanilla (2)	476(27.2%)	50.8(-215%)	0.307(-117%)	0.068(-32.6%)
n2n	457 (64.8%)	46.7 (20.5%)	0.348 (9.3%)	0.085 (0.5%)
v2v	459(60.8%)	47.2(-9.1%)	0.347(6.8%)	0.079(-11.9%)
S: GemNet-OC-small	344	31.3	0.524	0.51
T: GemNet-OC	286	25.7	0.598	1.063
Vanilla (1)	339(9.1%)	31.2(-1.0%)	0.525(1.3%)	0.51(0.4%)
Vanilla (2)	328(27.8%)	31.2(0.3%)	0.525(1.5%)	0.61 (18.3%)
n2n	310 (58.8%)	31.1(2.2%)	0.526(3.6%)	0.61(18.2%)
e2e	334(16.5%)	29.7 (27.6%)	0.543 (25.7%)	0.58(12.6%)
similar	S: SchNet	1308	65.1	0.204	0
T: PaiNN	440	45.3	0.376	0.139
Vanilla (1)	1214 (10.8%)	64.6(2.3%)	0.230 (15.2%)	0.003 (1.8%)
Vanilla (2)	1216(10.5%)	64.6 (2.5%)	0.229(14.5%)	0(0%)
n2n	1251(6.6%)	65.2(-0.5%)	0.223(11.1%)	0(0%)
different	S: PaiNN	440	45.3	0.376	0.139
T: GemNet-OC	286	25.7	0.598	1.063
Vanilla (1)	440(0.0%)	43.9(7.1%)	0.378(0.8%)	0.14(0.4%)
Vanilla (2)	419(13.6%)	114.8(-353%)	0.324(-23.8%)	0.127(-1.3%)
n2n	346 (60.8%)	42.8(12.8%)	0.393(7.4%)	0.262 (13.4%)
e2n	418(14.2%)	41.3 (20.5%)	0.405 (12.8%)	0.207(7.4%)
v2v	437(1.8%)	42.9(17.1%)	0.397(9.4%)	0.124(-1.6%)
results as high as  $62.5\%$  for some configurations (distilling PaiNN to PaiNN-small on COLL). A possible reason for this difference in improvement between energy and forces could be attributed to the nature of the supervised task - there are substantially more force labels (i.e., one 3D vector per atom, which could be hundreds per sample) than energy labels (i.e., one per sample). Consequently, we hypothesize it is easier for models to learn to make accurate force predictions, and, therefore, there is more room for improvement in the energy predictions, which we can target with KD.
All in all, our experiments demonstrate that, by applying knowledge distillation, we successfully enhance the performance of student models across all teacher-student configurations and datasets, confirming the effectiveness and robustness of the approach in the context of molecular GNNs. We reiterate the fact that no modifications to the student architectures are made, meaning we achieve an out-of-the-box boost in accuracy without impacting inference throughput.
Effect of KD on model similarity. We continue our investigation by exploring how KD affects student models and their features. To this end, we analyze how the CKA similarity between teacher and student models changes with the introduction of KD during training. We present the outcome of one of our analyses in Figure 3, where we summarize how the similarity between the node features of GemNet-OC and PaiNN changes with the implementation of  $n2n$  KD. We observe that KD introduces strong and specific similarity gains in the layers we use for KD, which also propagates along the
Table 4: Evaluation results on the COLL test set. Numbers in brackets represent the proportion of the gap between the student (S) and the teacher (T) that has been closed by the respective KD strategy (in %). Best results are given in **bold**.
Model	COLL test set
Energy MAE meV ↓	Force MAE meV/Å ↓	Force cos ↑	EFwT % ↑
same	S: PaiNN-small	104.0	80.9	0.984	5.4
T: PaiNN	85.8	64.1	0.988	10.1
Vanilla (1)	106.1(-11.5%)	82.0(-6.5%)	0.984(2.3%)	4.46(-20.2%)
Vanilla (2)	86.4 (96.7%)	80.9(0%)	0.983(-2.3%)	4.3(-23.7%)
n2n	92.5(63.2%)	77.8(18.5%)	0.984(18.2%)	6.63 (26.5%)
v2v	90.4(74.7%)	70.4 (62.5%)	0.986 (45.5%)	5.8(8.4%)
similar	S: SchNet	146.5	121.2	0.970	2.75
T: PaiNN	85.8	64.1	0.988	10.1
Vanilla (1)	146.1(0.7%)	120.8(0.7%)	0.970(1.1%)	2.54(-2.9%)
Vanilla (2)	104.1 (69.9%)	120.9(0.5%)	0.970(1.1%)	6.45 (50.7%)
n2n	141.6(8.1%)	117.2 (7.0%)	0.971 (5.4%)	2.63 (-1.6%)
different	S: PaiNN	85.8	64.1	0.988	10.1
T: GemNet-OC	44.8	38.2	0.994	20.2
Vanilla (1)	86.2(-1.1%)	63.9(0.6%)	0.988(1.5%)	10.1(0.1%)
Vanilla (2)	61.4(59.5%)	62.9(4.6%)	0.988(5.2%)	13.0(29.2%)
n2n	60.4 (62.0%)	61.2 (11.3%)	0.989 (14.9%)	13.6 (34.6%)
e2n	77.3(20.8%)	63.3(3.0%)	0.988(7.9%)	11.0(9.2%)
v2v	81.2(11.2%)	63.3(3.1%)	0.988(3.4%)	10.5(4.6%)
student architecture. We notice similar behavior across other teacher-student configurations and KD strategies (see Figures 6 and 7 in Appendix G), allowing us to monitor and quantify the effect of KD on student models as we explore different KD settings and design choices.
Hyperparameter studies. We additionally conduct a thorough hyperparameter study to evaluate the effect of different design choices within our KD framework. We summarize our results below.
- Effect of distillation loss: Apart from our default MSE-based distillation loss  $\mathcal{L}_{\mathrm{feat}}$ , we also experimented with more advanced losses such as Local Structure Preservation (LSP) [34] and Global Structure Preservation (GSP) [35], as well as directly optimizing CKA. We observed the best results with our default MSE loss, with other options substantially hurting accuracy (see Appendix E.1).
- Effect of transformation function: We also investigated a number of different transformation functions and the effect they have on performance. The transformations we utilized include: the identity transformation (when appropriate); learned linear transformations, and MLP projection heads. Our results showed that our default linear mapping is a sufficiently flexible choice as it gives the best results (see Appendix E.2).
- Effect of feature selection: We additionally explored the effect of feature selection on KD performance. In particular, we analyzed the change in the predictive accuracy of PaiNN as a student model as we distill features from earlier layers in the teacher (GemNet-OC), or distill knowledge into earlier layers in the student. Our results suggest that using features closer to the output is the best-performing strategy (see Appendix E.3). We also performed CKA-based similarity analyses to monitor how model similarity changes as we vary the features we used for KD (see Figure 8 in Appendix G).
Data augmentation. As for most other applications, data labeling for molecular data is costly as it requires running computationally expensive quantum mechanical calculations to obtain ground truth energies and forces. Motivated by this, we explore two data augmentation techniques, which we use to generate new data points that we label with the teacher and use for KD. We briefly describe these below (see Appendix D for more details).
Figure 3: Similarity analysis between Gemnet-OC and PaiNN without KD (left) and with KD (right). The feature pair used during KD is indicated with  $\swarrow$ . Similarity analyses for other KD strategies and teacher-student configurations are presented in Appendix G.
- Random rattling: Adding noise to existing structures (also known as "rattling") is a form of data augmentation that has been used in the context of pretraining of molecular GNNs [42, 43], and as a regularization strategy [44]. Inspired by this, we utilized "rattling" in the context of KD, where we added random noise to the atomic positions of systems and used the teacher to derive energy and force labels for these rattled samples. We then combined these rattled structures with the original dataset during the training of student models. However, this approach did not provide significant improvements. Additionally, we tried using gradient ascent to find perturbations that maximize the discrepancy between the teacher and student predictions, similar to [45], but this did not show improvements over random noise, and also increased training time.
- Synthetic Data: Samples in OC20 S2EF originate from the same relaxation trajectory and are therefore correlated. To tackle this, we generated our own distilled dataset coined  $d1M$ , which consists of one million samples generated by sampling new systems (generated with the OC Datasets codebase  $^4$ ), running relaxations with our pre-trained GemNet-OC model, and then subsampling approximately  $10\%$  of the frames. We explored different ways of incorporating this new  $d1M$  dataset, all of which were based on joint training with the OC20 S2EF 2M data (similar to what we did with the rattled systems). To study different combinations of the ground truth DFT samples and the  $d1M$  samples during training, we defined two hyperparameters determining: (a) how many of the samples per batch originate from each of the datasets; and (b) how to weight the loss contributions based on the origin of data. Unfortunately, and contrary to similar approaches, e.g., in speech recognition [46], the results we observed did not significantly improve on the baseline models.
In this paper, we investigated the utility of knowledge distillation in the context of GNNs for molecules. To this end, we proposed four distinct feature-based KD strategies, which we validated across different teacher-student configurations and datasets. We showed that our KD protocols can significantly enhance the performance of different molecular GNNs without any modifications to their architecture, allowing us to run faster molecular simulations without substantially impairing predictive accuracy. With this work, we aim to elucidate the potential of KD in the domain of molecular GNNs and stimulate future research in the area. Interesting future directions include: the combination of KD strategies (e.g.,  $n2n$  and  $\nu 2\nu$ ); extending the framework to other types of features (e.g., tensorial features [47]), molecular tasks and datasets; better understanding the connection between KD performance and model expressivity (e.g., can model similarity inform KD design); and performing a more comprehensive stability analysis [48]. One caveat of our approach is that even though inference times are not affected, training times are, albeit not necessarily if a pretrained teacher model is available (see Appendix H). Finally, it is important to recognize that such technologies, while innovative, could be used for potentially harmful purposes, such as the simulation or discovery of toxic systems, or the development of harmful technologies.
<sup>4</sup>https://github.com/Open-Catalyst-Project/Open-Catalyst-Dataset
F.E.K. is financially supported by the Excellence Center at Linköping-Lund in Information Technology (ELLIIT). D.G. is supported by UK Research and Innovation [UKRI Centre for Doctoral Training in AI for Healthcare grant number EP/S023283/1]. Computing resources provided by: the Berzelius resource at the National Supercomputer Centre, provided by Knut and Alice Wallenberg Foundation; the Alvis resource provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) at Chalmers e-Commons at Chalmers (C3SE) partially funded by the Swedish Research Council through grant agreement no. 2022-06725; the Chair of Aerodynamics and Fluid Mechanics at Technical University of Munich. This research project was initially conceived at the 2022 LOGML summer school, and we would like to thank Guocheng Qian and I-Ju Chen for their contribution during the early conceptualizing stages of this project during and in the first weeks following the summer school. We also thank the Open Catalyst team for their open-source codebase, support and discussions. In particular, Muhammed Shuaibi for providing the COLL dataset in LMDB format. Figures assembled in BioRender.
[1] Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annual review of physical chemistry, 71:361-390, 2020.
[2] Yuyang Wang, Zijie Li, and Amir Barati Farimani. Graph neural networks for molecules, 2023.
[3] 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 Liö, 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, 2023.
[4] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, May 2022.
[5] Johannes Gasteiger, Janek Groß, and Stephan Gunnemann. Directional Message Passing for Molecular Graphs. In International Conference on Learning Representations, 2020.
[6] Johannes Gasteiger, Shankari Giri, Johannes T. Marggraf, and Stephan Gunnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurIPS, 2020.
[7] Johannes Gasteiger, Florian Becker, and Stephan Gunnemann. GemNet: Universal Directional Graph Neural Networks for Molecules. In Advances in Neural Information Processing Systems, volume 34, pages 6790-6802. Curran Associates, Inc., 2021.
[8] Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Gunnemann, Zachary Ward Ulissi, C. Lawrence Zitnick, and Abhishek Das. GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets. Transactions on Machine Learning Research, October 2022.
[9] Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems, 35:8054-8067, 2022.
[43] Rishikesh Magar, Yuyang Wang, and Amir Barati Farimani. Crystal twins: Self-supervised learning for crystalline material property prediction. npj Computational Materials, 8(1):1-8, November 2022.
[44] Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Velicković, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971, 2021.
[45] Daniel Schwalbe-Koda, Aik Rui Tan, and Rafael Gómez-Bombarelli. Differentiable sampling of molecular geometries with uncertainty-based adversarial attacks. Nature Communications, 12(1):5104, August 2021.
[46] Yu Zhang, James Qin, Daniel S Park, Wei Han, Chung-Cheng Chiu, Ruoming Pang, Quoc V Le, and Yonghui Wu. Pushing the limits of semi-supervised learning for automatic speech recognition. arXiv preprint arXiv:2010.10504, 2020.
[47] Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks, 2022.
[48] Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarelli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations, 2022.
A Training and hyperparameters
Models were trained on NVIDIA A100 40 GB and NVIDIA RTX A6000 48 GB GPUs, except GemNet-OC-small which were trained on NVIDIA A100 80 GB and NVIDIA RTX A6000 48 GB. All models were trained on single GPUs, except for SchNet when trained on OC20-2M, which required 3 GPUs. Inference throughput was profiled on A100 40 GB GPUs, with reported values representing approximate numbers averaged across three evaluations. We provide detailed information about the hyperparameters we used for each model in Tables 5, 6, and 7.
Moreover, we summarize the KD weighting factors  $\lambda$  we used for each model configuration in Table 8.
Table 5: SchNet hyperparameters.
Hyperparameter	OC20	COLL
Hidden channels	1024	128
Filters	256	128
Interaction blocks	5	6
Gaussians	200	50
Cutoff	6.0	12.0
Batch size	192	32
Initial learning rate	\( 10^{-4} \)	\( 10^{-3} \)
Optimizer	AdamW	AdamW
Scheduler	LambdaLR	LinearWarmupExponentialDecay
Learning rate decay factor	0.1	0.01
Learning rate milestones	52083, 83333, 104166	-
Warmup steps	31250	3750
Warmup factor	0.1	-
Force Coefficient	100	100
Energy Coefficient	1	1
Number of epochs	30	500
Table 6: PaiNN hyperparameters.Slash-separated values indicate PaiNN versus PaiNN-small hyperparameters.
Hyperparameter	OC20	COLL
Hidden channels	512/256	256/128
Number of layers	6/4	6/4
Number of RBFs	128	128
Cutoff	12.0	12.0
Max. num. neighbors	50	50
Direct Forces	True	True
Batch size	32	32
Optimizer	AdamW	AdamW
AMSGrad	True	True
Initial learning rate	10-4	10-3
Scheduler	LambdaLR	LinearWarmupExponentialDecay
Warmup steps	None	3750
Learning rate decay factor	0.45	0.01
Learning rate milestones (steps)	160000, 320000, 480000, 640000	-
Force coefficient	100	100
Energy coefficient	1	1
EMA decay	0.999	0.999
Gradient clip norm threshold	10	10
Epochs	16	375