TAM: Task-Aware Memory Distillation for Efficient Spatiotemporal Prediction

📄 arXiv: 2610.11617v1 📥 PDF

作者: Yuqi Li, Xiaoqin Feng, Fan Xu, Weilun Feng, Chuanguang Yang, Yingli Tian, Hao Wu

分类: cs.CV, cs.AI

发布日期: 2026-10-08

备注: 19 pages

💡 一句话要点

提出TAM框架以提升时空预测的效率与准确性

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱八：物理动画 (Physics-based Animation)

关键词: 知识蒸馏 时空预测 任务感知 记忆蒸馏 视频分析 天气预报 交通流预测

📋 核心要点

  1. 现有知识蒸馏方法未能充分利用跨样本的预测结构，导致时空预测效率低下。
  2. TAM框架通过组织教师知识为可检索的历史记录，结合任务特定选择规则，提升学生模型的学习效率。
  3. 在多个数据集上，TAM显著提高了SSIM和降低了MSE，验证了其在不同预测任务中的有效性。

📝 摘要（中文）

知识蒸馏通过将准确教师的知识转移到紧凑学生上，实现高效的时空预测。然而，现有方法独立匹配每个样本的输出或特征，未能充分利用跨样本的预测结构。为此，本文提出了TAM（任务感知记忆蒸馏）框架，组织冻结教师的知识为可检索的历史记录。记忆条目编码潜在特征、预测变化或流动残差，而任务特定选择规则识别相关历史参考。学生通过匹配教师在共享参考上的相似性分布或回归观察条件下的残差原型来进行学习。实验结果表明，TAM在视频预测、天气预报和交通流预测等任务上均显著提升了性能，验证了历史教师监督的有效性。

🔬 方法详解

问题定义：本文旨在解决现有知识蒸馏方法在时空预测中未能充分利用跨样本预测结构的问题，导致模型效率低下和预测精度不足。

核心思路：TAM框架通过将教师的知识组织为可检索的历史记录，结合任务特定的选择规则，使学生模型能够更有效地学习和利用历史信息。

技术框架：TAM的整体架构包括三个主要模块：冻结的教师模型、任务感知的记忆模块和辅助适配器。教师模型用于知识传递，记忆模块存储历史信息，适配器则在训练阶段辅助学生模型。

关键创新：TAM的创新在于通过任务感知的记忆蒸馏，利用历史教师监督来增强学生模型的学习能力，这与传统的独立样本蒸馏方法有本质区别。

关键设计：在设计中，记忆条目编码了潜在特征和预测变化，损失函数结合了教师的相似性分布和观察条件下的残差原型，确保学生模型在训练时能够有效利用历史信息。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

在实验中，TAM在六个视频数据集上均提高了SSIM，并在五个数据集上降低了MSE，相较于传统知识蒸馏基线，KittiCaltech数据集的MSE减少了1.86%，WeatherBench数据集减少了1.93%，TaxiBJ数据集减少了1.01%。这些结果表明TAM在不同预测任务中的有效性和优越性。

🎯 应用场景

该研究在视频预测、天气预报和交通流预测等领域具有广泛的应用潜力。通过提升模型的学习效率和预测准确性，TAM框架可为智能交通系统、气象预测和视频分析等实际应用提供更可靠的支持，推动相关领域的发展。未来，TAM的理念也可扩展到其他时空数据分析任务中，进一步提升模型性能。

📄 摘要（原文）

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.