[论文解读] Temporal Attention Unit: Towards Efficient Spatiotemporal Predictive Learning

Cheng Tan, Zhangyang Gao+5 more|arXiv (Cornell University)|Jun 24, 2022

Human Pose and Action Recognition被引用 8

一句话总结

本文提出时间注意力单元（TAU），一种可并行化的注意力机制，将时间建模分解为帧内静态注意力与帧间动态注意力，取代时空预测中的循环单元。该方法在多个基准测试中实现最先进性能，训练速度显著提升（V100上每轮仅2.5分钟），FLOPs更低，同时引入一种新型微分散度正则化，提升泛化能力，超越MSE损失。

ABSTRACT

Spatiotemporal predictive learning aims to generate future frames by learning from historical frames. In this paper, we investigate existing methods and present a general framework of spatiotemporal predictive learning, in which the spatial encoder and decoder capture intra-frame features and the middle temporal module catches inter-frame correlations. While the mainstream methods employ recurrent units to capture long-term temporal dependencies, they suffer from low computational efficiency due to their unparallelizable architectures. To parallelize the temporal module, we propose the Temporal Attention Unit (TAU), which decomposes the temporal attention into intra-frame statical attention and inter-frame dynamical attention. Moreover, while the mean squared error loss focuses on intra-frame errors, we introduce a novel differential divergence regularization to take inter-frame variations into account. Extensive experiments demonstrate that the proposed method enables the derived model to achieve competitive performance on various spatiotemporal prediction benchmarks.

研究动机与目标

• 解决由于其顺序性、不可并行化特性，导致在时空预测学习中循环单元计算效率低下的问题。
• 开发一种新型时间模块，实现在保持长期时间依赖建模能力的同时支持并行计算。
• 通过考虑帧间差异，超越帧内重建误差，提升泛化能力。
• 证明仅使用基于2D CNN的简单空间编码器/解码器配合TAU，即可超越复杂的循环架构。
• 建立一种基于注意力机制的新型高效、可扩展且有效的时空表征学习范式。

提出的方法

• 提出时间注意力单元（TAU），将时间注意力分解为两部分：帧内静态注意力（跨帧的空间自注意力）与帧间动态注意力（帧间的时间交叉注意力）。
• 通过解耦时间注意力计算，设计可并行化的架构，实现完整的GPU并行化，训练速度显著快于循环单元。
• 引入微分散度正则化（DDR），惩罚真实序列与预测序列之间在时间动态上的差异，作为MSE损失的补充。
• 使用2D卷积神经网络作为空间编码器与解码器，以保持结构简洁与计算高效。
• 端到端训练模型，采用联合损失：MSE用于像素级重建，DDR用于保证帧间运动一致性。
• 通过将预测帧迭代反馈作为输入，实现自回归推理，生成更长的未来序列。

实验结果

研究问题

• RQ1 可并行化的注意力机制是否能在保持或提升性能的前提下，替代时空预测中的循环单元？
• RQ2 将时间注意力分解为静态（帧内）与动态（帧间）两部分，是否能增强对长程依赖的建模能力？
• RQ3 微分散度正则化是否能通过捕捉像素级重建之外的帧间运动动态，提升泛化能力？
• RQ4 与循环基线相比，所提方法在训练速度与推理效率方面如何扩展？
• RQ5 当在某一数据集上训练并在不同数据集上测试时（如KITTI到Caltech行人数据集），模型是否具备良好的跨域泛化能力？

主要发现

• 所提出的基于TAU的模型在多个基准测试中达到最先进性能，包括在KTH数据集上（10→20帧）SSIM达0.911，PSNR达34.13，优于SimVP与E3D-LSTM等先前方法。
• 在Moving MNIST数据集上，模型仅用50轮即收敛至MSE 35.0，远快于CrevNet（每轮30分钟）与PhyDNet（每轮7分钟），且在单张V100上仅需2.5分钟/轮训练时间。
• 模型仅需15.96 GFLOPs，显著低于E3D-LSTM（298.85 GFLOPs）与SimVP（19.43 GFLOPs），展现出极高的计算效率。
• 消融实验证实TAU模块与微分散度正则化均不可或缺：若移除DDR，MSE升至21.1（完整模型为19.8）；若用标准卷积替代TAU，性能下降至58.9 MSE。
• 模型在跨域设置下泛化能力出色：在KITTI上训练、在Caltech行人数据集上测试时，仍保持强性能，表明其具备鲁棒的特征学习能力。
• 通过迭代自回归推理，模型成功生成长达40帧的未来序列，且质量保持一致，无明显退化。