Title: Accelerating Molecular Graph Neural Networks via Knowledge Distillation
Venue: NeurIPS 2023
Accepted: yes
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

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, stateof-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.

1 Introduction

In the last couple of years, the field of molecular simulations has undergone a rapid paradigm shift with the advent of new, powerful computational tools based on machine learning (ML) [1]. 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 [2, 3]. Consequently, 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) at a fraction of the computational cost [4, 5, 6, 7, 8, 9]. This has, in turn, significantly accelerated the modeling of molecular properties and the simulation of molecular systems, bolstering new research developments in many scientific disciplines, including material sciences, drug discovery and catalysis [10, 11, 12, 13, 14].

Nonetheless, this progress - largely coinciding with the development of bigger and more complex models, has naturally come at the expense of increased complexity [15, 9, 7, 16]. This has gradually

Dimitar Georgiev ∗

Johannes Gasteiger

OC20-2M

COLL

GemNet-OC

PaiNN

PaiNN

GemNet-OC

PaiNN

PaiNN

PaiNN-small

PaiNN-small

107

KD

264

264

KD

680

3.5K

KD

13K

13K

KD

29K

samples/sec

samples/sec

samples/sec

samples/sec

samples/sec

samples/sec

samples/sec

samples/sec

Energy

Force

Energy

Force

Energy

Force

Energy

Force

(baseline)

Student

MAE

440.2

20.5%

MAE

45.3

MAE

489.2

20.5%

MAE

47.1

MAE

85.8

t'tg %

MAE

61.2

MAE

104.0

MAE

80.9

KD

60.8%

41.3

64.8%

46.7

62.0%

62.5%

mev

meV/A

meV

meV/A

meV

meV/A

96.7%

mev

meV/A

(with KD)

Student

346.1

457.4

60.4

70.4

86.4

Teacher

285.5

25.7

440.2

45.3

44.8

38.2

85.8

64.1

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 a 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 area.

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 molecular GNNs. In 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 - a large-scale, 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-tonode (e2n) and vector-to-vector (v2v) knowledge distillation, which facilitate the distillation of hidden representations in 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 of student models while fully preserving their throughput (see Figure 1 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 .

2 Background

Molecular simulations. In this work, we consider molecular systems at an atomic level, i.e., N atoms represented by their atomic numbers z = { z 1 , ..., z N } ∈ Z N and positions X = { x 1 , . . . , x N } ∈ R N × 3 . Given a system, we want a model that can predict the energy E ∈ R of the system, and the forces F ∈ R N × 3 acting on each atom. Both of these properties are of high interest when simulating molecular systems. The energy of a 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 , z ) can be represented as a mathematical graph G = ( V , E ) , where the nodes V correspond to the set of atoms, and edges E are created between nodes by connecting the closest neighboring atoms (typically defined by a cutoff radius and/or a 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 a multi-output manner - i.e., ˆ E, ˆ F = GNN ( X , z ) .

The main problem when modeling molecules and molecular properties is 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., is 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 symmetries and other physical constraints. There is now a plethora of diverse molecular GNNs that reflect that, e.g., SchNet [18], DimeNet [5, 6], PaiNN [19], GemNet [7, 8], NequIP [4], and SCN [9], which have incrementally established a more holistic description of molecular systems by capturing advanced geometric features and physical symmetries. This has, however, come at 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 [22]. The main objective of KD is to 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, 25]. 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 state-of-the-art molecular GNN architectures containing complex, geometric node- and edge-level features, is still unexplored.

3 Knowledge distillation for molecular GNNs

Preliminaries. In the context of the aforementioned prediction task, we train molecular GNNs by enforcing a loss L 0 that combines both the energy and force prediction error as follows:

where E and F are the ground-truth energy and forces, ˆ E and ˆ F are the predictions of the model of interest, and L E and L F are some loss functions weighted by α E , α F ∈ R .

