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

📄 arXiv: 2401.11798v4 📥 PDF

作者: Mohammad Izadi, Mehran Safayani, Abdolreza Mirzaei

分类: cs.LG, cs.AI

发布日期: 2024-01-22 (更新: 2024-09-24)

DOI: 10.1080/1206212X.2024.2442699

💡 一句话要点

提出知识蒸馏方法以提升时空图卷积网络的交通预测效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 时空图神经网络 交通预测 模型压缩 实时预测 深度学习

📋 核心要点

  1. 现有的时空图神经网络在处理真实世界交通数据时，实时预测效率不足，难以满足动态交通环境的需求。
  2. 本文提出通过知识蒸馏技术，利用教师网络的复杂模式指导学生网络学习，从而提升ST-GNN的执行效率。
  3. 在PeMSD7和PeMSD8数据集上的实验结果显示，学生网络在仅保留3%参数的情况下，准确性仍接近教师网络。

📝 摘要（中文）

高效的实时交通预测对于减少交通时间至关重要。本文采用时空图神经网络（ST-GNN）对实时交通数据进行建模，尽管其能力强大，但在处理真实世界交通数据时仍面临实时预测效率的挑战。为此，本文引入知识蒸馏（KD）作为解决方案，以提高ST-GNN在交通预测中的执行时间。我们设计了一种成本函数，旨在通过蒸馏数据训练参数更少的网络（学生网络），同时保持其准确性接近复杂网络（教师网络）。通过利用教师网络的时空相关性，学生网络能够学习教师所感知的复杂模式。此外，本文提出了一种算法，通过成本函数计算剪枝分数，解决小型网络架构搜索问题，并在每个剪枝阶段联合微调网络。最终，我们在两个真实世界数据集PeMSD7和PeMSD8上评估了所提方法，结果表明，即使仅保留3%的网络参数，学生网络的准确性仍能接近教师网络。

🔬 方法详解

问题定义：本文旨在解决时空图卷积网络在实时交通预测中的效率不足问题。现有方法在处理动态交通数据时，往往无法提供及时的预测结果，影响交通管理和决策。

核心思路：通过引入知识蒸馏技术，利用复杂教师网络的输出指导参数更少的学生网络学习，从而提升模型的执行效率和准确性。该方法旨在在保持预测精度的同时，减少模型的计算负担。

技术框架：整体架构包括教师网络和学生网络两个主要模块。教师网络负责提取复杂的时空特征，而学生网络则通过知识蒸馏学习这些特征。算法通过计算剪枝分数来优化学生网络的结构，并在每个剪枝阶段进行联合微调。

关键创新：本文的创新点在于提出了一种新的成本函数，用于训练学生网络，同时保持其准确性接近教师网络。此外，算法有效解决了小型网络架构搜索的问题，提升了模型的实用性。

关键设计：在损失函数设计上，结合了教师网络的时空相关性，确保学生网络能够有效学习复杂模式。网络结构方面，学生网络的参数设置经过精心设计，以实现高效的剪枝和微调。实验中，学生网络在仅保留3%参数的情况下，依然能够保持较高的预测准确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，所提方法在PeMSD7和PeMSD8数据集上表现优异，学生网络在仅保留3%参数的情况下，准确性仍接近教师网络，显示出显著的性能提升。这一结果验证了知识蒸馏在时空图卷积网络中的有效性，具有较强的应用前景。

🎯 应用场景

该研究的潜在应用领域包括智能交通系统、城市交通管理和自动驾驶技术。通过提高交通预测的实时性和准确性，能够有效优化交通流量，减少拥堵，提升出行效率，具有重要的实际价值和社会影响。未来，该方法还可以扩展到其他时空数据预测任务，如气象预测和人流量预测等。

📄 摘要（原文）

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.