用轻量MLP蒸馏复杂图神经网络,提升城市时空预测效率与泛化能力。
EasyST: A Simple Framework for Spatio-Temporal Prediction
- 通过知识蒸馏让简单MLP学习复杂GNN的时空特征表示
- 在三个真实城市数据集上达到更高精度且训练更快
- 适合需要快速部署、低资源消耗的智能交通与环境监测场景
时空预测是数据驱动城市计算中的关键研究方向,涉及交通管理、公共安全和环境监测。然而,可扩展性和泛化能力仍是主要挑战。现有先进模型多依赖图神经网络(GNN)捕捉时空相关性,但在大规模数据下复杂度上升,递归式GNN消息传递机制限制了实际部署。此外,长周期大规模数据引入分布偏移,要求更强的泛化性能。为此,我们提出EasyST框架:通过有效蒸馏复杂时空GNN的知识,训练出轻量且鲁棒的多层感知机(MLP)。通过融合时空信息瓶颈与教师约束回归损失,过滤无关噪声,避免错误引导。进一步引入空间与时间提示,增强学生模型对下游任务的上下文理解能力。在三个城市计算任务的数据集上评估显示,EasyST在效率与准确性上均优于当前最优方法。代码已开源:https://github.com/HKUDS/EasyST。
原文摘要 · Abstract (English)
Spatio-temporal prediction is a crucial research area in data-driven urban computing, with implications for transportation, public safety, and environmental monitoring. However, scalability and generalization challenges remain significant obstacles. Advanced models often rely on Graph Neural Networks to encode spatial and temporal correlations, but struggle with the increased complexity of large-scale datasets. The recursive GNN-based message passing schemes used in these models hinder their training and deployment in real-life urban sensing scenarios. Moreover, long-spanning large-scale spatio-temporal data introduce distribution shifts, necessitating improved generalization performance. To address these challenges, we propose a simple framework for spatio-temporal prediction - EasyST paradigm. It learns lightweight and robust Multi-Layer Perceptrons (MLPs) by effectively distilling knowledge from complex spatio-temporal GNNs. We ensure robust knowledge distillation by integrating the spatio-temporal information bottleneck with teacher-bounded regression loss, filtering out task-irrelevant noise and avoiding erroneous guidance. We further enhance the generalization ability of the student model by incorporating spatial and temporal prompts to provide downstream task contexts. Evaluation on three spatio-temporal datasets for urban computing tasks demonstrates that EasyST surpasses state-of-the-art approaches in terms of efficiency and accuracy. The implementation code is available at: https://github.com/HKUDS/EasyST.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。