To perform knowledge distillation, we augment this training process by defining an auxiliary KD loss term L KD, which is added to L 0 (with a factor λ ∈ R ) to derive the final training loss function L :

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 a given (teacher) model carry valuable information that can complement the ground-truth 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 a KD loss L KD that forces the student to mimic the predictions of the teacher model. This is usually achieved by constructing a loss L KD = KL ( z s , z t ) based on the Kullback-Leibler (KL) divergence between the soft logits of the student z s and the teacher z t .

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 [20, 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 is substantial, and thus detrimental to KD performance [30].

Feature-based KD. To circumvent these shortcomings, we focus on feature-based KD - an extension of vanilla KD concerned with the distillation of knowledge across the intermediate layers of models

[31, 32, 23]. In particular, we perform knowledge distillation of intermediate representations by devising a loss on selected hidden features H s ∈ U s and H t ∈ U t in the student and teacher models respectively, which takes the form:

where M s : U s ↦→ U and M t : U t ↦→ U are transformations that map the hidden features to a common feature space U , and L feat : U × U ↦→ R + is some loss of choice. The transformations M s , M t can be the identity transformation, linear projections and multilayer perceptron (MLP) projection heads; whereas for the distillation loss L feat, 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 s and H t ). 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, 28], 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 in 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 a complex molecular graph. These are continually evolved by model-specific operators to infer molecular properties, such as energy and forces, in a 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 to extract and align representations corresponding to comparable features in both models.

In this work, we set out to devise KD strategies that are 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 in 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 s ∈ R d . These are used to predict the energy ˆ E . The force is then calculated as the negative gradient of the energy with respect to the atomic positions, i.e., ˆ F = -∇ ˆ E .

PaiNN [19]: A GNN based on equivariant message passing, which contains scalar node features x ∈ R d 1 - used for energy prediction; as well as geometric vectorial node features v ∈ R 3 × d 2 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 [8]: A GNN model that utilizes directional message passing between scalar node features h ∈ R d h and scalar edges features m ∈ R d m . After each block of layers, these are processed through an output block, resulting in scalar node features x ( i ) E and edge features x ( i ) F , where i is the block number. The output features from each block are aggregated into output features x E and x 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 1.

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 H node, we can distill knowledge in between these directly by defining a loss L KD given by

where H node , s and H node , t 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 feature-based KD in Romero et al. [31]). Other recently proposed losses L feat 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-level 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 on its edge features, which are a key component of the directional message passing employed in the architecture. As such, they can be a useful resource for KD. Hence, we also consider KD between edge features, which we accomplish by applying Equation (4) to the edge features H edge , s and H edge , t 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 H edge , ( i,j ) by first aggregating them as follows:

where i is some node index. The resulting features H edge2node ,i are scalar, node-level features, and we can, therefore, use them to transfer knowledge to the student node features H node , s as in Equation (4).

- vector-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 a 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 no such vectorial features, we transfer knowledge between (invariant) scalar edge features and (equivariant) vectorial node features by noting that scalar edge features sit on an equivariant 3D grid since they are associated with an edge between two atoms in 3D space. Hence, we can aggregate the edge features { H edge , ( i,j ) } j ∈N corresponding to a given node i into node-level equivariant vectorial features H vec ,i by considering the unit vector u ij = 1 | x j -x i | ( x j -x i ) that defines the direction of the edge ( i, j ) , such that

with the superscript k indicating the channel. Notice that the features H ( k ) vec ,i fulfill the condition of equivariance with respect to rotations as each vector u ij is equivariant to rotations, and H ( k ) edge , ( i,j ) - a scalar not influencing its direction. Consequently, it is important to use a loss L feat that encourages vectors to align in both magnitude and direction - e.g., MSE.

Additional KD strategies. 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:

- Vanilla (1): One way of adapting vanilla KD for regression is by steering the student to mimic the final output of the teacher directly:

where the subscripts s and t refer to the predictions of the student and teacher, respectively. Note that, unlike in classification, this approach does not provide much additional information in regression tasks, except for some limited signal about the error distribution of the teacher model [20, 29].

- Vanilla (2): One way to enhance the teacher signal during training is to consider the fact that many GNNs for molecular simulations make separate atom- and edge-level predictions, which are consequently aggregated into a final output. For instance, the total energy E of a system is usually defined as the sum of the predicted contributions from each atom ˆ E = ∑ i ˆ E i . Hence, we can extend the aforementioned vanilla KD approach by imposing a loss on these granular predictions instead:

These individual energy contributions are not part of the labeled data, but, when injected during training, can provide more fine-grained information than the aggregated prediction.

4 Experimental results

