Video World Models · 2026-10-09
TAM: Task-Aware Memory Distillation for Efficient Spatiotemporal Prediction
Yuqi Li, Xiaoqin Feng, Fan Xu, Weilun Feng, Chuanguang Yang, Yingli Tian, Hao Wu
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
TAM (Task-Aware Memory Distillation) organizes a frozen teacher’s knowledge into a bounded, retrievable history to provide cross-sample supervision for spatiotemporal prediction. Task-specific rules choose what representations to store (latent features, forecast changes, or flow residuals) and which historical entries to retrieve. The student learns either relational matching over shared references or residual-prototype regression; teacher/memory/adapters are training-only, so inference cost is unchanged.
Why it matters
The authors argue that independent per-sample output/feature matching underuses cross-sample predictive structure in forecasting. TAM explicitly structures historical teacher supervision via a task-aware memory (what to store + which references to use), and evaluates it across video prediction, weather forecasting, and traffic flow prediction with matched distillation baselines.
Method
- Build a bounded teacher memory used only during training; memory entries encode latent features, forecast changes, or flow residuals, chosen per task.
- Use task-specific reference selection (e.g., motion/occupancy for video, geographic location for weather, observation keys for traffic) to retrieve shared references.
- Train the student with either (i) relation matching via teacher/student similarity distributions over shared references or (ii) observation-conditioned residual-prototype regression for traffic; inference uses only the student.
Abstract (from 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.
Related papers
- Learning What to Recall: Adaptive Multi-Cue Episodic Memory for World Models
- KITE: A Knowledge-Injected Temporal Encoder for Time-Series Event Forecasting
- MORE: Learning Motion-Oriented Registers for Driving World Models and Planning
- ST-GRU PredRNN-V2: A GRU-Based Reformulation of PredRNN-V2 for Video Prediction
- Pretrained Features as Teachers, Not States: Compact Latent World Models for Visual Planning
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-08 · Indexed by Pub-AI: 2026-10-09 UTC
TAM: Task-Aware Memory Distillation for Efficient Spatiotemporal Prediction
Preprint posted: 2026-10-08 · Original record ↗
- arxiv
- 2610.11617
Explicit source versions (1)
- v1 · 2026-10-08
- Preprint available
Source date: 2026-10-08 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
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.