Geometric Knowledge Distillation: Topology Compression for Graph Neural Networks

TL;DR

This paper introduces a novel geometric knowledge distillation method for GNNs, encoding topological information via Neural Heat Kernel alignment to transfer knowledge from a complete graph model to a sparser one.

Contribution

It proposes Neural Heat Kernel-based geometric distillation, connecting thermodynamics and GNN topology, to improve knowledge transfer between models with different graph structures.

Findings

Effective in various topological knowledge transfer scenarios

Improves performance of student GNNs on sparser graphs

Demonstrates the importance of geometric information in GNN distillation

Abstract

We study a new paradigm of knowledge transfer that aims at encoding graph topological information into graph neural networks (GNNs) by distilling knowledge from a teacher GNN model trained on a complete graph to a student GNN model operating on a smaller or sparser graph. To this end, we revisit the connection between thermodynamics and the behavior of GNN, based on which we propose Neural Heat Kernel (NHK) to encapsulate the geometric property of the underlying manifold concerning the architecture of GNNs. A fundamental and principled solution is derived by aligning NHKs on teacher and student models, dubbed as Geometric Knowledge Distillation. We develop non- and parametric instantiations and demonstrate their efficacy in various experimental settings for knowledge distillation regarding different types of privileged topological information and teacher-student schemes.