To evaluate our proposed methods, we perform comprehensive benchmarking experiments on the OC20-2M [17] dataset (structure to energy and forces (S2EF) task) - a large and diverse catalyst dataset; and COLL [6] - a challenging molecular dynamics dataset. We use the model implementations provided in the Open Catalyst Project (OCP) codebase 3 (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, we 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, we also experiment with more lightweight versions of the two architectures - referred to as PaiNN-small and GemNet-OCsmall respectively, where we 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-trained model available within the OCP repository (July 2022).

We present our benchmarking results on OC20 S2EF in Table 2, which summarizes the performance of the five models with respect to the following four metrics: energy and force MAE (i.e., the mean absolute error between ground truth and predicted energies and forces); force cos (i.e., the cosine similarity between ground truth and predicted forces); and energy and forces within threshold (EFwT) - i.e., the percentage of systems whose predicted energies and forces are within a specified threshold from the ground truth [17]. Since force cos and EFwT are correlated with energy and force MAE, we focus on the latter throughout the paper but present all four for completeness.

Our results highlight the substantial trade-off between predictive accuracy and computational cost across the GNN architectures, and, therefore, the need for methods that can alleviate this limitation. We observe the same trend on the COLL dataset (see Appendix C).

Similarity analysis of baseline models. To make our analysis exhaustive, we set out to design experiments involving teacher and student architectures of a variable degree of architectural disparity. As a proxy of that, we derive similarity scores based on central kernel alignment (CKA) [39, 40, 35]. In particular, we calculate the pairwise CKA similarity between the node features of our trained SchNet, PaiNN and GemNet-OC models. The results of this analysis are summarized in Figure 2. Focusing on intramodel similarities first (plots on the diagonal), we observe that, while representations from different layers within PaiNN and SchNet have a generally high degree of similarity, GemNet-OC exhibits the opposite behavior, with features extracted at each layer being significantly different from those captured across the rest of the architecture. This is consistent with the architectures of these three models, with features in PaiNN and SchNet being iteratively updated by adding

SchNet

Layers

5

1.0

0

0.8

10

PaiNN

Layers

0.6

Similarity

5

0

0.4

10

GemNet-OC

Layers

0.2

5

0

0.0

0

5

0

110

10

SchNet

PaiNN

GemNet-OC

Layers

Layers

Layers

displacement features computed at each layer, while those in GemNet-OC representing separate output features. When examining inter-model similarity instead, we notice that, generally speaking, node features in SchNet and PaiNN are similar, whereas those between SchNet and GemNet-OC, and PaiNN and GemNet-OC, diverge significantly as we move deeper into GemNet-OC.

Knowledge distillation results. Based on the aforementioned analyses, we define the following teacher-student pairs, covering the whole spectrum of architectural disparity: PaiNN to PaiNN-small ( same architecture); PaiNN to SchNet ( similar architectures); GemNet-OC to PaiNN ( different architectures). We additionally explore KD from GemNet-OC to GemNet-OC-small ( same architecture) on OC20. We train student models by utilizing an offline KD strategy [23], where we distill knowledge from the more competent, pre-trained teacher model to the simpler, more lightweight student model during the training of the latter. We augment the training of each student model with our KD protocols and evaluate the effect on predictive accuracy against the models trained without KD. If not mentioned otherwise, we utilize the following setup: we use MSE as a distillation loss L feat ; a learned linear layer as a transformation function M s applied to the features of the student; and the identity transformation as M t . When distilling knowledge from/into GemNet-OC models, we use the aggregated node- and edge-level output features, which is reminiscent of the review-based setup proposed in [41]. For PaiNN and SchNet, we use the final node features.

The results of our experiments are summarized in Tables 3 and 4, presenting a comparative analysis of the predictive performance of different student models trained with and without the implementation of knowledge distillation on the OC20-2M and COLL datasets, respectively. Focusing on energy predictions first, we observe that, by utilizing KD, we achieve significant improvements in performance in virtually all teacher-student configurations. In particular, we manage to close the gap in performance between student and teacher models by ∼ 60% or more in six out of the seven configurations, reaching results as high as 96 . 7% (distilling PaiNN to PaiNN-small on the COLL dataset). Putting our results into context, we remark that our PaiNN model trained with n2n KD from GemNet-OC, for instance, provides more accurate energy predictions than more advanced models such as GemNet-dT [7, 8] which is substantially slower. The only setup where results are not as definite is when training SchNet on OC20-2M with KD from PaiNN, where we close 10 . 8% . However, it is noteworthy to highlight that this still corresponds to a significant absolute improvement (i.e., notice the big initial difference in performance between the two baseline models on this dataset).

We similarly observe an improvement in the accuracy of student models in force predictions in all teacher-student configurations. Although we observe force improvements (typically ∼ 5 -25% ) that are generally not as pronounced as those achieved in energy prediction, we note that we reach

results as high as 62 . 5% for some configurations (distilling PaiNN to PaiNN-small on COLL). A possible reason for this difference in improvement between energy and forces could be attributed to the nature of the supervised task - there are substantially more force labels (i.e., one 3D vector per atom, which could be hundreds per sample) than energy labels (i.e., one per sample). Consequently, we hypothesize it is easier for models to learn to make accurate force predictions, and, therefore, there is more room for improvement in the energy predictions, which we can target with KD.

All in all, our experiments demonstrate that, by applying knowledge distillation, we successfully enhance the performance of student models across all teacher-student configurations and datasets, confirming the effectiveness and robustness of the approach in the context of molecular GNNs. We reiterate the fact that no modifications to the student architectures are made, meaning we achieve an out-of-the-box boost in accuracy without impacting inference throughput.

Effect of KD on model similarity. We continue our investigation by exploring how KD affects student models and their features. To this end, we analyze how the CKA similarity between teacher and student models changes with the introduction of KD during training. We present the outcome of one of our analyses in Figure 3, where we summarize how the similarity between the node features of GemNet-OC and PaiNN changes with the implementation of n2n KD. We observe that KD introduces strong and specific similarity gains in the layers we use for KD, which also propagates along the

student architecture. We notice similar behavior across other teacher-student configurations and KD strategies (see Figures 6 and 7 in Appendix G), allowing us to monitor and quantify the effect of KD on student models as we explore different KD settings and design choices.

Hyperparameter studies. We additionally conduct a thorough hyperparameter study to evaluate the effect of different design choices within our KD framework. We summarize our results below.

- Effect of distillation loss : Apart from our default MSE-based distillation loss L feat, we also experimented with more advanced losses such as Local Structure Preservation (LSP) [34] and Global Structure Preservation (GSP) [35], as well as directly optimizing CKA. We observed the best results with our default MSE loss, with other options substantially hurting accuracy (see Appendix E.1).

- Effect of transformation function : We also investigated a number of different transformation functions and the effect they have on performance. The transformations we utilized include: the identity transformation (when appropriate); learned linear transformations, and MLP projection heads. Our results showed that our default linear mapping is a sufficiently flexible choice as it gives the best results(see Appendix E.2).

- Effect of feature selection : We additionally explored the effect of feature selection on KD performance. In particular, we analyzed the change in the predictive accuracy of PaiNN as a student model as we distill features from earlier layers in the teacher (GemNet-OC), or distill knowledge into earlier layers in the student. Our results suggest that using features closer to the output is the best-performing strategy (see Appendix E.3). We also performed CKA-based similarity analyses to monitor how model similarity changes as we vary the features we used for KD (see Figure 8 in Appendix G).

Data augmentation. As for most other applications, data labeling for molecular data is costly as it requires running computationally expensive quantum mechanical calculations to obtain ground truth energies and forces. Motivated by this, we explore two data augmentation techniques, which we use to generate new data points that we label with the teacher and use for KD. We briefly describe these below (see Appendix D for more details).

no KD

with KD

Similarity

Similarity

Similarity gain

10

1.00

10

1.00

10

GemNet-OC

0.5

0.75

n2n

0.75

Layers

CKA

CKA

△CKA

0.50

0.50

0.0

0.25

0.25

- 0.5

0

0.00

0

0.00

0

0

11

0

11

0

11

PaiNN

PaiNN

PaiNN

Layers

Layers

Layers

- Random rattling: Adding noise to existing structures (also known as 'rattling') is a form of data augmentation that has been used in the context of pretraining of molecular GNNs [42, 43], and as a regularization strategy [44]. Inspired by this, we utilized 'rattling' in the context of KD, where we added random noise to the atomic positions of systems and used the teacher to derive energy and force labels for these rattled samples. We then combined these rattled structures with the original dataset during the training of student models. However, this approach did not provide significant improvements. Additionally, we tried using gradient ascent to find perturbations that maximize the discrepancy between the teacher and student predictions, similar to [45], but this did not show improvements over random noise, and also increased training time.

- Synthetic Data: Samples in OC20 S2EF originate from the same relaxation trajectory and are therefore correlated. To tackle this, we generated our own distilled dataset coined d1M , which consists of one million samples generated by sampling new systems (generated with the OC Datasets codebase 4 ), running relaxations with our pre-trained GemNet-OC model, and then subsampling approximately 10% of the frames. We explored different ways of incorporating this new d1M dataset, all of which were based on joint training with the OC20 S2EF 2M data (similar to what we did with the rattled systems). To study different combinations of the ground truth DFT samples and the d1M samples during training, we defined two hyperparameters determining: (a) how many of the samples per batch originate from each of the datasets; and (b) how to weight the loss contributions based on the origin of data. Unfortunately, and contrary to similar approaches, e.g., in speech recognition [46], the results we observed did not significantly improve on the baseline models.

5 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 a more comprehensive stability analysis [48]. One caveat of our approach is that even though inference times are not affected, training times are, albeit not necessarily if a pretrained teacher model is available (see Appendix H). Finally, it is important to recognize that such technologies, while innovative, could be used for potentially harmful purposes, such as the simulation or discovery of toxic systems, or the development of harmful technologies.

Acknowledgments and Disclosure of Funding

F.E.K. is financially supported by the Excellence Center at Linköping-Lund in Information Technology (ELLIIT). D.G. is supported by UK Research and Innovation [UKRI Centre for Doctoral Training in AI for Healthcare grant number EP/S023283/1]. Computing resources provided by: the Berzelius resource at the National Supercomputer Centre, provided by Knut and Alice Wallenberg Foundation; the Alvis resource provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) at Chalmers e-Commons at Chalmers (C3SE) partially funded by the Swedish Research Council through grant agreement no. 2022-06725; the Chair of Aerodynamics and Fluid Mechanics at Technical University of Munich. This research project was initially conceived at the 2022 LOGML summer school, and we would like to thank Guocheng Qian and I-Ju Chen for their contribution during the early conceptualizing stages of this project during and in the first weeks following the summer school. We also thank the Open Catalyst team for their open-source codebase, support and discussions. In particular, Muhammed Shuaibi for providing the COLL dataset in LMDB format. Figures assembled in BioRender.

