Layer-to-Layer Knowledge Mixing in Graph Neural Network for Chemical Property Prediction

TL;DR

This paper introduces Layer-to-Layer Knowledge Mixing (LKM), a novel self-distillation method that enhances GNN accuracy for chemical property prediction by efficiently aggregating multi-scale information with minimal additional computational cost.

Contribution

The study proposes LKM, a new self-knowledge distillation technique that improves GNN performance on chemical datasets without significantly increasing computational complexity.

Findings

LKM reduces MAE by up to 9.8% on QM9.

LKM improves accuracy by 45.3% on MD17 Energy.

LKM enhances predictions on Chignolin by 22.9%.

Abstract

Graph Neural Networks (GNNs) are the currently most effective methods for predicting molecular properties but there remains a need for more accurate models. GNN accuracy can be improved by increasing the model complexity but this also increases the computational cost and memory requirement during training and inference. In this study, we develop Layer-to-Layer Knowledge Mixing (LKM), a novel self-knowledge distillation method that increases the accuracy of state-of-the-art GNNs while adding negligible computational complexity during training and inference. By minimizing the mean absolute distance between pre-existing hidden embeddings of GNN layers, LKM efficiently aggregates multi-hop and multi-scale information, enabling improved representation of both local and global molecular features. We evaluated LKM using three diverse GNN architectures (DimeNet++, MXMNet, and PAMNet) using datasets of quantum chemical properties (QM9, MD17 and Chignolin). We found that the LKM method effectively reduces the mean absolute error of quantum chemical and biophysical property predictions by up to 9.8% (QM9), 45.3% (MD17 Energy), and 22.9% (Chignolin). This work demonstrates the potential of LKM to significantly improve the accuracy of GNNs for chemical property prediction without any substantial increase in training and inference cost.