SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Filip Ekstrom Kelvinius*
Linköping University filip.ekstrom@liu.se
Dimitar Georgiev*  
Imperial College London  
d.georgiev21@imperial.ac.uk
Artur Petrov Toshev*  
Technical University of Munich  
artur.toshev@tum.de
Johannes Gasteiger  
Google Research  
johannesg@gmail.com
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.
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
*These authors contributed equally to this work. Order was determined by rolling a dice.
Figure 1: Using knowledge distillation, we manage to significantly boost the predictive accuracy of different student models on the OC20-2M [17] and COLL [6] datasets while fully preserving their inference throughput.
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-to-node  $(e2n)$  and vector-to-vector  $(\nu 2\nu)$  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<sup>2</sup>.
Molecular simulations. In this work, we consider molecular systems at an atomic level, i.e.,  $N$  atoms represented by their atomic numbers  $\mathbf{z} = \{z_{1},\dots,z_{N}\} \in \mathbb{Z}^{N}$  and positions  $\mathbf{X} = \{\pmb{x}_1,\dots,\pmb{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  $\pmb {F}\in \mathbb{R}^{N\times 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  $\mathcal{G} = (\mathcal{V}, \mathcal{E})$ , where the nodes  $\mathcal{V}$  correspond to the set of atoms, and edges  $\mathcal{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
$^{2}$ https://github.com/gasteigerjo/ocp/blob/main/DISTILL.md
atomic graphs  $\mathcal{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.,  $\hat{E},\hat{F} = \mathrm{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  $\mathcal{L}_0$  that combines both the energy and force prediction error as follows:
$$
\mathcal {L} _ {0} = \alpha_ {\mathrm {E}} \mathcal {L} _ {\mathrm {E}} (\hat {E}, E) + \alpha_ {\mathrm {F}} \mathcal {L} _ {\mathrm {F}} (\hat {\boldsymbol {F}}, \boldsymbol {F}), \tag {1}
$$
where  $E$  and  $\pmb{F}$  are the ground-truth energy and forces,  $\hat{E}$  and  $\dot{\pmb{F}}$  are the predictions of the model of interest, and  $\mathcal{L}_{\mathrm{E}}$  and  $\mathcal{L}_{\mathrm{F}}$  are some loss functions weighted by  $\alpha_{\mathrm{E}},\alpha_{\mathrm{F}}\in \mathbb{R}$ .
To perform knowledge distillation, we augment this training process by defining an auxiliary KD loss term  $\mathcal{L}_{\mathrm{KD}}$ , which is added to  $\mathcal{L}_0$  (with a factor  $\lambda \in \mathbb{R}$ ) to derive the final training loss function  $\mathcal{L}$ :
$$
\mathcal {L} = \mathcal {L} _ {0} + \lambda \mathcal {L} _ {\mathrm {K D}}. \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 [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  $\mathcal{L}_{\mathrm{KD}}$  that forces the student to mimic the predictions of the teacher model. This is usually achieved by constructing a loss  $\mathcal{L}_{\mathrm{KD}} = \mathrm{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_{\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} _ {\mathrm {K D}} = \mathcal {L} _ {\text {f e a t}} \left(\mathcal {M} _ {\mathrm {s}} \left(H _ {\mathrm {s}}\right), \mathcal {M} _ {\mathrm {t}} \left(H _ {\mathrm {t}}\right)\right), \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}_{\mathrm{feat}}: U \times U \mapsto \mathbb{R}^{+}$  is some loss of choice. The transformations  $\mathcal{M}_{\mathrm{s}}, \mathcal{M}_{\mathrm{t}}$  can be the identity transformation, linear projections and multilayer perceptron (MLP) projection heads; whereas for the distillation loss  $\mathcal{L}_{\mathrm{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 \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{\pmb{F}} = -\nabla \hat{E}$ .
- PaiNN [19]: A GNN based on equivariant message passing, which contains scalar node features  $\pmb{x} \in \mathbb{R}^{d_1}$  - used for energy prediction; as well as geometric vectorial node features  $\pmb{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 [8]: 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}_{\mathrm{E}}$  and  $\mathbf{x}_{\mathrm{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.
Table 1: An overview of the types of features available in the three models we use in this study.
SchNet	PaiNN	GemNet-OC
Scalar node features	✓	✓	✓
Scalar edge features	✓
Vectorial node features	✓
Output blocks	✓
Defining feature-based KD distillation strategies for molecular GNNs. In the context of the three models considered in this work, we devise the following KD strategies:
- node-to-node  $(n2n)$ : As all three models contain scalar node features  $H_{\mathrm{node}}$ , we can distill knowledge in between these directly by defining a loss  $\mathcal{L}_{KD}$  given by
$$
\mathcal {L} _ {K D} = \mathcal {L} _ {\text {f e a t}} \left(\mathcal {M} _ {s} \left(H _ {\text {n o d e , s}}\right), \mathcal {M} _ {t} \left(H _ {\text {n o d e , t}}\right)\right), \tag {4}
$$
where  $H_{\mathrm{node,s}}$  and  $H_{\mathrm{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  $\mathcal{L}_{\mathrm{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_{\mathrm{edge,s}}$  and  $H_{\mathrm{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_{\mathrm{edge},(i,j)}$  by first aggregating them as follows:
$$
H _ {\text {e d g e 2 n o d e}, i} = \sum_ {j \in \mathcal {N} (i)} H _ {\text {e d g e}, (i, j)}, \tag {5}
$$
where  $i$  is some node index. The resulting features  $H_{\mathrm{edge2node},i}$  are scalar, node-level features, and we can, therefore, use them to transfer knowledge to the student node features  $H_{\mathrm{node,s}}$  as in Equation (4).
- vector-to-vector  $(\nu 2\nu)$ : 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_{\mathrm{edge},(i,j)}\}_{j\in \mathcal{N}}$  corresponding to a given node  $i$  into node-level equivariant vectorial features  $H_{\mathrm{vec},i}$  by considering the unit vector  $u_{ij} = \frac{1}{|\pmb{x}_j - \pmb{x}_i|} (\pmb{x}_j - \pmb{x}_i)$  that defines the direction of the edge  $(i,j)$ , such that
$$
H _ {\text {v e c}, i} ^ {(k)} = \sum_ {j \in \mathcal {N} (i)} \boldsymbol {u} _ {i j} H _ {\text {e d g e}, (i, j)} ^ {(k)}, \tag {6}
$$
with the superscript  $k$  indicating the channel. Notice that the features  $H_{\mathrm{vec},i}^{(k)}$  fulfill the condition of equivariance with respect to rotations as each vector  $\pmb{u}_{ij}$  is equivariant to rotations, and  $H_{\mathrm{edge},(i,j)}^{(k)}$  - a scalar not influencing its direction. Consequently, it is important to use a loss  $\mathcal{L}_{\mathrm{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:
$$
\mathcal {L} _ {\mathrm {K D}} = \alpha_ {\mathrm {E}} \mathcal {L} _ {\mathrm {E}} \left(\hat {E} _ {\mathrm {s}}, \hat {E} _ {\mathrm {t}}\right) + \alpha_ {\mathrm {F}} \mathcal {L} _ {\mathrm {F}} \left(\hat {\boldsymbol {F}} _ {\mathrm {s}}, \hat {\boldsymbol {F}} _ {\mathrm {t}}\right), \tag {7}
$$
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  $\hat{E} = \sum_{i}\hat{E}_{i}$ . Hence, we can extend the aforementioned vanilla KD approach by imposing a loss on these granular predictions instead:
$$
\mathcal {L} _ {\mathrm {K D}} = \frac {1}{N} \sum_ {i = 1} ^ {N} \mathcal {L} _ {\mathrm {E}} \left(\hat {\mathrm {E}} _ {i, \mathrm {s}}, \hat {E} _ {i, \mathrm {t}}\right). \tag {8}
$$
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<sup>3</sup> (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-OC-small 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).