References

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.

Yuyang Wang, Zijie Li, and Amir Barati Farimani. Graph neural networks for molecules, 2023.

Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang, Haiyang Yu, YuQing Xie, Xiang Fu, Alex Strasser, Shenglong Xu, Yi Liu, Yuanqi Du, Alexandra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung, Minkai Xu, Chaitanya K. Joshi, Simon V. Mathis, Kamyar Azizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Günnemann, Jure Leskovec, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial intelligence for science in quantum, atomistic, and continuum systems, 2023.

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications , 13(1):2453, May 2022.

Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. In International Conference on Learning Representations , 2020.

Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurIPS , 2020.

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.

Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C. Lawrence Zitnick, and Abhishek Das. GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets. Transactions on Machine Learning Research , October 2022.

Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems , 35:8054-8067, 2022.

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):230903, 2021.

Patrick Reiser, Marlen Neubert, André Eberhard, Luca Torresi, Chen Zhou, Chen Shao, Houssam Metni, Clint van Hoesel, Henrik Schopmans, Timo Sommer, et al. Graph neural networks for materials science and chemistry. Communications Materials , 3(1):93, 2022.

Zehong Zhang, Lifan Chen, Feisheng Zhong, Dingyan Wang, Jiaxin Jiang, Sulin Zhang, Hualiang Jiang, Mingyue Zheng, and Xutong Li. Graph neural network approaches for drugtarget interactions. Current Opinion in Structural Biology , 73:102327, 2022.

