Task-aware memory improves compact models for spatiotemporal prediction
TAM: Task-Aware Memory Distillation for Efficient Spatiotemporal Prediction
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Spatiotemporal prediction means forecasting how things change over time and space, like weather or traffic flow. The authors show that teaching a smaller model by letting it remember patterns the bigger model noticed before helps improve its accuracy. They organize the teacher’s knowledge into a memory of past examples that the smaller model uses during training but not when making predictions. This method, TAM, works better than older techniques on different tasks without making the smaller model slower.
What this means in practice
- •For weather forecast teams: Improve compact weather models by training them with a memory of past teacher predictions to boost accuracy without extra prediction cost.
- •For traffic management teams: Develop smaller traffic flow prediction systems that learn spatiotemporal patterns from a teacher’s historical projections to enhance forecasting performance.
Authors
Yuqi Li, Xiaoqin Feng, Fan Xu, Weilun Feng, Chuanguang Yang, Yingli Tian, Hao Wu
Abstract
Knowledge distillation enables efficient spatiotemporal prediction by transferring knowledge from an accurate teacher to a compact student. However, matching outputs or features independently for each sample leaves cross-sample predictive structure underused. Exploiting this structure requires representations and historical references that reflect the dynamics of each task. We propose TAM, a Task-Aware Memory Distillation framework that organizes a frozen teacher's knowledge into a bounded, retrievable history. Memory entries encode latent features, forecast changes, or flow residuals, while task-specific selection rules identify relevant historical references. The student either matches the teacher's similarity distribution over shared references or regresses observation-conditioned residual prototypes. These objectives complement supervised prediction and conventional distillation. The teacher, memory, and auxiliary adapters are used only during training, leaving student inference unchanged. We evaluate TAM on video prediction, weather forecasting, and traffic flow prediction across multiple teacher-student configurations. Averaged over four paired runs, adding TAM improves SSIM on all six video datasets and reduces MSE on five relative to the corresponding KD baselines. Mean paired MSE reductions reach 1.86% on KittiCaltech, 1.93% on WeatherBench with a gSTA teacher, and 1.01% on TaxiBJ. These results demonstrate the utility of historical teacher supervision across distinct forecasting tasks without additional student inference cost.