Accelerating Molecular Graph Neural Networks via Knowledge Distillation
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.
Filip Ekström Kelvinius, Dimitar Georgiev, Artur Petrov Toshev, Johannes Gasteiger
# Accelerating Molecular Graph Neural Networks via Knowledge Distillation

###### Abstract

Recent advances in graph neural networks (GNNs) have allowed molecular simulations with accuracy on par with conventional gold-standard methods at a fraction of the computational cost. 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, for the first time, 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 demonstrate that they can boost the predictive accuracy of student models without altering their architecture. We also 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 as much as 59% of the gap in predictive accuracy between models like GemNet-OC and PaiNN with zero additional cost at inference.

## 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; 2]. 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; 4]. Consequently, a multitude of GNNs have been developed, which now offer predictive performance on par with conventional gold-standard methods like density functional theory (DFT) at a fraction of the computational cost at inference time [5; 6; 7; 8]. This has, in turn, significantly accelerated the modeling of molecular properties and the simulation of diverse molecular systems.

Nonetheless, this progress - largely coinciding with the development of bigger and more complex models, has naturally come at the expense of increased complexity [9; 10]. This has gradually limited the utility of state-of-the-art GNNs for large-scale molecular simulation applications, where inference throughput (i.e., how many samples can be processed for a given time) is critical for making fast continual predictions about the evolution of a system. Therefore, 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, for the first time, investigate the potential of knowledge distillation (KD) in enhancing the performance and scalability of state-of-the-art GNNs for molecular simulations (Fig. 1). To this end, we devise custom strategies for KD on molecular GNNs, which we call _node-to-node (n2n)_, _edge-to-node (e2n)_ and _vector-to-vector (v2v)_ knowledge distillation. These overcome common limitations of KD for regression tasks by facilitating the distillation of hidden representations in directional and equivariant GNNs. We evaluate the performance of our KD protocols in augmenting the training process of different student models trained to predict molecular properties like energy and forces, and demonstrate the effectiveness of our protocols across teacher-student configurations, datasets and KD settings. We show that they can substantially improve the performance of student models without modifying their architecture, allowing us to reduce the performance gap between models like GemNet-OC [11] and PaiNN [12] by as much as \(59.16\%\) in energy predictions and \(16.7\%\) in force predictions, while fully preserving the throughput of the simpler student model.

In summary, the contributions of this paper are as follows:

* a large-scale, multi-output regression task, challenging to address with most current KD methods.
* We design custom KD strategies, which facilitate the distillation of hidden representations in directional and equivariant GNNs. We demonstrate the effectiveness of our protocols across diverse teacher-student configurations and datasets.
* We conduct a comprehensive analysis of different components of our KD strategies, as well as explore data augmentation techniques for further improving performance.

## 2 Background and related work

**Molecular simulations.** In this work, we consider molecular systems at an atomic level, i.e., \(N\) atoms represented by their atomic number \(\mathbf{z}=\{z_{1},...,z_{N}\}\in\mathbb{Z}^{N}\) and position \(\mathbf{X}=\{\mathbf{x}_{1},\ldots,\mathbf{x}_{N}\}\in\mathbb{R}^{N\times 3}\). Given a system, we want a model that can predict the energy \(E\in\mathbb{R}\) of the system, and the forces \(\mathbf{F}\in\mathbb{R}^{N\times 3}\) acting on each atom. Both 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 \((\mathbf{X},\mathbf{z})\) can be represented as a mathematical graph, where the set of atoms corresponds to the nodes \(\mathcal{V}\), and edges \(\mathcal{E}\) 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

Figure 1: Using knowledge distillation, we manage to significantly boost the predictive accuracy of different student models without altering their architecture. Hence, we substantially lessen the tradeoff between speed and performance in molecular GNNs, allowing us to run faster molecular simulations while maintaining predictive accuracy.
context of molecular simulations, we can create a GNN that operates on atomic graphs \(G=(\mathcal{V},\mathcal{E})\) by propagating information between the atoms and the edges, and makes predictions about the energy and forces of each system in a multi-output manner - i.e., \(\hat{E},\hat{\mathbf{F}}=\text{GNN}(\mathbf{X},\mathbf{z})\).

