Accelerating Molecular Graph Neural Networks via Knowledge Distillation
Authors: Filip Ekström Kelvinius
Dimitar Georgiev
Artur Toshev
Johannes Gasteiger
Venue: NeurIPS 2023 poster (2023)
Keywords: GNN
graph neural networks
knowledge distillation
molecules
molecular simulations

Abstract: Recent advances in graph neural networks (GNNs) have enabled more comprehensive modeling of molecules and molecular systems, thereby enhancing the precision of molecular property prediction and molecular simulations. Nonetheless, as the field has been progressing to bigger and more complex architectures, state-of-the-art GNNs have become largely prohibitive for many large-scale applications. In this paper, we explore the utility of knowledge distillation (KD) for accelerating molecular GNNs. To this end, we devise KD strategies that facilitate the distillation of hidden representations in directional and equivariant GNNs, and evaluate their performance on the regression task of energy and force prediction. We validate our protocols across different teacher-student configurations and datasets, and demonstrate that they can consistently boost the predictive accuracy of student models without any modifications to their architecture. Moreover, we conduct comprehensive optimization of various components of our framework, and investigate the potential of data augmentation to further enhance performance. All in all, we manage to close the gap in predictive accuracy between teacher and student models by as much as 96.7\% and 62.5\% for energy and force prediction respectively, while fully preserving the inference throughput of the more lightweight models.