Jiacheng Xiong, Zhaoping Xiong, Kaixian Chen, Hualiang Jiang, and Mingyue Zheng. Graph neural networks for automated de novo drug design. Drug Discovery Today , 26(6):1382-1393, 2021.

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.

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 arXiv:2203.09697 , 2022.

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 Jul 2023.

Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane 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 , 11(10):6059-6072, May 2021.

Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions, 2017.

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.

Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang. Model compression and acceleration for deep neural networks: The principles, progress, and challenges. IEEE Signal Processing Magazine , 35(1):126-136, 2018.

Lin Wang and Kuk-Jin Yoon. Knowledge distillation and student-teacher learning for visual intelligence: A review and new outlooks. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2021.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 , 2019.

Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision , 129:1789-1819, 2021.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network, March 2015.

Chengming Hu, Xuan Li, Dan Liu, Xi Chen, Ju Wang, and Xue Liu. Teacher-student architecture for knowledge learning: A survey. arXiv preprint arXiv:2210.17332 , 2022.

Qing Xu, Zhenghua Chen, Mohamed Ragab, Chao Wang, Min Wu, and Xiaoli Li. Contrastive adversarial knowledge distillation for deep model compression in time-series regression tasks. Neurocomputing , 485:242-251, 2022.

Yijun Tian, Shichao Pei, Xiangliang Zhang, Chuxu Zhang, and Nitesh V. Chawla. Knowledge distillation on graphs: A survey, 2023.