The main problem when modeling molecules and molecular properties are the number of underlying symmetries to consider. 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 \(\mathbf{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. There is now a plethora of diverse molecular GNNs that achieve that, e.g., SchNet [13], DimeNet [6; 14], PaiNN [12], GemNet [7; 11], NequIP [5], and SCN [10].

**Knowledge distillation.** Knowledge distillation is a technique for compressing and accelerating ML models [15], which has recently demonstrated significant potential in domains like computer vision [16] and natural language modeling [17]. 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 [18]. Since the seminal work of Hinton _et al._[19], the field has drastically expanded methodologically, with the development of protocols that accommodate the distillation of "deeper" knowledge, more comprehensive transformation and fusion functions, as well as more robust distillation losses [18; 20]. Yet, these advances have mostly focused on classification, resulting in methods of limited utility in regression tasks. Moreover, most research in the area has been confined to grid-structured 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 [21]. And, in particular, the application of KD to state-of-the-art, directional and equivariant GNN architectures, where features can include complex geometric information and span both node- and edge-level features, as well as to real-world regression problems in molecular simulations, is still unexplored.

## 3 Knowledge distillation in molecular GNNs

**Vanilla KD.** The standard loss function when training molecular GNNs is a loss that combines both the energy and force prediction error as follows:

\[\mathcal{L}_{0}=\alpha_{\text{E}}\mathcal{L}_{\text{E}}(\hat{E},E)+\alpha_{ \text{F}}\mathcal{L}_{\text{F}}(\hat{\mathbf{F}},\mathbf{F}), \tag{1}\]

where \(E\) and \(\mathbf{F}\) are the ground-truth energy and forces, \(\hat{E}\) and \(\hat{\mathbf{F}}\) are the predictions of the model of interest, and \(\mathcal{L}_{\text{E}}\) and \(\mathcal{L}_{\text{F}}\) are some loss functions weighted by \(\alpha_{\text{E}},\alpha_{\text{F}}\in\mathbb{R}\).

In KD, we augment this training process by defining an auxiliary knowledge distillation loss term \(\mathcal{L}_{\text{KD}}\), which is added to \(\mathcal{L}_{0}\) (with a factor \(\lambda\in\mathbb{R}\)) to derive a new training loss function \(\mathcal{L}\) of the form:

\[\mathcal{L}=\mathcal{L}_{0}+\lambda\mathcal{L}_{\text{KD}}. \tag{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 a given (teacher) model carry valuable information that can complement the ground-truth labels in the training process of another (student) model [19]. Since then, this has become the standard KD approach - usually 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 \(\mathcal{L}_{\text{KD}}\) that enforces the student to mimic the predictions of the teacher model. This is commonly achieved by constructing a loss \(\mathcal{L}_{\text{KD}}=\text{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, this strategy - based on the distillation of the output of the teacher model only [18] - poses two significant limitations. Firstly, it is by design exclusively applicable to classification tasks, since there are no outputs analogous to logits in regression setups [15; 22]. This has consequently limited the utility of most KD methods for regression tasks. Secondly, this approach forces the student to emulate the final output of the teacher directly, which is often unattainable in regimes where the complexity gap between the two models is substantial, and thus detrimental to KD performance [23].

**Feature-based KD.** To circumvent these shortcomings, we focus on feature-based KD [18]. This is an extension of vanilla KD [24], which is concerned with the distillation of knowledge across the
intermediate layers of models. This allows more lightweight models to be trained to mimic features that are easier to assimilate compared to the final output directly [25]. In this paper, we perform knowledge distillation of intermediate representations by devising a loss on selected hidden features \(H_{\mathrm{s}}\in U_{\mathrm{s}}\) and \(H_{\mathrm{t}}\in U_{\mathrm{t}}\) in the student and teacher models respectively, which takes the form:

\[\mathcal{L}_{\text{KD}}=\mathcal{L}_{\text{feat}}(\mathcal{M}_{\mathrm{s}}(H_{ \mathrm{s}}),\mathcal{M}_{\mathrm{t}}(H_{\mathrm{t}})), \tag{3}\]

where \(\mathcal{M}_{\mathrm{s}}:U_{\mathrm{s}}\mapsto U\) and \(\mathcal{M}_{\mathrm{t}}:U_{\mathrm{t}}\mapsto U\) are transformations that map the hidden features to a common feature space \(U\), and \(\mathcal{L}_{\text{feat}}:U\times U\mapsto\mathbb{R}^{+}\) is some loss of choice. Possible options for the transformations \(\mathcal{M}_{\mathrm{s}},\mathcal{M}_{\mathrm{t}}\) include the identity transformation, linear projections and multilayer perceptron (MLP) projection heads; whereas for the distillation loss \(\mathcal{L}_{\text{feat}}\), typical functions are 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 appropriate representations (i.e., features \(H_{s}\) and \(H_{t}\)). One needs to ensure that paired features are similar both in modeling capacity 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) features in single-output classification tasks [26, 27, 28, 29], thereby reducing the problem to the selection of layers to pair across the student and the teacher. This is often further simplified by utilizing models of the same architecture.

**Feature-based KD in molecular GNNs.** Most GNNs for molecular simulations contain diverse features (scalars, vectors and/or equivariant higher-order tensors based on spherical harmonics) organized across nodes and edges within a complex graph topology. 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.

In this work, we set out to devise KD strategies that are representative and effective across various molecular GNNs. We consider GNNs that have diverse architectures and performance profiles, and can be organized in teacher-student configurations at different levels of architectural disparity. This is why we investigate the effectiveness of KD with respect to the following three GNN models:

* _SchNet_[30]: A simple GNN model based on continuous-filter convolutional layers, which only contains scalar node features \(\mathbf{s}\in\mathbb{R}^{d}\). These are used to predict the energy \(\hat{E}\). The force is then calculated as the negative gradient of the energy with respect to the atomic positions, i.e., \(\hat{\mathbf{F}}=-\nabla\hat{E}\).
* used for energy prediction; as well as geometric vectorial node features \(\mathbf{v}\in\mathbb{R}^{3\times 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_[11]: A GNN model that utilizes directional message passing between scalar node features \(\mathbf{h}\in\mathbb{R}^{d_{\mathrm{h}}}\) and scalar edges features \(\mathbf{m}\in\mathbb{R}^{d_{\mathrm{m}}}\). After each block of layers, these are processed through an output block, resulting in scalar node features \(\mathbf{x}_{\mathrm{E}}^{(i)}\) and edge features \(\mathbf{x}_{\mathrm{F}}^{(i)}\), where \(i\) is the block number. The output features from each block are aggregated into output features \(\mathbf{x}_{E}\) and \(\mathbf{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 distillation strategies.** As the three models considered in this work contain very different kinds of features, there is room for considering different ways of distilling the information

\begin{table}
\begin{tabular}{l c c} \multicolumn{1}{c}{SchNet} & \multicolumn{1}{c}{PaiNN} & \multicolumn{1}{c}{GemNet-OC} \\ \hline Scalar node features & ✓ & ✓ & ✓ \\ Scalar edge features & & ✓ & \\ Vectorial node features & ✓ & & \\ Output blocks & & & ✓ \\ \end{tabular}
\end{table}
Table 1: Molecular GNNs can have diverse features depending on their architecture. This is an overview of the types of features available in the three models we use in this study.
contained in the respective features between model architectures. In the context of the three model architectures we have considered and their types of features, we devise three different KD strategies:

_node-to-node (n2n):_ As all three models considered in this study contain scalar node features \(H_{\text{node}}\), we can distill knowledge in between these directly by defining a loss \(\mathcal{L}_{KD}\), such that

\[\mathcal{L}_{KD}=\mathcal{L}_{\text{feat}}(\mathcal{M}_{s}(H_{\text{node}, \text{s}}),\mathcal{M}_{t}(H_{\text{node},\text{t}})). \tag{4}\]

_edge-to-node (e2n):_ The GemNet-OC model heavily relies on its edge features, which are a key component in the directional message passing defined in the architecture and can be useful as a KD resource. However, the other models considered here do not have similar edge features to distill to. To accommodate that, we propose a KD strategy where we transfer information from GemNet-OC's edge features \(H_{\text{edge},(i,j)}\) by first aggregating them as follows:

\[H_{\text{edge2node},i}=\sum_{j\in\mathcal{N}(i)}H_{\text{edge},(i,j)}, \tag{5}\]

where \(i\) is the node index. The resulting vector \(H_{\text{edge2node},i}\) is a scalar, node-level feature, and we can, therefore, use it to transfer knowledge to the student node features \(H_{\text{node},\text{s}}\) as in Eq. (4).

_- vector-to-vector (v2v):_ Similarly, the PaiNN model defines special vectorial features, which differ substantially 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. This poses a new challenge when distilling knowledge from or onto PaiNN. To this end, we define a KD procedure, where we transfer knowledge between (equivariant) vectorial node features and (invariant) scalar edge features. We achieve that 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_{\text{edge},(i,j)}\}_{j\in\mathcal{N}}\) corresponding to a given node \(i\) into node-level equivariant vectorial features \(H_{\text{vec},i}\) by considering the unit vector \(\mathbf{u}_{ij}=\frac{1}{\left\lvert\mathbf{x}_{j}-\mathbf{x}_{i}\right\rvert}(\mathbf{x}_{j} -\mathbf{x}_{i})\) that defines the direction of the edge \((i,j)\), such that

\[H_{\text{vec},i}^{(k)}=\sum_{j\in\mathcal{N}(i)}\mathbf{u}_{i,j}H_{\text{edge},(i, j)}^{(k)}, \tag{6}\]

with the superscript \(k\) indicating the channel. This fulfills the condition of equivariance with respect to rotations as each vector \(\mathbf{u}_{ij}\) is equivariant to rotations, with \(H_{\text{edge},(i,j)}^{(k)}\) - a scalar not influencing its direction.

**Baseline KD strategies.** To validate the performance of our KD strategies, we evaluate their performance against 2 vanilla-based KD approaches suitable for regression tasks.

_Vanilla (1):_ As mentioned above, the main problem with using vanilla KD for regression is the lack of features analogous to logits. One way of adapting vanilla KD for regression is by steering the student to mimic the final output of the teacher directly [31]:

\[\mathcal{L}_{\text{KD}}=\alpha_{\text{E}}\mathcal{L}_{\text{E}}(\hat{E}_{ \text{s}},\hat{E}_{\text{t}})+\alpha_{\text{F}}\mathcal{L}_{\text{F}}(\hat{ \mathbf{F}}_{\text{s}},\hat{\mathbf{F}}_{\text{t}}), \tag{7}\]

where the subscripts \({}_{\text{s}}\) and \({}_{\text{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 [15, 22].

_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 a sum of the predicted contributions from each atom \(\hat{E}=\sum_{i}\hat{E}_{i}\). Hence, we note that we can extend the aforementioned vanilla KD approach by imposing a loss on these granular predictions instead. Following the energy definition above, the KD loss can be expressed as:

\[\mathcal{L}_{\text{KD}}=\frac{1}{N}\sum_{i=1}^{N}\mathcal{L}_{\text{E}}(\hat{ \text{E}}_{i,s},\hat{E}_{i,t}). \tag{8}\]

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

To evaluate our proposed methods, we perform experiments on the OC20-2M [32] dataset - a representative subset of the diverse OC20 dataset for molecular simulations [11]. We focus on the S2EF task, i.e., predicting energies and forces for a given system (atom types \(\mathbf{z}\) and positions \(\mathbf{X}\)), which is the task of broadest interest [32]. We use the models as implemented in the OC20 codebase2 (see Appendix A for detailed information about the training procedure and different model hyperparameters). We then validate the robustness of our approach by conducting additional experiments on the COLL dataset introduced in [14].

Footnote 2: [https://github.com/Open-Catalyst-Project/ocp](https://github.com/Open-Catalyst-Project/ocp)

**Benchmarking molecular GNNs.** We selected SchNet, PaiNN and GemNet-OC for our experiments as they cover most of the accuracy-complexity spectrum. To highlight this, we start by first benchmarking the predictive accuracy and computational cost of the three models on the OC20-2M [32] dataset. In conjunction with the default PaiNN model - referred to as PaiNN-big in this paper, we also utilize a more lightweight version of the architecture (referred to as PaiNN-small), where we reduce the number of hidden layers and their dimensionality. We train SchNet, PaiNN-small and PaiNN-big to convergence ourselves, whereas the GemNet-OC model we employ is the pre-trained model as available within the OC20 repository. The benchmarking results, summarized in Table 2, demonstrate the substantial tradeoff between predictive accuracy and computational cost across different GNN architectures, and motivate the need for methods that can alleviate this limitation.

**Similarity analysis.** To make our analyses 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) [33; 34]. In particular, we calculate the CKA similarity score across the hidden node representations of the trained SchNet, PaiNN and GemNet-OC (averaged over \(n=987\) nodes). The results of this analysis are summarised in Fig. 2. Focusing on the intra-model similarities first (the diagonal in Fig. 2), we observe that PaiNN and SchNet exhibit similar behavior, with representations having a generally high degree of similarity within the two models. This indicates that the information captured across the respective architectures is similar, which is consistent with the fact that the features of PaiNN and SchNet are iteratively updated by adding displacement features computed at each layer. GemNet-OC, however, displays completely different behavior, where the representations extracted at each layer are

\begin{table}
\begin{tabular}{l c c c c c}  & \begin{tabular}{c} Inference \\ Throughput \\ \end{tabular} & \multicolumn{5}{c}{OC20 S2EF Validation} \\ \cline{2-6}  & \begin{tabular}{c} Samples / \\ GPU sec. \(\uparrow\) \\ \end{tabular} & \begin{tabular}{c} Energy MAE \\ \(\mathrm{meV}\downarrow\) \\ \end{tabular} & \begin{tabular}{c} Force MAE \\ \(\mathrm{meV}/\mathrm{\AA}\downarrow\) \\ \end{tabular} & \begin{tabular}{c} Force cos \\ \(\uparrow\) \\ \end{tabular} & 
\begin{tabular}{c} EFwT \\ \% \(\uparrow\) \\ \end{tabular} \\ \hline SchNet & \(788.2\) & \(1308\) & \(65.1\) & \(0.204\) & \(0\) \\ PaiNN-small & \(618.2\) & \(489\) & \(47.1\) & \(0.345\) & \(0.085\) \\ PaiNN-big & \(237.8\) & \(440\) & \(45.3\) & \(0.376\) & \(0.14\) \\ GemNet-OC & \(75.8\) & \(286\) & \(25.7\) & \(0.598\) & \(1.06\) \\ \end{tabular}
\end{table}
Table 2: Evaluation of the performance of the four baseline models on the OC20 S2EF validation dataset (trained on OC20-2M), averaged across OC20’s four validation sets. The results show a clear tradeoff between accuracy and computational cost. For individual validation sets, see Appendix B.

Figure 2: Similarity analysis between the node features of SchNet, PaiNN and GemNet-OC using CKA.
significantly different from the rest of the architecture. This is likewise in line with the GemNet-OC model, where each layer embedding is a separate output feature. When examining the similarity between different models instead, we see that, generally speaking, the 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.

**KD setup.** Based on the aforementioned analyses, we define the following teacher-student pairs, covering the whole spectrum of architectural disparity: PaiNN-big to PaiNN-small (_same_ architecture); PaiNN-big to SchNet (_similar_ architectures); GemNet-OC to PaiNN-big (_different_ architectures). We train the student models by utilizing an offline KD strategy [18], where we distill knowledge from the more competent, pre-trained teacher model to the simpler, more lightweight student model. We augment the training of each student model with our KD protocols and evaluate the effect on predictive accuracy against the two vanilla KD approaches. If not mentioned otherwise, we utilize the following setup: we use MSE as a distillation loss \(\mathcal{L}_{\text{feat}}\); a learned linear as a transformation function \(\mathcal{M}_{s}\) on the features of the student; and the identity transformation as \(\mathcal{M}_{t}\). When distilling knowledge from GemNet-OC, we use the aggregated output block features \(\mathbf{x_{\text{E}}}\) and \(\mathbf{x_{\text{F}}}\) as node features and edge features, respectively. This is similar to the review-based setup proposed in [35]. For PaiNN and SchNet, we use the final node features. Other training details are available in Appendix A.

**Knowledge distillation results.** The results of the experiment are summarized in Table 3, presenting a comparative analysis of the predictive performance of different student models with and without the implementation of knowledge distillation. Especially notable was the improvement in energy predictions for PaiNN-big when trained with GemNet-OC as a teacher under the _n2n_ KD strategy. This configuration led to a reduction in PaiNN-big's energy MAE from 440 \(\mathrm{meV}\) to 349 \(\mathrm{meV}\). This equates to a decrease of over 20 % compared to the base PaiNN-big model, thus closing nearly 60 % of the accuracy gap between GemNet-OC and PaiNN-big in energy predictions. The _n2n_ KD strategy also led to an improvement in force predictions, albeit a less dramatic one, bridging 9.3 % of the

\begin{table}
\begin{tabular}{l l r r r r}  & & \multicolumn{4}{c}{OC20 S2EF Validation} \\ \cline{3-6}  & & Energy MAE & Force MAE & Force cos & EFwT \\ Model & \(\mathrm{meV}\downarrow\) & \(\mathrm{meV}/\mathrm{\AA}\downarrow\) & \(\uparrow\) & \% \(\uparrow\) \\ \hline \multirow{7}{*}{
\begin{tabular}{l} **S**-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** **\\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\ **PaiNN-** \\
performance gap between the two models. We manage to improve that with our custom _e2n_ and _v2v_ protocols, elevating the improvements in force predictions to 16.7 % and 15.6% respectively. Nonetheless, this was offset by a decline in energy prediction accuracy, which dropped to 0.6 % and 0.8 % respectively. In contrast, the _n2n_ proved to be the best strategy for both energy and force predictions when distilling PaiNN-big in PaiNN-small. For this teacher-student configuration, we were able to close the gap in energy and force prediction between the two models by \(\sim\)\(65\%\) and \(20\%\) respectively. The effect of _v2v_ KD on energy predictions was similar at 60.8 %. Yet, this approach was detrimental to force predictions (-9.8 %). In both cases, the performance of our two vanilla KD protocols was not as effective, either yielding smaller improvements or significantly impairing accuracy. Still, they showed the most promising results when distilling knowledge from PaiNN-big to SchNet. However, it is noteworthy to highlight that the effect of KD for this pair was tangibly lower at around 11 % for energy improvements and 2.5 % for forces. Based on previous research [23], we deduct this may be a result of the vast gap in complexity and performance between the two models, which affects the learning of the substantially more lightweight student model.

These experiments showcase the effectiveness of our KD strategies, which allowed us to consistently boost the accuracy of student models across all examined teacher-student configurations. Importantly, these findings remained consistent when broadening our experiment to the COLL dataset, where we similarly observed significant performance gains across all model configurations (see Appendix C). This further highlights the robustness of our training protocols and their utility across molecular datasets. Finally, we reiterate a central facet of our approach in that we enhance the performance of our student models solely by imposing an additional KD loss during training, without any structural alterations to their architecture. Hence, the inference throughput of the distilled student models remained completely unaffected.

**Effect of KD on model similarity.** We continue our investigation of KD on molecular GNNs by investigating how KD affects the student architectures in the various teacher-student configurations we have studied. To this end, we analyze how the CKA similarity between the teacher and the student models changes with the introduction of KD during training. We observed that KD introduces strong and specific similarity gains in layers that are used for KD, which also propagate along the student architecture (see Fig. 3). We saw similar behavior across KD strategies and teacher-student configurations (see Appendix F), allowing us to monitor and quantify the effect of KD on student models as we explore different KD settings and design choices. Hence, we believe that such similarity-based analyses can serve as a useful tool for quantitatively profiling KD procedures, which can elucidate the mechanism behind KD and inform the development of KD protocols.

**Hyperparameter studies** To increase the breadth of our work, we additionally conducted a thorough study to evaluate the effect of different design choices within our KD strategy. We present the result of this study here, hoping to provide valuable insights into the trends we have observed, which can inform future research in the area. The results are further detailed in Appendix E.

_Effect of distillation loss_: In terms of distillation losses \(\mathcal{L}_{\text{feat}}\), we explored: mean squared error (MSE); mean absolute error (MAE); and also directly optimizing CKA. Our analyses showed that MAE performed similarly to MSE, whereas using CKA substantially hurt accuracy.

Figure 3: Similarity analysis between Gemnet-OC and PaiNN-big without KD (left) and with KD (right). It illustrates a clear similarity gain within the feature pair used during KD (indicated with a \(\diagup\) in the middle heatmap), which also propagates to earlier layers of the student architecture. Similarity analyses for other KD strategies and teacher-student configurations are presented in Appendix F.
_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; as well as variants of the structure-preserving transformations described in [36]. Our results (see Appendix E.1) showed that the examined transformations either hurt performance or only provided a marginal improvement in performance compared to our default linear mapping.

_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); onto earlier layers in the student; across multiple layers in the teacher and the student. We observed that our default, review-based strategy was the best-performing KD protocol for this teacher-student configuration, with other approaches either hurting performance or only providing a marginal improvement (see Appendix E.3). We complemented our experiments by performing CKA-based similarity analyses, which allowed us to gain insights into how these protocols change the architecture of the student (see Appendix F).

## 5 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 energies and forces. We explored two protocols for knowledge distillation via training on novel data points labeled by the teacher, and briefly describe them here. See Appendix D for more details.

**Random rattling.** Apart from the relaxation trajectories in OC20, there is also a dataset called "Rattled", where random noise has been added to a subset of the systems along the relaxation trajectories, and the forces and energies have then been computed for these rattled systems using DFT. Adding noise to existing structures has also been used in the context of pretraining of molecular GNNs [37; 38] or as a regularization strategy [39]. However, adding noise to existing samples and aligning the predictions between teacher and student (without ground truth labels) did not show significant improvement. Additionally, we tried using gradient ascent to find perturbations that maximize the discrepancy between the teacher and student predictions, similar to [40], but this approach increased training time and did not show improvements over random noise.

**Synthetic Data.** One weakness of the OC20 datasets is that some samples originate from the same relaxation trajectory and are therefore correlated. To tackle this, we generated our own distilled dataset coined _d1M_ consisting of 1M samples, which were generated by sampling novel systems (generated with the OC Datasets codebase 3), running relaxations with the pre-trained GemNet-OC 2M model, and then subsampling approx. 10% of the frames. We tried different ways of incorporating this new d1M dataset, all of which are based on joint training on the 2M as well as the d1M datasets (similar to what we did with the rattled systems). To explore 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 should originate from each of the datasets, and (b) how to weight the loss contributions based on the dataset of origin. Unfortunately, and contrary to similar approaches e.g. in speech recognition [41], the results we observed didn't significantly improve on the baseline models.

Footnote 3: [https://github.com/Open-Catalyst-Project/Open-Catalyst-Dataset](https://github.com/Open-Catalyst-Project/Open-Catalyst-Dataset)

## 6 Conclusion

In this paper, we investigate the utility of knowledge distillation as a means of distilling larger, more computationally expensive GNNs for molecules into smaller, more efficient models. To this end, we propose three distinct feature-based KD strategies, which facilitate the distillation of intermediate representations across diverse GNN models. Our analyses demonstrate that our protocols can significantly enhance the performance of molecular GNN without any modifications to their architecture, allowing us to run faster molecular simulations without substantially impairing predictive accuracy. We validate our approach across different teacher-student configurations and datasets, showcasing the effectiveness and robustness of the method. One limitation of our approach is that even though inference times are not affected, we need to perform forward passes through the teacher during training, increasing training times.
With this work, we aim to elucidate the potential of KD in the domain of molecular simulations and stimulate future work in the area, which can bring about important advances in scientific disciplines as broad as material science, drug discovery and catalysis. However, it is noteworthy to highlight that such technologies can be potentially harmful if, e.g., used to simulate and discover systems that are toxic or intended to be used in harmful technology.

## Acknowledgments and Disclosure of Funding

This project would not have been possible without the computing resources provided by: the Knut and Alice Wallenberg Foundation at the National Supercomputer Centre (Berzelius resource); 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; as well as the Chair of Aerodynamics and Fluid Mechanics at Technical University of Munich. We are also grateful to the team behind 2022 London Geometry and Machine Learning Summer School (LOGML), where this research project was initially conceived. We would like to specially 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, we would like to thank Muhammed Shuaibi for providing the COLL dataset in LMDB format. Figures assembled in BioRender.com. Funding disclosure: Dimitar Georgiev is supported by UK Research and Innovation [UKRI Centre for Doctoral Training in AI for Healthcare grant number EP/S023283/1]; Filip Ekstrom Kelvinius is financially supported by the Excellence Center at Linkoping-Lund in Information Technology (ELLIIT).
