Distilling Expressive GNNs into MPNNs

Fabian Jogl, Pascal Welke, Thomas Gärtner

Keywords: Graph Neural Networks, Expressivity, Knowledge Distillation

TL;DR: We show that expressive graph neural networks can be distilled into simple message passing graph neural networks for significant improvements in predictive performance and speed.

Abstract: We investigate the distillation of expressive graph neural networks (GNNs) into simple message passing graph neural networks (MPNNs) for the case of molecular graph data. Under the standard training protocols used in prior work, expressive GNNs substantially outperform simple MPNNs on popular molecular benchmarks. We demonstrate that knowledge distillation closes 50 to 100% of this gap. Importantly, the distilled MPNNs are 2 to 33 times faster than their teachers. Our results suggest that for these molecular tasks, the performance gap is largely due to optimization challenges rather than fundamental expressivity limitations.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 5564