Jing Liu, Tongya Zheng, Guanzheng Zhang, and Qinfen Hao. Graph-based knowledge distillation: A survey and experimental evaluation. arXiv preprint arXiv:2302.14643 , 2023.

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.

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.

Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets, 2015.

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.

Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert 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.

Yiding Yang, Jiayan Qiu, Mingli Song, Dacheng Tao, and Xinchao Wang. Distilling knowledge from graph convolutional networks, 2021.

Chaitanya K. Joshi, Fayao Liu, Xu Xun, Jie Lin, and Chuan Sheng Foo. On representation knowledge distillation for graph neural networks. IEEE Transactions on Neural Networks and Learning Systems , pages 1-12, 2022.

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.

Cuiying Huo, Di Jin, Yawen Li, Dongxiao He, Yu-Bin Yang, and Lingfei Wu. T2-gnn: Graph neural networks for graphs with incomplete features and structure via teacher-student distillation. arXiv preprint arXiv:2212.12738 , 2022.

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.

Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, 2019.

Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021.

Pengguang Chen, Shu Liu, Hengshuang Zhao, and Jiaya Jia. Distilling knowledge via knowledge review. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 5008-5017, 2021.

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 arXiv:2206.00133 , 2022.

Rishikesh Magar, Yuyang Wang, and Amir Barati Farimani. Crystal twins: Self-supervised learning for crystalline material property prediction. npj Computational Materials , 8(1):1-8, November 2022.

Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veliˇ ckovi´ c, James Kirkpatrick, and Peter Battaglia. Simple gnn regularisation for 3d molecular property prediction & beyond. arXiv preprint arXiv:2106.07971 , 2021.

Daniel Schwalbe-Koda, Aik Rui Tan, and Rafael Gómez-Bombarelli. Differentiable sampling of molecular geometries with uncertainty-based adversarial attacks. Nature Communications , 12(1):5104, August 2021.

Yu Zhang, James Qin, Daniel S Park, Wei Han, Chung-Cheng Chiu, Ruoming Pang, Quoc V Le, and Yonghui Wu. Pushing the limits of semi-supervised learning for automatic speech recognition. arXiv preprint arXiv:2010.10504 , 2020.

Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks, 2022.

Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarelli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations, 2022.

A Training and hyperparameters

Models were trained on NVIDIA A100 40 GB and NVIDIA RTX A6000 48 GB GPUs, except GemNet-OC-small which were trained on NVIDIA A100 80 GB and NVIDIA RTX A6000 48 GB. All models were trained on single GPUs, except for SchNet when trained on OC20-2M, which required 3 GPUs. Inference throughput was profiled on A100 40 GB GPUs, with reported values representing approximate numbers averaged across three evaluations. We provide detailed information about the hyperparameters we used for each model in Tables 5, 6, and 7.

Moreover, we summarize the KD weighting factors λ we used for each model configuration in Table 8.

B Full validation results on OC20

Tables 9-12 present the extended results on OC20 across the 4 separate S2EF validation sets.

C Baseline results on COLL

In Table 13, we present the performance and inference throughput of the baseline models on COLL. As the systems are much smaller than those in OC20, the throughput is a lot larger than the one observed in Table 2. Qualitatively, however, we observe the same, clear trade-off between accuracy and throughput.

D Data augmentation

We investigated data augmentation as a way of distilling knowledge from GemNet-OC into PaiNN on the OC20 dataset.

D.1 Data jittering

To create additional data, we 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 δ , we obtain the noise atomic positions as X δ = X + δ . Let the student and teacher models be denoted as f s and f t respectively, we then obtained the noise δ ′ by initializing this as 0 and

