EasyST: A Simple Framework for Spatio-Temporal Prediction

TL;DR

EasyST tackles scalability and distribution-shift challenges in spatio-temporal forecasting by distilling knowledge from a heavy STGNN into a lightweight MLP. It combines a spatio-temporal information bottleneck with a teacher-bounded regression loss and augments learning with spatio-temporal prompts, optimizing a multi-term objective $L = L_{ ext{pre}}(\hat{Y},Y) + \lambda L_{ ext{kd}}(\hat{Y},Y^{T})$ to produce robust representations. The framework is model-agnostic with respect to the teacher and demonstrates superior efficiency and accuracy across traffic, crime, and weather datasets. Practical impact lies in enabling scalable, robust urban sensing with fast inference on resource-constrained devices.

Abstract

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.

Figures (6)

• Figure 1: Overall framework of the proposed EasyST.
• Figure 2: Predictive visualization of our EasyST with other baselines on PEMS traffic data.
• Figure 3: Model performance and inference time of representative methods on the test set of traffic and crime datasets.
• Figure 4: Performance evaluation w.r.t noisy (top) and missing (bottom) data.
• Figure 5: In model interpretation evaluation, KDE visualization for distribution of embeddings learned by DMSTGCN, StemGNN and the proposed EasyST.