Accelerating Molecular Graph Neural Networks via Knowledge Distillation
Filip Ekström Kelvinius" Linköping University filip. ekstsroeuuoooml
Dimitar Georgiev" Imperial College London
georgiev21@imperial. ac uk
Artur Petrov Toshev
Johannes Gasteiger
Technical University of Munich artur. towtur.ttshevttm.co
Google Research johanesg@googgee com
Abstract
Recent advances in graph neural networks (GNNs) have enabled more compre- hensive modeling of molecules and molecular systems, thereby enhancing the precision of molecular property prediction and molecular simulations. Nonetheless, the field has been progressing to bigger and more complex architectures, state- of-the-art GNNs have become largely prohibitive for many large-scale applications. In this paper, explore the utility of knowledge distillation (KD) for accelerating molecular GNNs. To this end, we devise KD strategies that facilitate the distillation of hidden representations in directional and equivariant GNNs, and evaluate their performance on the regression task of energy and force prediction. We validate our protocols across different teacher-student configurations and datasets, and demonstrate that they can consistently boost the predictive accuracy of student models without any modifications their architecture. Moreover, conduct com- prehensive optimization of various components of our framework, and investigate the potential of data augmentation further enhance performance. All in all, we manage to close the gap in predictive accuracy between teacher and student models by as much as 7% and 62.5% for energy and force prediction respectively, while fully preserving the inference throughput of the more lightweight models.
Introduction
In the last couple of years, the field of molecular simulations has undergone rapid paradigm shift with the advent of new, powerful computational tools based on machine learning (ML) [. At the forefront of this transformation have been recent advances in graph neural networks (GNNs), which have brought about architectures that more effectively capture geometric and structural information critical for the accurate representation of molecules and molecular systems 3 a multitude of GNNs have been developed, which now offer predictive performance approaching that of conventional gold-standard methods such as density functional theory (DFT) fraction of 5 6
Consequently,
g 9 This has, in turn, significantly accelerated the modeling of molecular properties and the simulation of molecular systems, bolstering new research developments
the computational cost 4
in many scientific disciplines, including material sciences, drug discovery and catalysis [10 12
13 [4].
Nonetheless, this progress largely coinciding with the development of bigger and more complex
models, has naturally come the expense of increased complexity [15 9 16). This has gradually
TTeee authors contributed equally to this work. Order was determined by rolling dice.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

OC20-2M
COLL
GemNet-OC
PaiNN
PaiNN
GemNet-OC
PaiNN
PaiNN
PaiNN-small
PaiNN-small
13K ARTORee/aa
Energy
Farce
Eneeaay
Porce
Eneray
Forte B3.0
Reerny
Force
Stucent
20 555
20.5%
$L.2
KD
60.8%
84--%%
$22 82.0%
62-5
Starest
Figure I: Using knowledge distillation, manage significantly boost the predictive accuracy of different student models on the OC20-2M [7] inference throughput.
and COLL 66 datasets while fully preserving their
limited the utility of state-of-the-art GNNs in large-scale molecular simulation applications (e.g.. molecular dynamics and high-throughput searches), where inference throughput (i.e., how many samples can be processed for given time) is critical for making fast predictions about molecular systems at scale. Hence, addressing the trade-off between accuracy and computational demand remains essential for creating more affordable tools for molecular simulations and expanding the transformational impact of GNN models in the
Motivated by that, in this work, we investigate the potential of knowledge distillation (KD) in advancing the speed-accuracy Pareto frontier and enhancing the performance and scalability of molccular GNNs. summary, the contributions of this paper are as follows:
We. for the first time, explore the utility of knowledge distillation for accelerating GNNs for molecular simulations lcage- multi-output regression task, challenging to address with common KD methods.
We design custom KD strategies, which we call node-to-node (n2n), edge-to-edge (e2e), edge-to- node (e2n) and veclor-to-veclor (v2v) knowledge distillation, which facilitate the distillation of hidden representations directional and equivariant molecular GNNs.
We demonstrate the effectiveness of our protocols across different teacher-student configurations and datasets, allowing us to substantially improve the performance student models while fully preserving their throughput (see Figure for an overview).
We conduct a comprehensive empirical analysis of different components of our KD strategies, as well as explore data augmentation techniques for further improving performance.
Associated code is available online
2 Background
Molecular simulations. In this work, consider molecular systems atomic level, i.e.. N atoms represented by their atomic numbers z = {.,...., E N and positions X = {x1..... xN e Nx3 Given system, want model that can predict the energy E R of the system, and the forces E R x3 acting on each atom. Both of these properties are of high interest when simulating molecular systems. The energy of system is essential for the prediction of its stability, whereas the forces are important for molecular dynamics simulations, where computed forces are combined with the equations of motion to simulate the evolution of the system over time.
GNNs for molecular systems. GNNs are a suitable framework for modeling molecular systems. Each molecular system (X, =) can be represented as mathematical graph G (V. E). where the nodes V correspond the set of atoms, and edges E are created between nodes by connecting the closest neighboring atoms (typically defined by cutoff radius and/or maximum number of neighbors). Hence, in the context of molecular simulations, we can create GNNs that operate on

---

atomic graphs G by propagating information between the atoms and the edges, and make predictions about the energy and forces of each system in multi-output manner i.e., E, F GNN(X, 2).
The main problem when modeling molecules and molecular properties the number of underlying symmetries to account for, most importantly rigid transformations of the atoms. For instance, the total energy E of a system is not affected by (i.e., invariant to) rotations and translations of the system. However, the forces F do change as we rotate a system i.e., they are equivariant to rotations. Therefore, to make accurate predictions about molecular systems, it is crucial to devise models that respect these symmetties and other physical constraints. There is now a plethora of diverse molecular GNNs that reflect that, e.g., SchNet [18], DimeNet (5 PaiNN [9]. GemNet 8. NequlP 4. and SCN 19, which have incrementally established more holistic description of molecular systems by capturing advanced geometric features and physical symmetries. This has, however, come the expense of computational efficiency.
Knowledge distillation. Knowledge distillation is a technique for compressing and accelerating ML models (20). which has recently demonstrated significant potential in domains such as computer vision 21 and natural language modeling 122 The main objective of KD is create more efficient models by means of transferring knowledge (e.g., model parameters and activations) from large, computationally expensive, more accurate models, often referred to as teacher models, to simpler, more efficient models called student models [23]. Since the seminal work of Hinton et al. [24], the field has drastically expanded methodologically with the development of protocols that accommodate the distillation of "deeper" knowledge, more comprehensive transformation functions, as well as more robust distillation losses [23] 225. Yet, these advances have mostly focused on classification, resulting in methods of limited utility in regression tasks [26]. Moreover, most research in the area has been confined to non-graph data (e.g., images, text, tabular data). Despite recent efforts to extend KD to graph data and GNNs, these have likewise only concentrated on classification tasks involving standard GNN architectures [27] 28]. And, in particular, the application of KD to large-scale regression problems in molecular simulations, which involve offte-of-the-arr molecular GNN architectures containing complex, geometric node- and edge-level features, is still unexplored.
Knowledge distillation for molecular GNNs
Preliminaries. In the context of the aforementioned prediction task, we train molecular GNNs by enforeing a loss Co that combines both the energy and force prediction error as follows:
Lo C F.) + aF
(1)
where E and F are the ground-trute -tuuh energy and forces, É and F are the predictions of the model of interest, and L. and L are some loss functions weighted by ap, Op E R.
To perform knowledge distillation, we augment this training process by defining an auxiliary KD loss term LXD. which is added to Lo (with a factor E R) to derive the final training loss function L:
Co + ALKD
(2)
This was originally proposed in the context of classification by leveraging the fact that the soft label predictions (i.e., the logits after softmax normalization) of given (teacher) model carry valuable information that can complement the ground-trute labels in the training process of another (student) model [24. Since then, this has become the standard KD approach- commonly referred to as vanilla KD in the literature, which is often the foundation of new KD protocols. The main idea of this technique is to employ KD loss LKD that forces the student mimic the predictions of the teacher model. This is usually achieved by constructing a loss LKD = KL(2,, 2t) based on the Kullback-Leibler (KL) divergence between the soft logits of the student 2, and the teacher 21.
However, such strategies based on the distillation of the output of the teacher model only pose two significant limitations. First, they are by design exclusively applicable to classification tasks, since there are no outputs analogous to logits in regression setups 22 29 This has consequently limited the utility of most KD methods in regression tasks. Second, this approach forces the student to emulate the final output of the teacher directly, which can be unattainable in regimes where the complexity gap between the two models substantial, and thus detrimental KD performance (30)
Feature-based KD. To cireumvent these shorteomings, we focus on feature-based KD an extension of vanilla KD concerned with the distillation of knowledge across the intermediate layers of models
3

---

In particular, we perform knowledge distillation of intermediate representations by devising a loss on selected hidden features H3 E U, and H, E U, in the student and teacher models respectively, which takes the form:
32 23
LKD - Leea (M.((
where M2 : Us + U and M. : U1 + U are transformations that map the hidden features to a common feature space U. and Lteat U x U + It is some loss of choice. The transformations M, can be the identity transformation, linear projections and multilayer perceptron (MLP) projection heads; whereas for the distillation loss Cisas typical options include mean squared error (MSE) and mean absolute error (MAE).
The fundamental design decision when devising a KD strategy based on the distillation of internal representations is the choice of features to distill between (i.e., features H. and H,). One needs to ensure that paired features are similar both in expressivity and relevance to the output. Most research on feature-based distillation on graphs has so far focused on models that only have one type of scalar (node) features in classification tasks (27 (8 reducing the problem to the selection of layers to pair across the student and the teacher. This is often further simplified by utilizing models that share the same architecture up to reducing the number of blocks/layers and their dimensionality.
Features molecular GNNs. In contrast, molecular GNNs contain diverse features (e.g., scalars, vectors and/or equivariant higher-order tensors based on spherical harmonics) organized across nodes and edges within complex molecular graph. These are continually evolved by model-specific operators to infer molecular properties, such as energy and forces, in multi-output prediction fashion. Therefore, features often represent different physical, geometric and/or topological information relevant to specific parts of the output. This significantly complicates the design of an effective KD strategy, especially when the teacher and the student differ architecturally, as one needs extract and align representations corresponding comparable features in both models.
In this work, we set out to devise KD strategies that representative and effective across various molecular GNNs. This is why we investigate the effectiveness of KD with respect to GNNs that have distinct architectures and performance profiles, and can be organized teacher-student configurations at different levels of architectural disparity. In particular, we employ the following three GNN models, ordered by computational complexity (ascending):
SchNet 33 A simple GNN model based on continuous-filter convolutional layers, which only contains scalar node features E IRd These are used to predict the energy É. The force is then calculated as the negative gradient of the energy with respect to the atomic positions, i.e..
F -.
PaiNN [9:: A GNN based equivariant message passing, which contains scalar node features x E R.d. used for energy prediction; as well as geometric vectorial node features e E R that are equivariant to rotations and can thus be combined with the scalar features to make direct predictions of the forces (i.e., without computing gradients of the energy).
GemNet-OC [ A GNN model that utilizes directional message passing between scalar node features h Rdh and scalar edges features m E RRd After each block layers, these are
(6) processed through an output block, resulting in scalar node features ze and edge features xpp
where is the block number. The output features from each block are aggregated into output features gg and F,, which are used to compute the energy and forces respectively.
An overview of the features of the three models can be found in Table
Table 1: An overview of the types of features available in the three models we use in this study.
SchNet
PaiNN
GemNet-OC
Scalar node features Scalar edge features Vectorial node features Output blocks
Defining feature-based KD distillation strategies for molecular GNNs. In the context of the three models considered in this work, we devise the following KD strategies:

---

node-to-node (n2n): As all three models contain scalar node features Hnode, we can distill knowledge in between these directly by defining loss CDD given by
LKD
Liea (M,(Hnoess) M, ()ode))))
where Hode,s and Hnodert represent the node features of the student and teacher, respectively. Note this is a general approach that utilizes scalar node features only, making it applicable to standard GNNs. Here, we want to force the student to mimic the representations of the teacher for each node (i.e., atom) independently, so we use a loss that directly penalizes the distance between the features in the two models, such as MSE (similar to the original formulation of brature-based KD in Romero et al. [31).. Other recently proposed losses Lient for the distillation of node features in standard GNNs specifically include approaches based on contrastive learning [34 35 36 37 and adversarial training (38). We do not focus on such methods as much since they are better suited for (node) classification tasks (e.g., contrasting different classes of nodes), and not for molecule-lveel predictions.
To take advantage of other types of features relevant to molecular GNNs, we further devise three additional protocols, which we outline below.
edge-to-edge (e2e): The GemNet-OC model heavily relies its edge features, which key component of the directional message passing employed in the architecture. As such, they can be a useful resource for KD. Hence, also consider KD between edge features, which we accomplish by applying Equation (4) to the edge features Hedge,s and Hedge, of the student and teacher, respectively.
edge-to-node (e2n): However, not all models considered in this study contain edge features to distill to. To accommodate that, we propose a KD strategy where we transfer information from GemNet-OC's edge features (idg) by first aggregating them follows:
Hedge2node,
Σ
Hedge:(is)>
(5)
JEN(()
where is some node index. The resulting features HedgeZmndd,, scalar, node-level features, and we can, therefore, use them to transfer knowledge to the student node features Hnode,s as in Equation (4)
veckor-to-vector (v2v): Similarly, the PaiNN model defines custom vectorial node features, which differ from the scalar (node and edge) features available in the other models. These are not scalar and invariant to rigid transformations of the atoms, but geometrical vectors that are equivariant with respect to rotations. As these carry important information about a given system, we also want to define procedure to distill these. When we perform KD between two PaiNN models, we can directly distill information between these vectorial features just as in Equation 4 In contrast, when distilling knowledge into PaiNN from our GemNet-OC teacher that has such vectorial features, transfer knowledge between (invariant) scalar edge features and (equivariant) vectorial node features by noting that scalar edge features sit edge between two atoms in 3D space. Hence, we can aggregate the edge features { (Hodge hjew corresponding to a given node i into node-level equivariant vectorial features H by considering the unit vector Ukj
equivariant 3D grid since they associated with an
[2 a. (x.
æ:4) that defines the direction of the edge (i, such that
(k) H
Σ
Uiy HiyH!M (k)
(A)
(6)
vOo,
JEN(()
vee.; fulfill the condition of equivariance with respect to rotations as each vector Ui is equivariant rotations, and H(() scalar not influencing its direction. Consequently, it is important to loss Lisat that encourages vectors to align in both magnitude and direction c.g., MSE.
with the superscript & indicating the channel. Notice that the features H (k)
(f.j)
Additional KD strategies. We further evaluate two additional KD approaches inspired by the vanilla baset- KD used in classification, which we augment to make suitable for regression tasks:
Vanilla ( (I): One way of adapting vanilla KD for regression is by steering the student to mimic the hnal output of the teacher directly:
LXD CELE(EE En) + apLf(F F.).
5

---

where the subscripts and, refer to the predictions of the student and teacher, respectively. Note that, unlike in classification, this approach does not provide much additional information regression tasks, except for some limited signal about the emorr distribution of the teacher model 20 29]
Vanilla (2): One way enhance the teacher signal during training is consider the fact that many GNNS for molecular simulations make separate atom- and edge-level predictions, which are consequently aggregated into a final oulput. For instance, the total energy E' of system is usually defined as the sum of the predicted contributions from each atom E Σ, E. Hence, we can extend the aforementioned vanilla KD approach by imposing a loss on these granular predictions instead:
N
CKD
CB Eis).
(8)
These individual energy contributions are not part of the labeled data, but, when injected during training, can provide more fine- fine-grained information than the aggregated prediction.
Experimental results
To evaluate proposed methods, we perform comprehensive benchmarking experiments the OC20-2M [17] dataset (structure to energy and forces (S2EF) task) - a large and diverse catalyst dataset; and COLL 66 challenging molecular dynamics dataset. We use the model implementations provided in the Open Catalyst Project (OCP) codebase (see Appendix A for detailed information about training procedure and model hyperparameters)
Benchmarking baseline models. We start by first evaluating the baseline performance of the models we employ in this study. As previously mentioned, we select SchNet, PaiNN and GemNet-OC for our experiments as they cover most of the accuracy-complexity spectrum, with the last representing the state-of-the-art on OC20 S2EF and COLL at the time of experimentation. To demonstrate this, benchmark the predictive accuracy and inference throughput of the models on the two aforementioned datasets. In conjunction with the default PaiNN and GemNet-OC models, also experiment with more lightweight versions of the two architectures referred to as PaiNN-small and GemNet-OC- small respectively, where reduce the number of hidden layers and their dimensionality. We train all models to convergence ourselves, except for the GemNet-OC model on OC20-2M, where we utilize the pre-traines pre-trained model available within the OCP repository (July 2022).
We present our benchmarking results on OC20 S2EF in Table 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.c. the percentage of systems whose predicted energies and forces are within a specified threshold from the ground truth [17]. Since force cos and EFwT 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 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
Inference
OC20 S2EF Validation
Throughput
Samples GPU sec.
Energy MAE
Force MAE mcV/A /
Force cos
EFwT
Model
meV 4
1
%
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
GemNett-O-smman
158
344
31.3
0.524
0.51
GemNet-OC
107
286
25.7
0.598
1.06

---

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 COL.I dataset (see Appendix C
Similarity analysis of baseline models. To make our analysis exhaustive, we set out to de- sign experiments involving teacher and student architectures of a variable degree of architec- tural disparity. As a proxy of that, we derive similarity scores based on central kernel align- ment (CKA) [3] 40 (3]. In particular, we cal- culate 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 intra-
1.0
SchvVer
Layers
0
0.8
PaINN
Layers
-066
Similntty
0
-04
10
GemNet-OC Layers
model similarities first (piots on the diagonal), we observe that, while representations from dif- ferent layers within PaiNN and SchNet have generally high degree of similarity, GemNet-OC exhibits the opposite behavior, with features ex- tracted at each layer being significantly different from those captured across the rest of the archi- tecture. This is consistent with the architectures of these three models, with features in PaiNN and SchNet being iteratively updated by adding displacement features computed 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.
5
0.0
0
11
10
SchNet Layers
PaiNN Layers
GemNet-DC Layers
Figure 2: Similarity analysis between the node features of SchNet, PaiNN and GemNet-OC using CKA (averaged over n = 987 nodes).
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 PaiNN (different archi- tectures). We additionally explore KD from GemNet-OC to GemNet-OC-smaa (same architecture) on OC20. We train student models by utilizing offline KD strategy 23 where we distill knowl- edge 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 Crear a learned linear layer as transformation function MM applied to the features of the student; and the identity transformation as M,. When distilling knowledge from/into GemNet-OG models, we use the aggregated node- and edge-level output features, which is reminiscent of the review-bar based setup proposed in [41. For PaiNN and SchNet, we use the final node features.
The results of our experiments are summarized in Tables and presenting a comparative analysis of the predictive performance of different student models trained with and without the implemen- tation of knowledge distillation on the OC20-2M and COLL datasets, respectively. Focusing on energy predictions first, we observe that, by utilizing KD, 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 ~ 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 Ge 7 g 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 ~--25% that are generally not as pronounced as those achieved in energy prediction, we note that we reach
7

---

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 brackets represent the proportion of the gap between the student (S) and the teacher (7) 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
OC20 S2EF Validation
Energy MAE mcV
Force MAE
Force cos
EFwT % 1
Model
mcV/A
S: PaiNN-small PaiNN
489
47.1
0.345
0.085
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% 22)
50.8(-215%)
0.307(-117%)
0.068(-32.66
n2n
457 (64.8%)
46.7 (20.5%)
0.348 (9.3%)
0.085 (0.5%)
same
v2v
(60.8%)%
47.2(-9.1%)
0.347(6.8%)
0.079(-11.9)
S: GemNet-OC-smee GemNet-OC
344
31.3
0.524
0.51
286
25.7
0.598
1.063
Vanilla (1)
(9.1%)
31.2(-1.0%)
0..25
0.51(0.4%)
Vanilla (2)
(28.87.8%
31.2(0.3%)
0.525
0.61 (18.3%)
n2n
310 (58.8%)
(2.29)
0.526
0.61 (18.2%)
e2e
334 (16.5%)
29.7 (27.6%)
0.543 (25.7%)
0.58(12.6%)
S: SchNet
1308
65.1
0,204
0
similar
T: PaiNN
440
45.3
0.376
0.139
Vanilla (1)
1214 (10.8%)
64.6(2.3%)
0.230 (1).2%)
0.003 (1.8 %)
Vanilla (2)
1216(10.5%)
64.6 (2.5%
(1.29(14.3%)
0(0%)
n2n
1251(6.6%)
65.2(-0.5%)
0.223
0(0%)
S: PaiNN
440
45.3
0.376
0.139
T: GeGnNet-OC
286
25.7
0.598
1.063
different
Vanilla (1)
440(0.0%)
43.9(7.1%)
0.378
0.14(0.4%)
Vanilla (2) n2n
419(13.6%)
114.8(-353%)
0.324(-23.8%)
0.127(-1.3%)
346 (60.8%)
42.8(12.8%)
0.393(7.44%
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.11)
0.397(9.4%)
0.124(-1.6%)
results as high as 62.5% for some configurations (distilling PaiNN to PaiNN-small on COLL). 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, casier for models to learn to make accurate force predictions, and, therefore,
we hypothesize 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 ceacher-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 where we summarize how the similarity between the node features of Ge-Net-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
8

---

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 (7) that has been closed by the respective KD strategy (in %). Best results are given in bold.
COLL test
Energy MAE meV
Foree MAE meV meV/AA
Force cos
EFwT % 1
Model
S: PaiNN-small
104.0
80.9
0.984
5.4
T: PaiNN
85.8
64.1
0.988
10.1
same
Vanilla (1)
106.1(-11.5%)
82.0(-6.5%)
0.984(2.3%)
4.46(-20.2%% 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.77)
70.4 (62.5%)
0.986 (45.5%
5.8(8.4%)
SchNet
146.5
121.2
0.970
2.75
similar
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%)
(-6.6%)
S: PaiNN
85.8
64.1
0.988
10.1
T: GemMNetOOC
44.8
38.2
0.994
20.2
different
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.87) 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 and Appendix G). allowing us to monitor and quantify the effect of KD on student models as we explore different KD settings and design choices.
Hy pyperaaraeeeer 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 Lieat. we also exper- imented 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.I)
Efffect of Transformation function: We also investigated number of different transformation func- tions 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 sufficiently flexible choice as it gives the best results(see Appendix E.2).
Effect of feature selection: We additionally explored the effect feature selection on KD perfor- mance. In partieular, analyzed the change in the predictive accuracy of PaiNN as 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 the output is the best-performing strategy (see Appendix EE.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 g in Appendix G).
Data augmentation. As for most other applications, data labeling for molecular data is costly as 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 label with the teacher and use for KD. We briefly describe these below (see Appendix for more details).

---

no KD
with KD
Similarity
Similarity
Similarity gain
1.00
1.00
GemNet-OC Layers
n2n
AOKA
0.30
0.50
0.25
0.25
-05
0%
PaINN
PaiNN
PaiNN Layers
Layers
Layers
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 / Similarity analyses for other KD strategies and teacher-student configurations are presented in Appendix
Random rattling: Adding noise to existing structures (also known as "rattling") is form of data
augmentation that has been used in the context of pretraining of molecular GNNs 42 45 and as 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 dIM, which consists of one million samples generated by sampling new systems (generated with the OC Datasets codebase } 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 d/M 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 the ground truth DFT samples and the dIM 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 bascline models.
Conclusion
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 v2v); 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 more comprehensive stability analysis (48) One caveat of our approach that though inference times not affected, training times are, albeit not necessarily if pre- trained teacher model is available (see Appendix H.. Finally, is important to recognize that such technologies, while innovative, could be used for potentially harmful purposes, such the simulation or discovery toxic systems, or the development harmful technologies.
Dataset
10

---

Acknowledgments and Disclosure of Funding
F.E.K. financially supported by the Excellence Center at Linköping-Lund in Information Technol- ogy (ELLIIT). D.G. is supported by UK Research and Innovation [UKRI Centre for Doctoral Training in Al for Healthcare grant number EP/SO23283/1J. 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- 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 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- open-source codebase, support and discussions. In particular, Muhammed Shuaibi for providing the COLL dataset in LMDB format. Figures assembled in BioRender.
References
[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 Zitmik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Günnemann, Jure Leskovee, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial intelligence for science 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. B(3)-qquivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, May 2022.
[5] Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. In International Conference on Learning Representations, 2020.
[6] Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-awarr directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurlPS, 2020.
[7] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. 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 Günnemann, Zachary Ward Ulissi, C. Lawrence Zitnick, and Abhishek Das. Ge: 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. 2022.
Advances Neural Information Processing Systems,

---

[10] Julia Westermayr, Michael Gastegger, Kristof T Schütt, and Reinhard J Maurer. Perspective on integrating machine learning into computational chemistry and materials science. The Journal of Chemical Physics, 154(23):230033, 2021.
[11] Patrick Reiser, Marlen Neubert, André Eberhard, Luca Torresi, Chen Zhou, Chen Shao, Hous- sam Metni, Clint van Hoesel, Henrik Schopmans, Timo Sommer, et al. Graph neural networks for materials science and chemistry. Communications Materials, 3(1):93, 2022.
[12] Zehong Zhang, Lifan Chen, Feisheng Zhong, Dingyan Wang, Jiaxin Jiang, Sulin Zhang, Hualiang Jiang, Mingyue Zheng, and Xutong Li. Graph neural network approaches for drug- target interactions. Current Opinion in Structural Biology, 73:10232) 2022.
[13] Jiacheng Xiong, Zhaoping Xiong. Kaixian Chen, Hualiang Jiang, and Mingyue Zheng. Graph neural networks for automated de novo drug design. Drug Diseovery Today, 26(6):1382-1393, 2021.
[14) Mengying Sun, Sendong Zhao, Coryandar Gilvary, Olivier Elemento, Jiayu Zhou, and Fei Wang. Graph convolutional networks for computational drug development and discovery. Briefings in bioinformatics, 21(3):919-935, 2020.
[15] Anuroop Sriram, Abhishek Das, Brandon M Wood, Siddharth Goyal, and C Lawrence Zitnick. Towards training billion parameter graph neural networks for atomic simulations. arXiv preprint MxXiv:2003006997 2022.
[16] Saro Passaro and C. Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pages 27420- 27438. PMLR, 23-29 29 Jul 2023.
[17] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Mor- gane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. The Open Catalyst 2020 (OC20) Dataset and Community Challenges. ACS Catalysis, ((10))6059-6072, May 2021.
[18] Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robers Müller. Schnet: A continuous-files convolutional neural network for modeling quantum interactions, 2017.
[19] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In Proceedings of the 38th International Conference on Machine Learning, pages 9377-9388. PMLR, July 2021.
[20] Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang. Model compression and acceleration for deep neural networks: The principles, progress, and challenges. IEFE Signal Processing Magazine, 35(1):126-136, 2018.
[21] Lin Wang and Kuk-Jin Yoon. Knowledge distillation and student-teacher learning for visual intelligence: A review and new outlooks. IEEE Transactions Pattern Analysis and Machine Intelligence, 2021.
[22] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 108, 2019.
[23] Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao. Knowledge distillation: survey. International Journal of Computer Vision, 119:1789-111 18189918111111 2021.
[24] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network, March 2015.
[25] Chengming Hu, Xuan Li, Dan Liu, Xi Chen, Ju Wang, and Xue Liu. Teacher-student architecture for knowledge learning: A survey. arXiv preprint 20. 2022.
12

---

[26] Qing Xu, Zhenghua Chen, Mohamed Ragab, Chao Wang, Min Wu, and Xiaoli Li. Contrastive adversarial knowledge distillation for deep model compression time-series regression tasks. Neurocomputing, 085:242-251, 2022.
[27] Yijun Tian, Shichao Pei, Xiangliang Zhang, Chuxu Zhang, and Nitesh V. Chawla. Knowledge distillation on graphs: A survey, 2023.
[28] Jing Liu, Tongya Zheng, Guanzheng Zhang, and Qinfen Hao. Graph-based knowledge distilla- tion: A survey and experimental evaluation. arXiv preprint arXiv:2302.14643, 2023.
[29] Muhamad Risqi U Saputra, Pedro PB De Gusmao, Yasin Almalioglu, Andrew Markham, and Niki Trigoni. Distilling knowledge from a deep pose regressor network. In Proceedings of the IEEE/CVF international conference on computer vision, pages 263-272, 2019.
[30] Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4794 4802, 2019.
[31] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets, 2015.
[32] Gustavo Aguilar, Yuan Ling, Yu Zhang, Benjamin Yao, Xing Fan, and Chenlei Guo. Knowledge distillation from internal representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 7350-7357, 2020.
[33] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmicla, Alexandre Tkatchenko, and Robeert Müller. SchNet: A continuous filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[34] Yiding Yang. Jiayan Qiu, Mingli Song, Dacheng Tao, and Xinchao Wang. Distilling knowledge from graph convolutional networks, 2021.
[35] Chaitanya K. Joshi, Fayao Liu, Xu Xun, Jie Lin, and Chuan Sheng Foo. On representation knowledge distillation for graph neural networks. JEEE Transactions on Neural Networks and Learning Systems, pages 1-12, 2022.
[36] Lu Yu, Shichao Pei, Lizhong Ding, Jun Zhou, Longfei Li, Chuxu Zhang, and Xiangliang Zhang. Sail: Self-augmented graph contrastive learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8927-8935, 2022.
[37] Cuiying Huo, Di Jin, Yawen Li, Dongxiao He, Yu-Bin Yang, and Lingfei Wu. T2-gmn: Graph neural networks for graphs with incomplete features and structure via teacher-student distillation. arXiv preprint arXiv:2212. 12738, 2022.
[38] Huarui He, Jie Wang, Zhanqiu Zhang, and Feng Wu. Compressing deep graph neural networks via adversarial knowledge distillation. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 534-544, 2022.
[39] Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, 2019.
[40] Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy, Do vision transformers see like convolutional neural networks? In A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[41] Pengguang Chen, Shu Liu, Hengshuang Zhao, and Jiaya Jia. Distilling knowledge via knowl- edge review. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5008-5017, 2021.
[42] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint ar arXiv:2206.00/33 2022.
13

---

[43] Rishikesh Magar, Yuyang Wang, and Amir Barati Farimani. Crystal twins: Self-supervised learning for crystalline material property prediction. npj Computational Materials, 811):1-8, November 2022.
[44] Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzaler Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint (xXi::106.009777 2021.
[45] Daniel Schwalbe-Koda, Aik Rui Tan, and Rafael Kömme-Bommarel/e Differentiable sampling of molecular geometries with uncertainty-based adversarial attacks. Nature Communications, 12(1):5104, August 2021.
[46] Yu Zhang, James Qin, Daniel 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 and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molccular simulations, 2022.
14

---

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, summarize the KD weighting factors À used for each model configuration in Table g
Table 5: SchNet hyperparameters.
Hyperparameter Hidden channels
OC20
COLL
1024
128
Filters
256
128
Interaction blocks
Gaussians
200
50
Cutoff
6.0
12.0
Batch size
192
32 10-3 3
Initial learning rate Optimizer
10-4
AdamW W LambdaLR
AdamW LinerWarruuuyyoenntaalaaaay
Scheduler
Learning rate decay factor
0.1
0.01
52083, 83333,
Learning rate milestones
104166
Warmup steps
31 31250
3750
Warmup factor
0.1
Force Coefficient
100
100
Energy Coefficient
Number of epochs
30
500
15

---

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 None
Linear WarmpppxpoonennnaDee..
Warmup steps Learning rate decay factor
3750
0.45
0.01
160000, 320000,
Learning rate milestones (steps)
480000, 640000
Force coefficient
100
100
Energy coefficient
1
EMA decay
0.999
0.999
Gradient clip norm threshold
10
10
Epochs
16
375
16

---

Table 7: GemNet-OC hyperparameters. Slash-separated values indicate GemNet-OC versus GemNet-
OC-small hyperparameters.
Hyperparameter
OC20
COLL
No. spherical basis
7
No. radial basis
128
128
No. blocks
4/3
4
Atom embedding size
256/128
128
Edge embedding size
512/256
256
Triplet edge embedding input size Triplet edge embedding output size Quadruplet edge embedding input size Quadruplet edge embedding output size Atom interaction embedding input size Atom interaction embedding output size Radial basis embedding size
64
64
64
64
32
32
32
32
64
64
64
64
16
16
Circular basis embedding size
16
16
Spherical basis embedding size
32
32
No. residual blocks before skip connection
2
2
No. residual blocks after skip connection
2
No. residual blocks after concatenation
No. residual blocks in atom embedding blocks
3
No. atom emhedding output layers
3
3
Cutoff
12.0
12.0
Quadrupiet cutoff
12.0
12.0
Atom edge interaction cutoff
12.0
12.0
Atom interaction cutoff
12.0
12.0
Max interaction neighbors
30
30
Max quadruplet interaction neighbors
8
8
Max atom edge interaction neighbors
20
20
Max atom interaction neighbors
1000
1000
Radial basis function Circular basis function
Gaussian
Gaussian
Spherical harmonics
Spherical Harmonics
Spherical basis function Quadruplet interaction Atom edge interaction
Legendre Outer
Legendre Outer
True
True
True
True
Edge atom interaction
True
True
Atom interaction
True
True
Direct forces
True
True
Activation Optimizer
Silu
Silu
AdamW
AdamW
LinearWarmup
Scheduler
ReduaeL_ROnPlateau
ExponentialDecay
Force coefficient
100
100
Energy coefficient
1
EMA decay
0.999
0.999
Gradient clip norm threshold
10
10
Initial learning rate
5 x 10
10
Epochs
80/9
165
17

---

Table 8: Choice of the weighting factor À of the KD loss for the different teacher-student configura- tions and KD strategies.
Teacher
Student
KD
OC20
COLL
GemNet-OC
PaiNN
Vanilla (1)
1.0
0.2
GemNet-OC
PaiNN
Vanilla (2)
500
100
GemNet-OC
PaiNN
n2n
10000
1000
GemNet-OC
PaiNN
e2n
1000
GemNet-OC
PaiNN
v2v
50000
100
GemNet-OC
GemNet-OC-small GemNet-OOC-sall
Vanilla (7)
0.2
GemNet-OC
Vanilla (2)
10.0
GemNet-OC
GemNet-OC-small
n2n
1000.0
GemNer-OC
GemNet-OO-small
e2e
100000
PaiNN
PaiNN-small
Vanilla (1)
PaiNN
PaiNN-small
Vanilla (2)
200
100
PaiNN
PaiNN-small
n2n
100
100
PaiNN
PaiNN-small
v2v
1000
10000
PaiNN
SchNet
Vanilla (1)
0.1
PaiNN
SchNet
Vanilla (2)
0.1
100
PaiNN
SchNet
n2n
1000
100
18

---

B
Full validation results on OC20
Tables
12
present the extended results on OC20 across the separate S2EF validation sets.
Table 9: Evaluation results on the OC20 S2EF in-distribution validation set.
OC20 S2EF Validation (in-distribution)
Energy MAE meV.
Force MAE
Force
EFwT
Model
meV V
% f
S: PaiNN-small
409
41.6
0.357
0.16
PaiNN
358
38.5
0.390
0.25
Vanilla (1) Vanilla (2)
1366-33.79)
42.6(-32.3%)
0.35(-21.9%%
0.12(-44.4%)
45.7(-132.3%)
0.316(-126.%%
0.11(-55.6%)
n2n
393 (31.2%)
41.7(-2.5%)
0.359 (4.7%)
0.15(-7.8%)
same
v2v
40665.6%
42.1(-16.9%)
0.357(-2.6%)
0.13(-32.6%)
S: GemNettOC-small T: GemNet-OC
292
27.7
0.534
0.90
226
22.5
0.610
1.09
Vanilla (1)
202
27.7(0.0%)
0.535(1.1%)
0.92(1.8%)
Vanilla (2)
283 283 (13.5%)
27.7(-0.4%)
0.535(0.1%)
1.01(18.8%)
n2n
252 (61.1%)
27.5(3.8%)
0.536(1.6%)
1.09 (19.2%)
e2e
285 (10.1%)
26.4 (25.3%)
0.551 (22.5%)
1.01(10.9%)
S: SchNet
1237
62.2
0.214
similar
P:PaiN
358
38.5
0.390
0.25
Vanilla (1)
1139 (10.6%
52.9 (13.1%)
0.2422 (16%)
0(0%)
Vanilla (2)
1140(10.5%)
(9.2(12.7%)
0.241(15.1%)
0(0%)
n2n
1170(7%)
(9(9.3%)
0.235(11.9%
0(0%)
S: PaiNN
358
38.5
0.390
0.25
T: GemNet-OO
226
22.5
0.61
1.89
different
Vanilla (1)
356(1.7%)
38.3(1.1%)
0.392(1.2%)
0.258(0.5%)
Vanilla (2)
357(0.7%)
43.5(-31.4%)
0.334(-25.4%)
0.210(-2.4%)
n2n
271 (66.0%)
37.3(7.5%)
0.408(8.299
0.477 (13.9%)
e2n
(30(2).8%
36.3(14.0%
0.419 (13,4%)
0.371
v2v
369(-8.2%
37.2(8.0%)
0.409(8.9%)
0.217(-2.0%)
19

---

Table 10: Evaluation results on the OC20 S2EF out-of-distribution (adsorbates) validation set.
OC20 S2EF
Validation (out-of-distribbttion (adsorbates)
Energy MAE meV
Force MAE
Force
EFwT
Model
meV/À.
%
S: PaiNN-small
469
47.9
0.334
0.03
PaiNN
437
44.5
0.369
0.043
Vanilla (1)
519(-156.2%
47.4(14.7%)
0.334(-1.2%)
0.003(0%)
Vanilla (2)
477(-23.6%
54.5(-106.2%%
0.297 99%)
0.002(-76.9%)
n2n
443 (80.9%)
47.2 (19.8%)
0.337 (9.3%)
0.003(-10%)
same
v2v
441 (87.5%)
47.9(-1%)
0.337(8.68%)
0.04 (81.6%)
S: GemNet-OCCsmall T: GemNet-OC
325
31.0
0.521
0.190
258
25.2
0.600
0.45
Vanilla (1) Vanilla (2)
319(19
31.0(-0.7%)
0.522(2.1%)
0.19(-0.9%)
309 24 309.24.22 2%
(1.4%)
0.523
0.20(2.7%)
n2n
282 (63.7% %
30.9(1.7%)
0..23(5.8%)
0.22(11.5%)
e2e
315 (14.8%)
29.3 (28.8%)
0.542 (26.3%)
0.23 (16.5%
S: SchNet
1344
58
0.196
similar
T: PaiNN
437
44.5
0.369
0.043
Vanilla (1)
1247(10.7%
64.5(-49.7%)
0.221 (14.7% %
0(0%)
Vanilla (2)
1245 (10.9%)
64.5(-48.2)
0.22(14%)
0,0%)
n2n
1286(6.3%)
65(-51.%%
0.213
0(0%)
S: PaiNN
437
44.5
0.369
0.043
T: GemNet-OC
258
25.2
0.6
0.45
different
Vanilla (1)
424 (7.2%)
44.5(-0.2%)
0.370(0.5%
0.052(2.3%)
Vanilla (2)
408 (15.99)9%
49.2(-24.3%)
0.315(-23.19%
0.036(-1.79)
n2n
321 (64.9%)
43.2(6.9%)
0.387
0.084 (10.0%)
e2n
407 (16.9%%
41.6 (15.3%)
0.498 (12.9%)
0.081(9.3%)
v2v
418
42.0(13%)
0.391(9.9%)
0.058(3.7%)
20

---

Table 11: Evaluation results on the OC20 S2EF out-of-distribution (catalysts) validation set.
OC20 S2EF
Validation (out-of-distribution (catalysts))
Energy MAE
Foree MAE
Force cos
EFwT
Model
meV
meV V
%
S: PaiNN-smal
467
42
0.341
0.13
PaiNN
412
39.2
0.369
0.23
Vanilla (1)
466(4.9%
42.8(-28.4) 44%
0.336(-16.9%%
0.11(-20%)
Vanilla (2)
(52.449%%
45.4(-120.6%)
0.306(-120.87)
0.12(-10%)
n2n
437 (56.3%)
42.0 (1%)
0.343 (8.2%)
0.14 (11.2%
same
v2v
444 (43.4%%
42.4(-12.8%)
0.342(3.6%)
0.12(-12%)
S: GemNet-OC-smaill T: GemNet-OC
335
28.9
0.506
0.85
288
24.0
0.576
1.68
Vanilla (1)
339 3-9(-9.3%
29.0(-1.1%)
0.50%
0.85(-0.4%)
Vanilla (2)
318(35.4%)
28.9(-0.1%)
0.507
1.05 (24.1%
n2n
309 (54.4%)
28.8(2.0%)
0.508
1.02(20.5%%
e2e
(24.61))%
27.6 (26.5%)
0.524 (25.1% ))
(19.
S: SchNet
1205
61.6
0.205
similar
PaiNN
412
39.2
0.369
0.23
Vanilla (1)
1122 (10.6%)
58.7 (12.9%
0.234 (15.9%)
0.01 (4.3%
Vanilla (2)
1122(10.5%)
58.8(12.5%)
0.23
0(0%)
n2n
1159().9%)
59.4(9.8%)
0.225(12.3%)
0(0%)
S: PaiNN
412
39.2
0.369
0.23
T: GC
288
24
0.576
1.68
different
Vanilla (1)
423(-8.8%)
39.1(0.8%)
0.371(1.1%)
0.230(0.0%)
Vanilla (2)
49.5%
43.6(-29%)
0.320(-23.5%)
0.23(0.2%)
n2n
345 (54.0%)
38.5(4.7%)
0.383
0.433 (14%)
e2n
(8.9%)
37.4 (11.9%)
0.395 (12.3%)
0.317(6.0%)
v2v
424(-10.7%)
38.2(6.5%)
0.386 (0.6(88222 (8.2%)
0.187(-3%)
21

---

Table 12: Evaluation results on the OC20 S2EF out-of-distribution (both) validation set.
OC20 S2EF Validation (out-of-distributtio (both))
Energy MAE meV
Force MAE
Force cos
EFwT
Model
meV/At /
%
S: PaiNN-small
610
56.8
0.346
0.02
PaiNN
554
59.2
0.379
0.03
Vanilla (1)
648(-67.6%)
61.1(-179.2%%
0.056(-900.37)
0.02(0%%
Vanilla (2)
592(32.2%)
60.6(-158.3%
0.310(-112.4%%
0.02(0%)
n2n
557 (94.2%
56.0 (33%)
(13.8(14.8%)
0.018(-15.8%
same
v2v
547 (112.6%)
56.5(12.39%
0.352 (17.3%)
0.025 (43.3)
S: GemNet-OC-sman T: GemNet-OC
424
37.4
0.533
0.11
370
31.0
0.606
0.23
Vanilla (1) Vanilla (2)
412(23.1%) 1%)
37.5(-1.8%)
0.534(1.2%
0.11(-2.3%)
401 2%)
37.4(0.2%)
0.535 (2.0%)
0.12(8.5%)
n2n
395 (3)
37.3(1.69)
0.537(44.9)
(8.5))
e2e
(12.22.)%)
35.5 (29.5%)
0.554 (29.0%)
0.13 (19.9%)
S: SchNet
1450
78.4
0.202
0
similar
T: PaiNN
554
59.2
0.379
0.03
Vanilla (1) Vanilla (2)
1350 (11.2%
75.9(13%)
0.227 (14.3%)
0.0025 (1.8% % 0(0%)
(358(10.2%)
75.7 (14.1 %
0.226 (3.8%)
n2n
(396(6.1%)
76.2(11.5%)
0.220(10.49%
0(0%)
S: PaiNN
554
59.2
0.379
0.03
T: GemNet-OC
370
31
0.606
0.23
differenn
Vanilla (1)
558(-2.39) -2.3%)
53.8(19.0%)
0.380(0.6%)
0.031(-0.5%)
Vanilla (2)
511 (23.3%)
323.1 7%)
0.326(-23.2%
0.027(-2.6%)
n2n
448 (57.4%)
52.4(24.11)
(.394(2.4%%
0.056(12.1%)
e2n
59.6%))%
50.1 (32.4%)
0.407 (12.5%)
0.057 (12.6%)
v2v
538(8.5%
50.5(31.0%)
0.402(10.4%)
0.035(..7%)
22

---

C Baseline results on COLL
In Table present the performance and inference throughput the baseline models on COLL. As the systems are much smaller than those in OC20, the throughput is lot larger than the one observed in [ Qualitatively, however, we observe the same, clear trade-off between accuracy and throughput.
Table 13: Evaluation of the performance the four baseline models on the COLL dataset.
Inference Throughput
COLL test set
Samples /
Energy MAE
Force MAE
Force cos
EFwT
Model
GPU sec. f
meV +
mcV +
% t
SchNet
44000
146.5
121.2
0.970
2.75
PaiNN-small
29000
104.0
80.9
0.984
5.4
PaiNN
13000
85.8
64.1
0.988
10.1
GemNet-OC
3520
44.8
38.2
0.994
20.2
23

---

D Data augmentation
We investigated data augmentation as way of distilling knowledge from GemNet-OC into PaiNN on the OC20 dataset.
D.1 Data jittering
To create additional data, added noise to the atomic positions of the training samples and then used the teacher to label the newly derived samples. We tried two different approaches: Random noise, and optimizing the positions using gradient ascent such that the difference between the predictions of the student and teacher was maximized as done in [45]. Denoting the noise as 6, we obtain the noise atomic positions Xs X + ó. Let the student and teacher models be denoted as fs and fr respectively, we then obtained the noise 8' by initializing this as O and
CKD = ( x), f((6,,)))
= 6 + V&LKD.
However, this becomes computationally expensive, it requires additional gradients for LKLD-- Hence, we settled on using a single step, where we fixed the norm of to avoid going too far away from the real structure. We experimented with different norms, with the smallest being 0.1 Á. We compared this to using random directions with fixed norm, and we did not any improvements when using the more computationally expensive gradient ascent approach. both cases, the noise was added to all the samples in the batch.
D.2 Synthetic data
Combined dataset 2M+dIM. We generated 1M synthetic samples by first drawing 100k random adsorbate and catalyst combinations (systems) and then running relaxations with a pre-trained GemNet-OC model. Out of these relaxations with 100 steps on average (200 max), we randomly draw approx. 10% to obtain 1M samples.
In the next step, we combine the 1M samples with the 2M OCP dataset, which based on DFT relaxations. Directly working with this combined dataset means iterating over 1-to-2 ratio of samples from each subset in an epoch of 3M samples. To control this ratio, we define the target ratio of samples from the synthetic dataset during training ttarget E [0. 1]. Setting Otarget = 0.5 means that per epoch we iterate over the 1M dataset 1.5 times and over the 2M dataset 0.75 times.
Different weighting in loss depending on origin. Next to specifying sampling ratio of samples from the synthetic dataset, we can also specify how to weight the contribution of samples to the loss based on their origin (DFT or synthetic). To achieve this, we specify the weighting ratio of synthetic to DFT samples Tsifft was/ R+ In each batch, we compute a weighting factor in front the synthetic and DFT wan samples satisfying the conditions
W3
Obbatch + Wan (1 - (baacc)) = 1,
where Obatch is the ratio of synthetic to DFT samples in batch. Hence, when we combine DFT and synthetic data- Obatch E (0, 1), derive the following weights:
Wiat (1 Obatch
Obbach
(10)
WC = Twan Wan-
(11)
Likewise, when Obbrch {0, 1} i.e., we either train on DFT or synthetic data exclusively, the corresponding weighting coefficient (w, wan) is naturally equal to 1.
24

---

E
Hyperparameter studies
We additionally investigated different aspects of our KD protocols, which we present below. We have performed these experiments when distilling GemNer-OC into PaiNN on the OC20-2M dataset.
E.1 Effect of losses
In our experiments, we have used MSE as the loss Lreat- However, in the general framework in Equation (3), there are other choices that possible, e-g-, more advanced losses like Local Structure Preservation 34] and Global Structure Preservation (GSP) [35]. We therefore initially experimented with using these alternative losses when distilling GemNet-OC into PaiNN on OC20. However, the initial experiments showed that MSE worked well, and in particular, a lot better than the more advanced GSP and LSP losses. We therefore settled on using MSE as our Lieat- In Table [4] we present the average performance over all four validation splits when using these different losses in the n2n and e2n settings.
We also experimented with trying to optimize the CKA directly (as we saw that the CKA similarity improved when using distillation), but did not work and did not pursue any further.
Table 14: Comparing different loss functions Lieat with GemNet-OC teacher and PaiNN as student on OC20, using the n2n and e2n KD protocols. In the case of LSP e2n, the model completely failed when predicting forces on the ood (both) validation set, leading to force MAE of 464 meV /Å, which is almost a factor 10 larger than the other models on the same split. We have therefore written this value as "-" When evaluating checkpoint for model which was trained half as long, the error on this split was 53.7 meV V/A, AA, and the average over all four validation splits was 44.8 me V /A.
OC20 validation set
Energy MAE
Force MAE
Force cos
EFwT % 1
Loss
meV 4
meV / L
MSE
346
42.8
0.393
0.262
n2n
GSP
427
46.1
0.356
0.124
LSP
398
44.8
0.367
0.159
MSE
430
41.3
0.405
0.195
e2n
GSP
463
45.2
0.363
0.119
LSP
441
0.380
0.134
E.2 Effect of transformations
We have evaluated using different transformations My, i.c.. transformations of the student features before applying the loss Lreat. We tried either using the identity function, (i.c., not using a transforma- tion at all), a linear transformation (i.c., multiplication with matrix and adding bias vector), or using a multilayer perceptron (MLP) with one hidden layer. We conducted our experiments when distilling G We found that using an MLP worsened the results, and for e2n, there was not big difference between using a linear layer and no transformation all. For n2n, the node features in PaiNN and GemNet-OC are of different dimensions, and we can therefore not use the identity transformation when using the MSE loss.
The results from the experiments are presented in Table 15
E.3 Effect of feature selection
GemNet-OC consists of an initial embedding layer, followed by series of interaction layers. The result of each embeddingg/nterration layer used as input into the next layer, while copy is also processed by an "output layer". To make the final prediction, the results of the different output layers are concatenated and processed by a final MLP. This means that, for each embeddingginneraation layer, we have two features that could potentially be distilled: the feature used as input for the next layer, or the result of the output layer. Additionally, we could use features from inside the final MI.P which make the prediction by processing the concatenated output features.
25

---

Comparing different transformation functions Ms (Identity, a linear layer and an MLP with one hidden layer) with GemNet-OC as teacher and PaiNN as student, using the n2n and e2n KD protocols. As the node features in GemNet-OC and PaiNN are of different dimensions, we cannot use the identity transformation when using n2n.
Table 15:
OC20 validation set
Energy MAE meV
Force MAE meV V
Force cos
EFwT % t
Loss
t
Identity
n2n
Linear
346
42.8
0.393
0.262
MLP
363
45.1
0.367
0.147
Identity
430
41.3
0.405
0.195
e2n
Linear
418
41.3
0.405
0.207
MLP
427
43.2
0.387
0.161
Initially, used the feature after the final interaction layer when distilling knowledge from GemNet- OC. However, we found that using the feature just before the final linear layer in the final MI.P gave drastic improvement performance. We, therefore, out to investigate how the choice of features impacted the results in more detail.
We performed these experiments when distilling GemNet-OC into PaiNN OC20 using n2n strategy, and the results presented here are on a set of 30 thousand samples sampled from the in- distribution validation set. We did not perform any extensive hyperparameter tuning, but chose such that AC (the distillation loss term) was initially roughly the same for all choices.
present the training feature in PaiNN and varying the choice of features in GemNet-OC. The overall trend is that closer to the output is better: even using the features from the early output layers is better than using features from later interaction layers. Our results suggest that for forces, it is better to use features from earlier output layers. However, think this could be due the choice of X, found that the weighting of the loss term in n2n could offer trade-off between energy and force performance.
Choice of GemNet-OC layer. In Figure
when fixing the choice of
have empirically
800
60
Embeddingt Layer
Guigul Layer
100
55-
Enwegy M55 [mx]]]
600
PANE (nvewvk)
50-
Firal
500
45
400
300
35
200
30
300000
200003
300000
300000
100000
200000
300600
400000
Training
Figure 4: Evaluation error as we vary the features to distill from in the teacher model energy MAE (leff), and force MAE (right). n2n KD from GemNet-OC PaiNN. Performance is evaluated a validation subset comprising 30k samples. The numbers 0 to 4 indicate at what stage the feature has been extracted, with 0 meaning after the embedding layer, and 1 to 4 after the corresponding interaction layer. Solid and dashed lines indicate if the feature is the result of an embedding//nteractinn layer, or an output layer, respectively. "Tinal" refers the feature extracted right before the final linear prediction layer.
26

---

Choice of PaiNN layer. PaiNN consists of sequence of blocks, where each block consists of message layer and an update layer. In Figure we present training curves when fixing the choice of features in GemNet-OC (the feature just before the final linear layer), and varying the choice of features PaiNN (choice of block, and either the feature after the corresponding message or update layer). The results here indicate that using deeper features leads to better results.
Atter uodate
Block
Bleck
Stock
Energy MSS [mx]]]
IorewA1
IC600R
200003
200000
400009
100300
200000 200000
400000
Training step
Trailing step
Figure 5: Evaluation eгrоr as we vary the features distill into in the student model energy MAE (left), and force MAE (right). n2n KD from GemNet-OC to PaiNN. Performance is evaluated validation subset comprising 30k samples. The different colors indicate after which block the features have been extracted, and dashed and solid lines indicate if features were extracted after the message or update layers, respectively.
Conclusion. The conclusion we draw from this study is that using features close to the output as possible improves KD performance in our setup. However, these results are only empirical, and more investigation could be done. For example, it is possible to beforehand determine which pairs of features should be used (and not having rely on trial-and-error).
27

---

Error bars
To get an idea of the stability of our KD protocols, we perform additional experiments distilling GemNet-OC into PaiNN using three different seeds and compute standard deviations. We present these results in Table 16
Table 16: Performance of KD from GemNet-OC into PaiNN across different seeds, averaged over all validation splits. The numbers presented mean + one standard deviation. The missing force error for the baseline model due to one of the seeds completely failing on the out-of-distribution (both) split, drastically increasing the error. The other two seeds had force MAEs of 43.8 and 45.3 meV/ / A, respectively.
OC20 validation set
Energy MAE mcV +
Force MAE
Force cos
EFwT % 1
Loss
mcV/À +
None (baseline)
440 + 8 346 + 0.7
0.376 + 0.0018
0.143 + 0.0051 0.256 + 0.011
n2n
43.2 + 0.6
0.392 + 0.0017
28

---

G Explainability
We utilize CKA similarity scores to monitor the effect of KD throughout our studies, Here, present selection of the analyses we have performed, summarizing how KD influences the similarity between teacher and student models across saacher-studen configurations (Figure 6 across KD protocols (Figure 7): and feature selections (Figure 8 We found out that such similarity metrics can effectively used to examine and profile different KD approaches, as well as as a potential debugging tool. We also explored the utility of CKA (in conjunction with measures of the predictive ability of individual features) as means inform the design of (optimal) KD strategies and feature selection protocols priori, but the results were not conclusive include here.
no KD
with KD
Similarity
Similarity
Similarity gain
10
1.00
1.00
GemNet-OC Layers
0.75
n2n
0.75
0.5
0.50
g
0.50
0.0
ACKA
0.25
-0.5
0,00
0.00
11
PaiNN-big big Layers
PaiNN-big Layers
PaiNN-big Layers
Similarity
Similarity
Similarity gain
1.0
0.8
0.8
PaiNN big Layers
n2n
0.1
0.6
0.6
g
0.0
0.4
0.4
0
-0.2
PaiNN-small Layers
small Layers
PaiNN-small Layers
Similarity
Similarity
Similarity gain
11
1.0
0.02
PaiNN-big Layers
n2n
0.6
ACKA
0
0.00
0.4
0.4
0.2
-0.02
0.0
SchNet
SchNet Layers
Layers
Layers
Figure 6: We explore the effect of n2n KD the feature similarity between different student-teacher configurations: (a) GemNet-OC PaiNN; (b) PaiNN -> PaiNN-small; (c) PaiNN -> SchNet. The layer pair that was used in each experiment is indicated with Note the scale difference in the Similarity gain plots.
29

---

no KD
with KD
Similarity
Similarity
Similarity gain
10
1.00
1.00
GemNettOO Node Layers
0.75
n2n
0.75
g
0.50
g
0.25
0.25
0.00
0.00
11
11
11
PaiNN-big
PaiNN-big Node Layers
PaiNN-big Node Layers
Node Layers
b
Similarity
Similarity
Similarity gain
10
1.00
GemNet-O6 Edge Layers
0.75
e2n
0.75
0.50 g
0.50 g
0.00
0.25
-0.25
0.00
0.00
11
11
PaiNN-big
PaiNN-big
PaiNN-big Node Layers
Node Layers
Node Layers
Figure 7: We explore the effect of different KD strategies n2n and e2n KD on the feature similarity between the student and the teacher models. This is computed for GemNet-OC PaiNN: (a) n2n; (b) e2n. The layer pair that was used in each experiment is indicated with /. Note the scale difference the Similarity gain plots.
30

---

with KD
a
Similarity
Similarity gain
CemMNtt- OC
0.75
GemMettOO Layers
0.2
ACKA
0.50
0.0
0.25
0.00
PaiNN big Layers
PaiNN big Layers
b
Similarity
Similarity gain
1.00
0.5
GemNet-OC Layers
GemNet-OC Layers
0.50
g
0.0
ACKA
0.25
no KD
PaiNN
0.00
PaiMN Layers
Layers
Similarity
C
Similarity
Similarity gain
10
1.00
1.00
GomNet-Or OC
GemNet-Oe OC Layers
GemNet-OC Layers
0.5
0.50 g
0.50 g
0.0
ACKA
0.25
0.00
0.00
big Layers
big Layers
big Layers
d
Similarity
Similarity gain
1.00
GemNet-Oc Layers
0.75
GemNet-Oc Layers
0.25
0.50
ACK
0.25
-0.25
0.00
PaiNN-big
PaiNN big Layers
Layers
Similarity
Similarity gain
GemNett-0 OC
0.75
GemNet:OC Layers
0.25
PaiNN- big Layers
PaiNN--ing -ig
Layers
Figure 8: We explore the effect of feature selection in KD on feature similarity between the student and the teacher models: (a) H4->U6; (b) CONCAT+MLP H4->U5; (c) X2->M2. The layer pair that was used in each experiment is indicated with /. Note the scale difference in the Similarity gain plots.
(c) MLNNCTTMMPP>M-HH (d)
31

---

H
Training times
One caveat of knowledge distillation that it inherently increases the training time of the student model. In our offline KD setup, we need to perform additional forward passes through the teacher to extract representations distill to the student. However, it important note that, despite increasing the computational time per training step, we observed that models trained with KD can outperform their baseline counterparts even when compared at the same training time point (Figure D despite the latter having been trained for more steps/epoch in total. This means that, all in all, we can use KD to enhance the predictive aceuracy in models without necessarily impacting training times.
However, we make the following remark. In this experiment, we utilized publicly available pre-trained Gemnet-OC model weights, and therefore did not have train the teacher model ourselves. However, when access to a pre-trained teacher model is not available, one should also account for the time required train the teacher.
PaiNN (no KD)
1.0
PaiNN (with KD)
0.9
0.8
Energy MAE
0.7
0.6
0.5
0.4
0.3
0
25
50
75
100
125
150
175
Training time (hours)
Figure 9: Energy validation error of PaiNN without (blue) and with (orange) knowledge distillation from GemNet-OC, trained for the same number of steps (1 million). Validation on random sample of size 30k samples from the midisttibuttion OC20 validation set.
32


Reviews:
[{"review_id": "9u7NejjSBt", "reviewer": "Reviewer_o8B6", "summary": "the paper proposes new knowledge distillation strategies to enhance the hidden representation in the molecular GNN, downstream regression task of energy and force predicition shows promising results", "questions": "see above", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. the paper is well written and easy to follow\n2. the experiments show the knowledge distillation can improve the regression performance.", "weaknesses": "1. to handle various structures in the molecule dataset, the paper proposes to use some GNN model to extract hidden features and do feature-based KD on top of that. the novelty is limited.\n2. need to include more benchmark datasets to demonstrate the effectiveness. Also, in table 3, vanilla KD can outperform the proposed method in certain scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. We are glad to hear that you like the presentation of our work and highlight the improvement in performance we achieve! We also appreciate your insightful comments. We carefully respond to your concerns below.\n\n### We have an additional benchmark dataset: COLL (W2.1)\nWe would like to highlight that we have actually run experiments on ***2 distinct datasets*** - ***OC20-2M*** (the biggest and most diverse dataset in the area) and ***COLL*** (a challenging dataset for molecular dynamics). Due to space constraints, we had to put the results on COLL in the appendix, which might have made them easy to miss. We now moved these results into the main paper. ***Please see the discussion in our general response.***\n\n### Vanilla KD is part of our investigation, not a baseline to beat (W2.2)\nWe believe there might have been a misunderstanding about the inclusion of the two vanilla KD approaches in our study due to how we had originally phrased some parts of the text. The two vanilla KD strategies (inspired by the vanilla, logit-based KD used in classification problems (Hinton *et al.* (2015)) ***have not been investigated in this field before***, so analyzing their performance is as much part of our work as are the other strategies. In other words, these are additional strategies we study, and not really baselines we want to compare with and beat necessarily. As such, the fact that they can sometimes outperform n2n, e2n, and v2v does not undermine our contribution - to show that KD is a viable strategy in the context of molecular GNNs and compare different methods.\n\nHowever, we agree that this may not be clear from the phrases we use in the paragraph where we define these two methods. ***To make that more clear in the text, we change***: \n- the caption of the paragraph from *\"Baseline KD strategies\"* to *\"Additional KD strategies\"*; \n- and the introduction of that paragraph from *\"To validate the performance of our KD strategies, we evaluate their performance against 2 vanilla-based KD approaches suitable for regression tasks.\"* to *\"We further evaluate two additional KD approaches inspired by the vanilla logit-based KD used in classification, which we augment to make suitable for regression tasks.\"*\n\n### Novelty of our work (W1)\nThe reviewer does not mention a reason behind their assertion that our work is of limited novelty, so we are not sure what the problem may be. ***We believe our work is a novel contribution to the field because:***\n- *We introduce KD to molecular modelling*: we explore the utility of KD for molecular GNNs for the first time; \n- *Technical novelty*: we adapt feature-based KD to explore custom KD strategies (n2n, e2n, v2v, 2xVanilla) for the distillation of representations in equivariant and directional molecular GNNs. These are indeed built upon the established framework of feature-based KD, but they have been significantly adapted to the requirements of this field, e.g. how to respect the physics of the problem, which features to use, how to perform distillation).\n- *Comprehensive empirical analysis*: we conduct extensive empirical analyses and ablation studies of our framework and its components (features to distill, transformation functions, losses, data augmentation methods). \n\nAs such, it perfectly ***aligns with the definition of novelty outlined in NeurIPS' guidelines***. \n"}, {"review_id": "z7yh1Umxz0", "reviewer": "Reviewer_t1R8", "summary": "The main contribution of this work is the exploration of knowledge distillation as a means to enhance the performance and scalability of GNNs for molecular simulations. The authors introduce custom KD strategies, namely node-to-node, edge-to-node, and vector-to-vector \n distillation, to overcome limitations of KD for regression tasks in GNNs. The performance of the KD protocols is evaluated by training student models to predict molecular properties such as energy and forces. The results show improvement in the performance of student models.\n", "questions": "It would have been interesting seen these methods applied to model with Tensor features like e3nn, which are much heavier to train than the models in the paper. ", "limitations": "The authors discuss the limitation of their approach. ", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and the topic is of practical importance. ", "weaknesses": "See questions. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the concise, positive review! We are glad to hear that they appreciate the presentation of our work and acknowledge its importance!\n\n### Applying KD to tensorial models\nWe also thank the reviewer for their suggestion to apply our methods to tensorial models like e3nn, but we kindly regard that as beyond the scope of this paper. The focus of this work was to conduct an empirical evaluation of the utility of KD in the area, which we show across 3 different architectures that represented the SOTA at the time of experimentation. Recognizing the positive findings of our work, we agree that it would be an interesting extension to also explore how they translate to tensorial models. To carry this message to the reader, ***we have added this as an additional point in our discussion on future work***."}, {"review_id": "6FbmzD8eEU", "reviewer": "Reviewer_xfxb", "summary": "This paper investigates the use of knowledge distillation (KD) to accelerate molecular dynamics using graph neural networks (GNNs) and improve their predictive accuracy. They describe their experiments with various KD protocols such as node to node, edge to node, and vector to vector KD for four different molecular GNN architectures, namely, PaiNN, SchNet, GemNet-OC. They demonstrate that KD can improve the speed with minor compromise on accuracy of molecular GNNs without altering their architecture by studying the OC20 catalysts.", "questions": "1. The choice of the specific architectures is not clear. Although the authors have tried to be extensive by the selection of three different architectures, none of these architectures are giving SOTA performance. Accordingly, the reasons for choosing these architectures are not clear. It would be interesting to try the approaches on SOTA architectures such as equiformer, NequIP, or allegro.\n2. The baselines chosen seem to be not really fair and give mostly negative results which suggests that they are incapable of any knowledge distillation. Fair baselines, should be implemented based on regression tasks and appropriate loss functions. \n3. More extensive studies on different and challenging datasets should be carried out. This is available in most of the papers on machine learned potentials such as 3BPA, AcAc, alanine dipeptide, or LiPS.\n4. Stability of the KD-based model need to be analyzed more carefully to understand the limitations and applicability of the models.", "limitations": "Authors have not included a detailed discussion on the limitations of the study. This should be included.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper demonstrates the approach of knowledge distillation for developing faster GNN-based interatomic potentials for molecular simulations. This is an important problem to be addressed in the field.", "weaknesses": "There are several weaknesses for the paper as outlined below.\n\n1. The authors have selected three different GNNs as the teacher models. However, the performance of the GNNs are not comparable to the state-of-the-art for the same architectures on the same dataset (if I am not mistaken) as confirmed by other publications and OC20 dashboard. \n2. Authors selected only dataset for the evaluation, which is OC20. However, OC20 is not primarily a molecular simulation dataset as I understand. It is a dataset for energy and force prediction primarily based on DFT minimization. Although the task of force and energy prediction is similar molecular simulations, the dynamics task is not tested. Other more complex datasets on dynamics are not included in testing.\n3. Although, this is one of the first attempts for KD in molecular simulations, the improvement in performance of most of the models is only marginal and not substantial. \n4. The stability of these models on molecular simulation is not tested. Further, error evaluation metrics are not exhaustively presented [1].\n\n[1] Fu, X., Wu, Z., Wang, W., Xie, T., Keten, S., Gomez-Bombarelli, R. and Jaakkola, T., 2022. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. arXiv preprint arXiv:2210.07237.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and are glad to hear that they acknowledge the importance of the problem we are tackling! We also appreciate your insightful comments. We carefully respond to your concerns below.\n\n### We have an additional benchmark dataset: COLL (W2+Q3)\nWe wish to highlight that ***we have actually run experiments on 2 distinct datasets, including a molecular dynamics dataset - COLL!*** ***Please see the discussion in our general response.***\n\n### OC20 is a highly relevant molecular simulation dataset (W2+Q3)\nWe are confused about the statement that OC20 is not a molecular simulations dataset - ***it is a dataset concerned with the simulation of catalyst-adsorbate systems***. In this work, we tackle the task of predicting energy and forces, which is the most general task in OC20 and has the broadest applicability across catalysis and related fields (Chanussot *et al.* (2021)). Also, note that performing DFT minimization in high throughput searches is of high interest in the materials science domain, and inference speed is of high importance to enable covering as large a part of chemical space as possible. ***OC20 is the largest dataset for this task***, and we thus think this is a highly relevant dataset.\n\n### The models we study were SOTA at the time of research (W1+Q1)\nWe also want to highlight that ***the models investigated in this paper were SOTA on both OC20 and COLL at the time of research and experimentation***. Looking at the relevant S2EF leaderboard, even today there are only three models better than GemNet-OC: SCN, eSCN and EquiformerV2. However, note that the public weights of SCN and eSCN were only released in February and March this year, in parallel to our work (Aug 2022-May 2023), and the weights for EquiformerV2 were released in July 2023 (after our submission). Hence, GemNet-OC was SOTA at the time of our work (***and even today is still very representative of the SOTA***). The other models we study (PaiNN and SchNet) are chosen as they represent models that are faster, but not SOTA performance. ***To make that more clear to the reader***, we include that information in the text.\n\n### Performance improvements are not marginal (W3)\nWe respectfully disagree with the assertion that most of the improvements we achieve are marginal. Using the KD strategies we study, ***we consistently close >60% of the gap*** between teacher and student models for energy predictions, and 10-20% for forces. Especially on the COLL dataset, which the reviewer might have missed (we hope this is now more apparent after moving these results to the main text), where these numbers go ***up to 96.7% and 62.5% for energy and force prediction, respectively*** (refer to the COLL table in the one-page pdf document attached to our response). \n\nAnd still, we want to reiterate that this ***improvement is out-of-the-box*** with respect to inference throughput, which is the bottleneck we tackle in this paper. In other words, the improvement does not come at the expense of slower models at inference, meaning that even a small improvement is still useful. Moreover, ***this is the first work in the area***, and we hope we can inspire more work to further close the gap.\n\n### Vanilla KD are fair and not arbitrary (Q2)\nWe are not sure why our 2 vanilla KD approaches have come across as unfair to the reviewer.  Our vanilla KD methods are ***not arbitrary***. These are strategies inspired by vanilla KD from classification, adapted to our context and undergone significant optimization as part of our work.We agree that the results we achieve with vanilla KD are not always perfect, but ***they do prove promising in many experiments, often outperforming other techniques***. ***Especially on the COLL dataset***, which the reviewer might have missed (we apologize again and hope this is more apparent now with the results moved to the main text), where looking at force prediction, Vanilla (2) outperforms the other methods on 2 of the 3 teacher-student configurations, and achieves just slightly worse than n2n on the third.\n\nAlso, we want to highlight that there might have been a misunderstanding about the inclusion of the two vanilla KD approaches in our study due to how we had originally phrased some parts of the text. The two vanilla KD strategies ***have not been investigated in this field before***, so analyzing their performance is as much part of our work as are the other strategies. In other words, these are additional strategies we study, and not really baselines we want to compare with or beat necessarily. \n\nHowever, we agree that this may not be clear from the phrases we use in the paragraph where we define these two methods. ***We make that more clear in the text (see general response)***.\n\n### Discussing limitations and future work\nWe agree that our previous discussion on limitations/future work was somewhat narrow, only concerning potential issues around increasing training times. Accordingly, ***we have expanded this section and provided a more comprehensive account of the limitations of our methods and future work***, mentioning future directions like combining KD strategies (e.g. n2n and v2v); extending the framework to other types of features, molecular tasks and datasets; improving the theoretic intuition behind KD; and performing a more comprehensive stability analysis of the approach.\n\n### Stability analysis and error evaluation metrics (W4+Q4)\nWe concur that a stability analysis would be indeed valuable. However, considering that this is a first work on KD for molecular GNNs, we have used well-established metrics, and have categorized stability analysis as beyond the present research. We thank you for the suggestion and have ***included this aspect in our limitations and future work discussions***.\n\nIn terms of error evaluation metrics, we do provide all the metrics that are relevant for the datasets and tasks we benchmark on - OC20 and COLL."}, {"review_id": "VgkpQXOpVn", "reviewer": "Reviewer_n7yW", "summary": "This paper aims to improve the performance of resource-efficient GNNs for molecular simulation, an area where top models are becoming increasingly larger and more cumbersome.\n\nSeveral knowledge distillation approaches are proposed with the aim of regressing the smaller student model's embeddings onto those of a larger teacher model.\n\nThe paper conducts a unified empirical benchmark of various distillation approaches and demonstrates that smaller molecular GNN performance can be non-trivially boosted when trained using KD from a larger teacher model.\n", "questions": "- The Introduction or elsewhere in the Experiments could provide a key quantification of whether/how much the proposed techniques boosted GNNs over training from scratch and baseline KD techniques. Figure 1 gives a nice illustration, but it is unclear which KD technique was used to obtain each of the three results. In each of the cases, was it the baseline or one of yours?\n\n- For line 136, please elaborate a bit on what 'complications' we need to account for? Why can't we just use the scalar component of features to perform distillation in these networks? (As my reading of this paper's results seem to suggest that distilling from vector channels was not actually that useful...)\n\n- For Table 2's EFwT column, an application-oriented question: Yes, we can distill into small and efficient models like SchNet and boost their performance slightly, but is this scientifically relevant? For instance, would you use the SchNet model after your distillation based training for running molecular simulations?\n\n- The formulation of v2v via equations 3 and 6 was a bit worrying to me. If the student's final vector feature is regressed onto the vector feature of the teacher using mean absolute/square error, won't doing this not account for how these quantities are equivariant/geometric vectors in 3D space? \n    - For example, there could be a very high MSE between two rotated versions of the same set of vectors. However, if you found the optimal alignment and then computed MSE, it would be zero.\n    - Thus, comparing two sets of geometric vectors requires first finding an optimal alignment/rotation matrix and then computing the RMSE. Just computing unaligned RMSE may not be very informative. This is standard practice in fields such as protein structure prediction.\n    - Have I misunderstood this aspect of the work?", "limitations": "There is some discussion on limitations and potential negative societal impact. I would have been interested in deeper discussions around both the technical aspects as well as application-specific limits of this line of work.\n\nFor instance, whether GNN model expressivity, especially for geometric and equivariant GNNs, plays a role in the choice of teacher-student pairings? There may be classes of functions that certain student models can provably not learn ever (eg. distinguishing types of neighborhoods or entire geometric graphs). The paper does not propose any principles for selecting teacher-student pairings for molecular simulations. (This is related to my question on whether a distilled SchNet is ever really useful in practice, too.)\n\nAdditionally, for these OCP datasets, could it be interesting to distill across tasks? What may happen if we distill from S2EF models into IS2RE models (a task with significantly lesser data), and is this interesting? Did you try this?", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "## Significant motivation to accelerate GNNs for molecular simulation\n- Improving scalability and efficiency of GNNs specialized for molecular simulations is a worthwhile research question. The best architectures for this area are increasingly becoming very large in terms of compute requirements, so improving the performance of smaller models is certainly worthwhile.\n- I agree that this is the first paper to propose distillation as an approach to boost smaller models - the novelty of the overall contribution also makes this paper interesting.\n\n## Experiments are performed in a fair and unified manner under the same pipeline\n- This is true to the best of my understanding and without having looked at the code.\n- Empirically benchmarking and demonstrating to what extend various KD techniques can boost molecular GNNs is worthwhile. \n\n## Overall clear and well-structured presentation", "weaknesses": "## Limited technical novelty\n- Equation 3 essentially projects teacher and student node features to a common dimension and performs regression to align the embedding spaces. And perhaps this is 'all we need' for doing effective KD for molecular GNNs.\n- However, in my opinion, this technical idea is of limited novelty.\n    - I will caveat this by saying that empirical studies which may not introduce new methodology are very important. I would not reject a paper for this. I think the empirical benchmark is a strength.\n- I don't think anything about Equation 3 is very specific to GNNs for molecular simulation. While this paper may be the first to apply and evaluate these ideas for molecular simulations (which is worthwhile), I don't think the technical ideas are sufficiently novel or tailored to this area.\n\n## Missing contextualization w.r.t. general GNN distillation literature\n- The study is certainly valuable for its focus on molecular simulation applications. However, better contextualization and (possibly) comparison with existing work on general purpose GNN representation knowledge distillation could further improve it.\n- The introduction (line 37) claims that KD is limited for regression tasks, but feature-based KD or representation distillation is a generally applicable technique regardless of what downstream task one is interested in. I currently do not see a good reason to not compare to at least some standard baseline on general purpose representation distillation, or (better) to existing work on GNN representation distillation ([LSP, Yang et al.](https://openaccess.thecvf.com/content_CVPR_2020/papers/Yang_Distilling_Knowledge_From_Graph_Convolutional_Networks_CVPR_2020_paper.pdf), [G-CRD, Joshi et al.](https://arxiv.org/abs/2111.04964)). Joshi et al. could also be cited for introducing CKA-based embedding similarity analysis in the context of GNNs, which this paper's analysis also uses.  \n- Line 87 states that distillation approaches for standard GNN architectures have only concentrated on classification tasks. I do not agree with this claim, since representation distillation approaches are agnostic to the downstream task.\n\n**Overall, I felt that the paper does not give enough credit to existing, more general purpose solutions on GNN distillation in the literature. Yes, the goal of the paper is to be an empirical study specific to molecular simulations and this is a positive (see Strengths). However, the proposed techniques are seemingly not bespoke to molecular simulation, so better contextualization and comparison to existing work is warranted, in my opinion.**", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the detailed feedback! Happy to hear you appreciate the importance and novelty of our work, and the quality of our empirical results.  We also appreciate your insightful comments. We carefully respond to your concerns below.\n\n### Technical novelty of our work and contextualization (W1+W2)\nWe agree that putting our work into context is central and made multiple improvements to address this gap.\n\nWe have ***improved the contextualization*** in our paper based on the discussion points below (see general response for the exact changes).\n\nFollowing the suggestion of the reviewer, we ran ***additional experiments with LSP and GSP (see table in attached pdf)***.\n\nWe added a reference to Joshi et al. (2021) as a paper also using CKA.\n\nWe furthermore want to highlight that:\n1) Eq. 3 represents a general framework for feature-based KD, which we agree is not new on its own. KD research is focused on finding new methods based on that framework-i.e finding appropriate representations (neurons, layers, combinations of layers, relationships); transformations (projections, fusion, fission, etc), and losses (MSE, MAE, LSP, GSP, G-CRD). Thus, not every new method that belongs to the class of feature-based KD is not novel. Actually, GSP and G-CRD methods (Joshi et al. (2021)) mentioned are new types of KD loss, which also fall under the framework defined by Eq 3. \n2) Here, we do not create a new KD loss but tackle the question of what type of features to distill in molecular GNNs (a separate, unique degree of freedom in our context). We define and study 5 strategies for distilling different features across molecular GNNs. They are indeed built upon the overarching framework of feature-based KD, but adapted to the requirements of this field, e.g. how to respect the physics of the problem, which types of features to use, how to transform equivariant and directional representations. \n3) As such, they are not directly comparable to methods like LSP, and G-CRD. Actually, our strategies can be regarded as general frameworks that can use LSP, GSP, G-CRD as a KD loss. In our study, we decided to use MSE loss as we found it more appropriate for our setup. Contrastive learning (the backbone of LSP, GSP and G-CRD) is useful for node classification where you want to contrast different classes of nodes, but less clear how to use for graph-level regression tasks on molecules (i.e. atoms are a part of the whole). We had already experimented with GSP as a loss, available in the appendix where we perform ablation studies of layers, transformations and losses. MSE seemed to work significantly better.\n4) We agree that as a general approach, feature-based KD is mostly agnostic to the downstream task. Yet, ***prior work on applying feature-based KD to regression task is quite limited. Most feature distillation methods are tailored towards classification tasks*** - e.g. still use logit information (e.g. LSP, GSP, and G-CRD).\n\n### Definition of v2v KD (Q4)\nWe believe this is a misunderstanding. Note that the vectorial features of *both* the student and the teacher are equivariant by construction, and PaiNN uses this to make equivariant force predictions. It is hence desirable that also the directions of the vectorial features are correct, and rotating the vectors before applying the loss would not encourage this. It is therefore ***important to use a loss that encourages vectors to align in both magnitude and direction***. To make that more clear in the text, ***we add the following in connection to Eq. 6***: \"Note that as these features are equivariant to rotations, it is important to use a loss that encourages vectors to align in both magnitude and direction - e.g. MSE.\"\n\n### Experimental results are available in Table 3 (Q1)\nWe are puzzled by the comment that we should: \"provide a key quantification of whether/how much the proposed techniques boosted GNNs over training from scratch and baseline KD techniques.\". We provide detailed experimental results in Table 3 (and Table 12 for COLL), which give all the details mentioned in this comment. Could you clarify what you think is missing?\n\nWe indeed summarize these in Fig 1, where we depict the gap between student and teacher models. We are happy the reviewer found this figure useful. We agree that the caption of Figure 1 was not clear enough. Now it explicitly says that percentages represent our best results for each teacher-student configuration depicted.\n\n### Clarifying line 136 (Q2)\nWe apologize for any confusion here. ***We have made that more explicit:*** by adding *\"as one needs to extract and align representations corresponding to comparable features in both models.\"*\n\n### SchNet is scientifically relevant (Q3)\nWe see the point with respect to OC20. But it is likely relevant for cases with a well-sampled chemical space. Also, simple, well-established methods protect us from overfitting to highly-engineered models. We decided to use SchNet as it is still one of the most widely known models in the field.\n\n### Discussing limitations\nWe thank the reviewer for the very helpful suggestions as potential limitations. \n- ***\"defining teacher-student pairings based on their (provable) expressivity\"*** - understanding the theory behind KD is a very active area of research (how it works, when it works, etc.). It is definitely an interesting future research avenue, so ***we added that to future work***.\n- ***\"distillation across OCP tasks\"*** - This is not something we have considered, but we think you have a point that this extra data could be used. However, we don't envision this within a **KD framework**, but rather a **pretraining framework**. The reason for this is that the two tasks are very different, and we think it is unclear how the knowledge of a teacher should be distilled into a student.\n\nWe have ***expanded our discussion into limitations/future work*** and added the aforementioned points, as well as other suggestions from the reviewers."}, {"review_id": "qvCj8OZ8ff", "reviewer": "Reviewer_6ZfK", "summary": "* This paper explores knowledge distillation (KD) on speeding up molecular GNNs, which has challenges of being regression instead of classification, and having both scalar and vector outputs\n* They tried a few different methods on a few different teacher-student combinations, with the best results closing 65% of the gap between the student and teacher in energy, and 21% in forces, slightly less for different student-teacher models.\n* Analyzed performance with respect to similarity of models, \n* Tried data augmentation techniques, but with no improvements", "questions": "* Why are forces much worse than energy improvement? Any analysis into this?\n* Error bars? I’m not sure how much of Table 3 is noise\n* I am unsure how much it matters about how similar the teacher and student model are. Isn’t it often that you use the same architecture for the student, just smaller?\n* Why does vanilla KD do significantly worse on PaiNN-big to PaiNN-small? This doesn’t seem intuitive.\n* Data augmentation: how much data did you add? Is the amount of rattling reasonable (e.g. too much movement may be OOD for the teacher)\n* Suggestion for baseline: how good are your KD techniques when the student model is the exact same as the teacher?", "limitations": "The authors acknowledged limitations and potential negative impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This paper tackles an important problem that is an unstudied area: applying KD to GNNs of 3D space with regressions\n* Tackles challenges of KD: regression instead of classification, vector outputs in addition to scalar\n* Paper is clearly written, with good related background to my knowledge\n* Despite the weaknesses described below, I believe this is an important area for KD to extend into, and this paper makes good initial progress at it.\n", "weaknesses": "* Despite good improvement for the small model for energy, it’s still quite a ways from the performance of the large model. Force metrics are especially bad, considering force MAE is almost 2x even with KD for some models.\n* Beyond that, it is hard to get a sense of how good these improvements are. How does it compare to the improvement of KD in other fields? One suggestion is to plot speed vs. accuracy to see the tradeoff, which may help determine the usefulness in downstream applications (e.g. hypothetically, a small model is twice as fast but 10% less accurate, but perhaps you can make up for that by running a bit less than 2x the number of inference calls, so then that’s considered “good”).\n* I would like to see other teacher-student combinations - for example, why 2 sizes of PaiNN but not other models? I am curious about big gemnet on small gemnet, since gemnet is the better model in this selection.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We want to thank the reviewer for their comments and feedback! We are extremely happy to hear they appreciate the importance of our work, the unique challenges we address in the paper, and the quality of the presentation. We also appreciate your insightful comments. We carefully respond to your concerns below.\n\n### Performance improvements are substantial (W1)\nWe respectfully disagree with the assertion that the improvements not being substantial. Using the KD strategies we study, ***we consistently close >60% of the gap*** between teacher and student models for energy predictions, and 10-20% for forces. Especially on the COLL dataset (which we now have moved to the main text) where these numbers go ***up to 96.7% and 62.5% for energy and force prediction, respectively*** (see the COLL table in the attached pdf). \n\nAnd still, we want to reiterate that this ***improvement is out-of-the-box*** with respect to inference throughput, which is the bottleneck we tackle in this paper. In other words, the improvement does not come at the expense of slower models at inference, meaning that even a small improvement could be useful. And moreover, ***this is the first work in the area***, and we hope we can inspire further work to further close the gap.\n\n#### Force predictions (W1+Q1)\nGiven the amount of training data on energies (one scalar per system) vs forces (one vector per atom), we think that learning to predict energies is a harder problem. We hypothesize that by distilling knowledge from a teacher model \"knowing\" more about the energy surface of a molecule, ***there is more potential to improve the energy predictions***.\n\nBecause of the tradeoff between energy and force accuracy, people often develop 2 separate models optimized for energy and force predictions, tuning different hyperparameters (Gasteiger *et al.* (2021)). Here, we do not perform any such optimization specifically. But, we still achieve force improvements up to 62.5%, which combined with our energy prediction gains, constitute a good overall improvement for our initial exploration in the area. Future work can focus on better optimizing for force predictions. \n\n### Putting performance improvements into context (W2)\nWe believe this is hard to quantify as previous research is mostly on classification, where percentage gains are normally smaller. To put the performance into some context, closing almost 60% of the gap in energy predictions between GemNet-OC and PaiNN on OC20 gives an energy error lower than that of GemNet-dT (which is severely slower), trained on the same OC20-2M dataset (Gasteiger *et al.* (2022)).\n\nWe thank the reviewer for the suggestion to plot speed vs. accuracy. We created such a plot, but due to the different scales, the performance was less visible and we instead plotted in the format of Figure 1, which conveys the same information. We also have that information in Table 2, which summarises the tradeoff between the models. We now edited Figure 1 to include the results on the COLL dataset and benchmarked the models on COLL, which we added to the appendix. \n\n### Importance of model similarity (Q3)\nWe can see two aspects in the question:\n1) Our results do not in the end show that model similarity is important for KD.\n2) You were already under the impression that model similarity shouldn't be a necessity for successful KD before our analysis.\n\nRegarding 1), we think that this is a contribution of our work; we initially thought that it would be easier to distill between more similar models, but the results suggest otherwise.\n\nAbout 2), we were not of this impression before seeing the final results, and although model similarity doesn't seem to be a requirement for knowledge distillation, our results (e.g. Figure 3) indeed indicate that model similarity increases with KD.\n\n### Vanilla KD on PaiNN-big to PaiNN-small (Q4)\nNote that Vanilla KD appears to not do well only on OC20 between PaiNN-big to PaiNN-small. We believe that is because the gap between the force prediction accuracy of the 2 baseline models used in this scenario is already rather small, meaning the extra signal we provide during Vanilla KD is not as significant. \n        \n### Data augmentation (Q5)\nIn the case of random rattling, we added noise to all samples in the batch. To avoid adding too much noise (and going OOD), we experimented with adding noise of a fixed norm, and we tried different values starting from 0.1 Å. We have added this information in the appendix.\n\nFor the synthetic data, we mixed it with the real data, with the fraction of synthetic data in a single batch being \\{0%, 5%, 10%, 20%, 50%, 80%\\} on average.\n\n### Error bars (Q2)\nWe thank the reviewer for the comment. We have run additional runs for GemNet-OC -> PaiNN on the OC20 dataset. A table can be found in the attached pdf, also added to the appendix and we add a reference in the results section. The force results for the baseline model are missing, as one of the runs completely failed on the ood_both task. The force MAEs for the baseline were 45.3, 43.7, and 115.7 meV/Å.\n\n### Other teacher-student combinations (W3)\nWe agree that distilling into a smaller GemNet-OC model would be an interesting experiment. However, given the limited time, we have not been able to finalize these experiments. Preliminary, it looks like energy predictions see a substantial improvement when using the n2n loss.\n\n### Baseline suggestion (Q6)\nWe thank the reviewer for the suggestion for another baseline, but we are not sure what the reviewer means here by *\"exact same\"* - distilling to the exact same model configuration - e.g. PaiNN-big to PaiNN-big. That would be considered teacher-free KD which we haven't explored, but seems to be an established and well-performing method [1].\n\n[1]Yuan et al, Revisiting Knowledge Distillation via Label Smoothing Regularization, CVPR 2020"}]

Author rebuttal:
We wish to take this opportunity to thank the reviewers for their time and effort in assessing our work. We are deeply grateful for the reviews we have received. We were extremely happy to hear that the reviewers recognized the importance of our work and the quality of our results and presentation!

We also appreciate the reviewers' insightful comments and feedback, which significantly contributed to the improvement of our presentation. We have taken the time to carefully respond to all of their concerns separately. Below, we include some of the central points of the reviews, which we believe are important to clarify.

### Additional benchmark dataset: COLL
We would like to highlight that we have actually run experiments on ***2 distinct datasets*** - ***OC20-2M*** (the biggest and most diverse dataset in the area) and ***COLL*** (a challenging dataset for molecular dynamics). We apologize if our ***results on COLL***, originally presented in Table 12 in the appendix due to space constraints, escaped the reviewer's attention. Recognizing the importance of such additional validation, ***we have decided to move Table 12 to the main text*** and extend our discussion of the results on the COLL dataset. We have ***also revisited Figure 1 to include results on COLL*** to make the presence of an additional benchmark dataset more apparent. We have also included the table summarizing our results on COLL in the one-page pdf document attached to our response. Note that we had previously had an error in the percentage calculations for GemNet-OC -> PaiNN-big on COLL, resulting in significantly smaller percentage improvements than those actually achieved. We have fixed that and updated the table.

### Clarification on Vanilla KD strategies
We believe there might have been a misunderstanding about the inclusion of the two vanilla KD approaches in our study due to how we had originally phrased some parts of the text. The two vanilla KD strategies (inspired by the vanilla, logit-based KD used in classification problems (Hinton *et al.* (2015)) ***have not been investigated in this field before***, so analyzing their performance is as much part of our work as are the other strategies. In other words, these are additional strategies we study, and not really baselines we want to compare with and beat necessarily. As such, the fact that they can sometimes outperform n2n, e2n, and v2v does not undermine our general line of work - to demonstrate that KD is a viable strategy in the context of molecular GNNs.

However, we agree that this may not be clear from the phrases we use in the paragraph where we define these two methods. ***To make that more clear in the text, we change***: 
- the caption of the paragraph from *"Baseline KD strategies"* to *"Additional KD strategies"*; 
- and the introduction of that paragraph from *"To validate the performance of our KD strategies, we evaluate their performance against 2 vanilla-based KD approaches suitable for regression tasks."* to *"We further evaluate two additional KD approaches inspired by the vanilla logit-based KD used in classification, which we augment to make suitable for regression tasks.*

### Other changes to the paper:
- ***Additional experiments (see one-page pdf)***
    - Additional seeds for GemNet-OC -> PainNN-big (see response to reviewer 6ZfK)
    - Additional losses, LSP and GSP (see response to reviewer n7yW)
    - Additional results GemNet-OC -> GemNet-OC-small (currently training - see response to reviewer 6ZfK)
- ***We have revisited parts of our methods section to improve the contextualization*** of the proposed techniques (see response to reviewer n7yW). More notably, we have added the following to the definition of *n2n*: *"Note this is a general approach that utilizes node features only, making it applicable to standard GNNs. Here, we want to enforce the student to mimic the representations of the teacher for each node (i.e. atom) independently, so we use a loss that directly penalizes the distance between the features in the two models, such as MSE (similar to the original formulation of feature-based KD in Romero et al. (2014)). Other recently proposed losses L_feat for the distillation of node features in standard GNNs specifically include approaches based on contrastive learning (Yang et al. (2021), Joshi et al. (2021), Yu et al. (2022) Huo et al. (2022)) and adversarial training (He et al. (2022)). We do not focus on such methods as much since they are better suited for (node) classification tasks (e.g. contrasting different classes of nodes), and not for molecule-level predictions."*
- ***We have expanded our discussion on limitations and future work*** to accommodate the suggestions of the reviewers, mentioning future directions like combining KD strategies (e.g. n2n and v2v); extending the framework to other types of features, molecular tasks, and datasets; improving the theoretic intuition behind KD and its applicability; and performing a more comprehensive stability analysis of the approach."

Meta-review: Throughout the reviews, most of the reviewers agree that this submission contributed a dedicated empirical study of a range of methods on the distillation for molecular force field models which provides evidence for future research, while there are also concerns on e.g., technical novelty and statement on related work, absolute error and usability of the distilled model, and significance on the choice of the models and evaluation datasets/metrics. In follow-up discussions with the reviewers, they finally tend to accept this work considering that the empirical study on distillation methods on molecular tasks could be a worthy contribution. I hence recommend an accept. Nevertheless, the authors should be aware of and try to address the concerns that the limited number of datasets and not using SOTA models (e.g., GemNet (https://arxiv.org/abs/2106.08903), NequIP (https://www.nature.com/articles/s41467-022-29939-5)) weaken the solidity of the empirical study, and that the intention of distillation is not clearly described. For the latter point, the authors should respond to "if the intention is a lighter-weighted model to run molecular dynamics, modeling chemical reactions, or finding stable structures, then the method should be evaluated on these workloads, for which the sacrificed accuracy may or may not corrupt the stability". Ultimately, the method is for an application in the science domain but not a purely machine learning task. The authors should respect that in the paper.
Decision: Accept (poster)