FreeKD: Free-direction Knowledge Distillation for Graph Neural Networks

TL;DR

FreeKD introduces a reinforcement learning-based framework for graph neural network knowledge distillation that eliminates the need for a well-optimized teacher, enabling two shallow GNNs to exchange knowledge dynamically and effectively.

Contribution

This work is the first to propose a free-direction knowledge distillation framework for GNNs using reinforcement learning, removing the dependency on a deep teacher model.

Findings

FreeKD outperforms baseline GNNs on five benchmark datasets.

FreeKD achieves comparable or better results than traditional KD methods with deep teachers.

The framework is compatible with various GNN architectures.

Abstract

Knowledge distillation (KD) has demonstrated its effectiveness to boost the performance of graph neural networks (GNNs), where its goal is to distill knowledge from a deeper teacher GNN into a shallower student GNN. However, it is actually difficult to train a satisfactory teacher GNN due to the well-known over-parametrized and over-smoothing issues, leading to invalid knowledge transfer in practical applications. In this paper, we propose the first Free-direction Knowledge Distillation framework via Reinforcement learning for GNNs, called FreeKD, which is no longer required to provide a deeper well-optimized teacher GNN. The core idea of our work is to collaboratively build two shallower GNNs in an effort to exchange knowledge between them via reinforcement learning in a hierarchical way. As we observe that one typical GNN model often has better and worse performances at different nodes during training, we devise a dynamic and free-direction knowledge transfer strategy that consists of two levels of actions: 1) node-level action determines the directions of knowledge transfer between the corresponding nodes of two networks; and then 2) structure-level action determines which of the local structures generated by the node-level actions to be propagated. In essence, our FreeKD is a general and principled framework which can be naturally compatible with GNNs of different architectures. Extensive experiments on five benchmark datasets demonstrate our FreeKD outperforms two base GNNs in a large margin, and shows its efficacy to various GNNs. More surprisingly, our FreeKD has comparable or even better performance than traditional KD algorithms that distill knowledge from a deeper and stronger teacher GNN.