Knowledge Distillation on Spatial-Temporal Graph Convolutional Network for Traffic Prediction

TL;DR

This paper introduces a knowledge distillation approach to improve the efficiency of spatial-temporal graph neural networks for traffic prediction, achieving high accuracy with significantly fewer parameters.

Contribution

It proposes a novel cost function and pruning algorithm for designing lightweight student networks that retain the teacher network's accuracy in traffic prediction tasks.

Findings

Student networks retain close accuracy to teachers with only 3% of parameters.

The method effectively models spatial-temporal correlations in real-world traffic data.

Achieves efficient real-time traffic prediction on PeMS datasets.

Abstract

Efficient real-time traffic prediction is crucial for reducing transportation time. To predict traffic conditions, we employ a spatio-temporal graph neural network (ST-GNN) to model our real-time traffic data as temporal graphs. Despite its capabilities, it often encounters challenges in delivering efficient real-time predictions for real-world traffic data. Recognizing the significance of timely prediction due to the dynamic nature of real-time data, we employ knowledge distillation (KD) as a solution to enhance the execution time of ST-GNNs for traffic prediction. In this paper, We introduce a cost function designed to train a network with fewer parameters (the student) using distilled data from a complex network (the teacher) while maintaining its accuracy close to that of the teacher. We use knowledge distillation, incorporating spatial-temporal correlations from the teacher network to enable the student to learn the complex patterns perceived by the teacher. However, a challenge arises in determining the student network architecture rather than considering it inadvertently. To address this challenge, we propose an algorithm that utilizes the cost function to calculate pruning scores, addressing small network architecture search issues, and jointly fine-tunes the network resulting from each pruning stage using KD. Ultimately, we evaluate our proposed ideas on two real-world datasets, PeMSD7 and PeMSD8. The results indicate that our method can maintain the student's accuracy close to that of the teacher, even with the retention of only 3% of network parameters.