[Paper Review] Accelerating Molecular Graph Neural Networks via Knowledge Distillation

Filip Ekström Kelvinius, Dimitar Georgiev+2 more|arXiv (Cornell University)|Jun 26, 2023

Machine Learning in Materials Science4 citations

TL;DR

This paper proposes knowledge distillation (KD) strategies—node-to-node, edge-to-edge, edge-to-node, and vector-to-vector—to accelerate molecular graph neural networks (GNNs) without architectural changes. By distilling hidden representations from large, accurate teacher GNNs into smaller student models, the method achieves up to 96.7% and 62.5% of teacher performance on energy and force prediction, respectively, while preserving the student model’s inference speed.

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.

Motivation & Objective

• Address the growing computational cost of state-of-the-art molecular GNNs, which limits their use in large-scale simulations.
• Improve the speed-accuracy trade-off in molecular property prediction by accelerating lightweight GNNs without architectural modifications.
• Develop specialized KD strategies tailored for directional and equivariant GNNs used in molecular simulations.
• Evaluate the effectiveness of KD across diverse teacher-student configurations and datasets, including OC20-2M and COLL.
• Investigate data augmentation and training protocol optimization to further enhance student model performance.

Proposed method

• Propose four novel KD strategies: node-to-node (n2n), edge-to-edge (e2e), edge-to-node (e2n), and vector-to-vector (v2v) distillation for equivariant GNNs.
• Use a pre-trained, high-accuracy teacher GNN to generate target representations (energy, forces, hidden states) for knowledge distillation.
• Train student GNNs using a multi-component loss combining standard regression loss and KD losses from hidden representations.
• Apply distillation to both node-level and edge-level representations, including scalar and vector features, to preserve geometric and structural information.
• Optimize training hyperparameters, including data mixing ratios and loss weighting between real and synthetic data.
• Explore data augmentation via random rattling and synthetic data generation (d1M) to improve student generalization.

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.

Experimental results

Research questions

• RQ1 Can knowledge distillation effectively accelerate molecular GNNs while preserving inference speed?
• RQ2 How do different feature-based distillation strategies (n2n, e2e, e2n, v2v) impact student model performance on energy and force prediction?
• RQ3 To what extent can distillation close the accuracy gap between large teacher models and small student models in molecular simulations?
• RQ4 Does data augmentation, such as random rattling or synthetic data, further improve student model performance?
• RQ5 How do hyperparameters like data mixing ratio and loss weighting affect the final performance of distilled models?

Key findings

• The proposed KD strategies consistently improve student model performance across multiple datasets and teacher-student configurations without architectural changes.
• The method achieves up to 96.7% of the teacher model’s accuracy on energy prediction and 62.5% on force prediction, significantly closing the accuracy gap.
• Inference throughput remains fully preserved, as the student models are lightweight and inference-efficient.
• Data augmentation via random rattling did not yield significant improvements, and gradient-based perturbations were less effective than random noise.
• Synthetic data generation (d1M) did not substantially improve performance over baseline joint training with real data, contrary to expectations from other domains.
• Comprehensive ablation studies confirmed the effectiveness of individual KD components, with vector-to-vector (v2v) and node-to-node (n2n) strategies showing strong gains.

Figure 2: Similarity analysis between the node features of SchNet, PaiNN and GemNet-OC using CKA (averaged over $n=987$ nodes).