# TAM: Task-Aware Memory Distillation for Efficient Spatiotemporal Prediction

arXiv: https://arxiv.org/abs/2610.11617 · HTML: https://mlsanity.com/paper/2610.11617/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Multimodal distillation, one of 12 directions in Knowledge distillation.
- Builds on 39 of its 50 references that are in the map.
- First posted on arXiv 8 Oct 2026.

## Placed in

- Knowledge distillation / [Multimodal distillation](https://mlsanity.com/distill/multimodal-distillation/)

## Builds on (most cited first, 30 of 39)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [FitNets: Hints for Thin Deep Nets](https://mlsanity.com/paper/1412.6550/) (2014)
- [Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting](https://mlsanity.com/paper/1506.04214/) (2015)
- [Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer](https://mlsanity.com/paper/1612.03928/) (2016)
- [Relational Knowledge Distillation](https://mlsanity.com/paper/1904.05068/) (2019)
- [Contrastive Representation Distillation](https://mlsanity.com/paper/1910.10699/) (2019)
- [Similarity-Preserving Knowledge Distillation](https://mlsanity.com/paper/1907.09682/) (2019)
- [Unsupervised Learning of Video Representations using LSTMs](https://mlsanity.com/paper/1502.04681/) (2015)
- [On the Efficacy of Knowledge Distillation](https://mlsanity.com/paper/1910.01348/) (2019)
- [Deep Spatio-Temporal Residual Networks for Citywide Crowd Flows Prediction](https://mlsanity.com/paper/1610.00081/) (2016)
- [Skillful Precipitation Nowcasting using Deep Generative Models of Radar](https://mlsanity.com/paper/2104.00954/) (2021)
- [Distilling Knowledge via Knowledge Review](https://mlsanity.com/paper/2104.09044/) (2021)
- [Structured Knowledge Distillation for Dense Prediction](https://mlsanity.com/paper/1903.04197/) (2019)
- [WeatherBench: A benchmark dataset for data-driven weather forecasting](https://mlsanity.com/paper/2002.00469/) (2020)
- [Unsupervised Learning for Physical Interaction through Video Prediction](https://mlsanity.com/paper/1605.07157/) (2016)
- [Knowledge Transfer via Distillation of Activation Boundaries Formed by Hidden Neurons](https://mlsanity.com/paper/1811.03233/) (2018)
- [Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control](https://mlsanity.com/paper/1812.00568/) (2018)
- [A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay](https://mlsanity.com/paper/1803.09820/) (2018)
- [Disentangling Physical Dynamics from Unknown Factors for Unsupervised Video Prediction](https://mlsanity.com/paper/2003.01460/) (2020)
- [UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning](https://mlsanity.com/paper/2201.04676/) (2022)
- [SimVP: Simpler yet Better Video Prediction](https://mlsanity.com/paper/2206.05099/) (2022)
- [Channel-wise Knowledge Distillation for Dense Prediction](https://mlsanity.com/paper/2011.13256/) (2020)
- [Earthformer: Exploring Space-Time Transformers for Earth System Forecasting](https://mlsanity.com/paper/2207.05833/) (2022)

## Abstract (arXiv)

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.