However, this becomes computationally expensive, as it requires additional gradients for ∇ δ L KD. 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 a fixed norm, and we did not see any improvements when using the more computationally expensive gradient ascent approach. In both cases, the noise was added to all the samples in the batch.

D.2 Synthetic data

Combined dataset 2M+d1M. 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 is based on DFT relaxations. Directly working with this combined dataset means iterating over a 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 α target ∈ [0 , 1] . Setting α target = 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 a 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 r s/dft = w s /w dft ∈ R + . In each batch, we compute a weighting factor in front of the synthetic w s and DFT w dft samples satisfying the conditions

where α batch is the ratio of synthetic to DFT samples in a batch. Hence, when we combine DFT and synthetic data α batch ∈ (0 , 1) , we derive the following weights:

Likewise, when α batch = { 0 , 1 } - i.e., we either train on DFT or synthetic data exclusively, the corresponding weighting coefficient ( w s or w dft) is naturally equal to 1 .

E Hyperparameter studies

We additionally investigated different aspects of our KD protocols, which we present below. We have performed these experiments when distilling GemNet-OC into PaiNN on the OC20-2M dataset.

E.1 Effect of losses

In our experiments, we have used MSE as the loss L feat . However, in the general framework in Equation (3), there are other choices that are 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 L feat . In Table 14, 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 it did not work and we did not pursue it any further.

Table 14: Comparing different loss functions L feat with GemNet-OC as 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 a 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 a checkpoint for a model which was trained half as long, the error on this split was 53 . 7 meV / Å, and the average over all four validation splits was 44 . 8 meV / Å.

E.2 Effect of transformations

We have evaluated using different transformations M s, i.e., transformations of the student features before applying the loss L feat. We tried either using the identity function, (i.e., not using a transformation at all), a linear transformation (i.e., multiplication with matrix and adding a 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 a big difference between using a linear layer and no transformation at 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 a series of interaction layers. The result of each embedding/interaction layer is used as input into the next layer, while a 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 embedding/interaction 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 MLP which make the prediction by processing the concatenated output features.

OC20 validation set

Initially, we used the feature after the final interaction layer when distilling knowledge from GemNetOC. However, we found that using the feature just before the final linear layer in the final MLP gave a drastic improvement in performance. We, therefore, set out to investigate how the choice of features impacted the results in more detail.

We performed these experiments when distilling GemNet-OC into PaiNN on OC20 using our n2n strategy, and the results presented here are on a set of 30 thousand samples sampled from the indistribution validation set. We did not perform any extensive hyperparameter tuning, but chose λ such that λ L KD (the distillation loss term) was initially roughly the same for all choices.

Choice of GemNet-OC layer. In Figure 4, we present the training curves when fixing the choice of 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, we think this could be due to the choice of λ , as we have empirically found that the weighting of the loss term in n2n could offer a trade-off between energy and force performance.

800

60

Embedding/lnteraction Layer

Output Layer

0

700

55

1

2

3

600

50

4

MAE [meV]

[meV/A]

Final

500

MAE

45

Energy

Force

400

40

300

35

200

30

0

100000

200000

300000

400000

0

100000

200000

300000

400000

Training step

Training step

Choice of PaiNN layer. PaiNN consists of a sequence of blocks, where each block consists of a message layer and an update layer. In Figure 5, 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 in 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.

600

60

After message block

After update block

550

Block 1

55

Block 2

Block3

500

Block 4

Block 5

50

MAE [meV]

450

[meV/A]

Block6

400

Force MAE

45

Energy

350

40

000

35

250

200

30

0

100000

200000

300000

400000

0

100000

200000

300000

400000

Training step

Training step

Conclusion. The conclusion we draw from this study is that using features as 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, if it is possible to beforehand determine which pairs of features should be used (and not having to rely on trial-and-error).

F 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 3 different seeds, averaged over all validation splits. The numbers are presented as mean ± one standard deviation. The missing force error for the baseline model is 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 / Å, respectively.

G Explainability

We utilize CKA similarity scores to monitor the effect of KD throughout our studies. Here, we present a selection of the analyses we have performed, summarizing how KD influences the similarity between teacher and student models across teacher-student configurations (Figure 6); across KD protocols (Figure 7); and feature selections (Figure 8). We found out that such similarity metrics can be 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 a means to inform the design of (optimal) KD strategies and feature selection protocols a priori , but the results were not conclusive to include here.

no KD

with KD

a

Similarity

Similarity

Similarity gain

10

1.00

1.00

GemNet-OC

0.5

0.75

n2n

0.75

Layers

△CKA

0.50

0.50

0.0

0.25

0.25

-0.5

0

0.00

0.00

11

0

11

0

11

PaiNN-big

PaiNN-big

PaiNN-big

Layers

Layers

Layers

b

Similarity

Similarity

Similarity gain

11

1.0

T

1.0

0.2

0.8

0.8

PaiNN-big

Layers

0.6

CKA

n2n

0.6

CKA

0.1

0.0

△CKA

0.4

0.4

-0.1

0.2

0.2

0

0.0

0.0

-0.2

0

7

0

7

0

7

PaiNN-small

PaiNN-small

PaiNN-small

Layers

Layers

Layers

C

Similarity

Similarity

Similarity gain

11

1.0

1.0

0.02

0.8

0.8

PaiNN-big

Layers

0.6

CKA

n2n

0.6

CKA

△CKA

0.00

0.4

0.4

0.2

0.2

-0.02

0

0.0

0.0

0

5

0

5

0

5

SchNet

SchNet

SchNet

Layers

Layers

Layers

no KD

with KD

a

Similarity

Similarity

Similarity gain

10

1.00

1.00

GemNet-OC

0.5

0.75

n2n

0.75

0.50

CKA

0.50

0.0

△CKA

0.25

0.25

-0.5

0

0.00

0.00

11

0

11

0

11

PaiNN-big

PaiNN-big

PaiNN-big

Similarity

Similarity

Similarity gain

10

1.00

1.00

GemNet-OC

EdgeLayers

0.75

e2n

0.75

0.25

△CKA

0.50

0.50

CKA

0.00

0.25

0.25

-0.25

0.00

0.00

0

11

0

11

0

11

PaiNN-big

PaiNN-big

PaiNN-big

with KD

a

Similarity

Similarity gain

1.00

GemNet-OC

Layers

0.75

GemNet-OC

Layers

0.2

0.0

△CKA

0.25

-0.2

0.00

PaiNN-big

PaiNN-big

Layers

Layers

Similarity

Similarity gain

1.00

GemNet-OC

Layers

0.75

GemNet-OC

Layers

0.5

△CKA

0.50

CKA

0.0

0.25

-0.5

no KD

PaiNN-big

0.00

PaiNN-big

Layers

Layers

Similarity

C

Similarity

Similarity gain

GemNet-OC

Layers

10

0.75

1.00

GemNet-OC

Layers

T1.00

0.75

GemNet-OC

Layers

0.5

0.50

CKA

0.50

CKA

0.0

△CKA

0.25

0.25

-0.5

0

0.00

0.00

PaiNN-big

PaiNN-big

PaiNN-big

Layers

Layers

Layers

Similarity

Similaritygain

1.00

GemNet-OC

Layers

0.75

GemNet-OC

Layers

0.25

△CKA

0.50

0.00

0.25

0.25

PaiNN-big

0.00

PaiNN-big

Layers

Layers

Similarity

Similarity gain

11.00

GemNet-OC

Layers

0.75

GemNet-OC

Layers

0.05

△CKA

0.50

CKA

0.00

0.25

-0.05

0.00

0

11

0

11

PaiNN-big

PaiNN-big

Layers

Layers

H Training times

One caveat of knowledge distillation is 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 to distill to the student. However, it is important to 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 9), 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 accuracy 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 to 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 to train the teacher.

PaiNN (no KD)

1.0

PaiNN (with KD)

0.9

0.8

MAE

0.7

Energy

0.6

0.5

0.4

0.3

0

25

50

75

100

125

150

175

Training time (hours)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\n* 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\n\nSTRENGTHS:\n* 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\nWEAKNESSES:\n* 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\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe authors acknowledged limitations and potential negative impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n## 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\n\nWEAKNESSES:\n## 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.**\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThere 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?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nAuthors have not included a detailed discussion on the limitations of the study. This should be included.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

comments: 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 
 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.

STRENGTHS:
The paper is well written and the topic is of practical importance.

WEAKNESSES:
See questions.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11

comments: 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

STRENGTHS:
1. the paper is well written and easy to follow
2. 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.
2. need to include more benchmark datasets to demonstrate the effectiveness. Also, in table 3, vanilla KD can outperform the proposed method in certain scenarios.

QUESTIONS:
see above

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-15