[論文レビュー] 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

ひとこと要約

本稿では、時空間予測における再帰的ユニットの代わりに、時系列モデリングをフレーム内静的自己注意とフレーム間動的クロス注意に分解する並列化可能なアテンション機構、Temporal Attention Unit (TAU) を提案する。この手法は、複数のベンチマークで最先端の性能を達成し、V100で1エポックあたり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を組み合わせた単純なモデルが、複雑な再帰的アーキテクチャを上回ることを示す。
• アテンション機構を用いた、効率的でスケーラブルかつ効果的な時空間表現学習の新しいパラダイムを確立する。

提案手法

• 時系列アテンションを2つのコンponentに分解するTemporal Attention Unit (TAU) を提案する：フレーム内静的自己注意（フレーム内での空間的自己注意）とフレーム間動的クロス注意（フレーム間での時系列的クロス注意）。
• 時系列アテンション計算を分離することで、並列化可能なアーキテクチャを設計し、再帰的ユニットよりも高速なGPU並列化と高速な学習を実現する。
• 予測シーケンスと真値シーケンスの間の時系列的ダイナミクスの差異をペナルティ化するため、微分発散正則化（DDR）を導入する。これはMSE損失を補完する。
• 空間エンコーダおよびデコーダとして2D畳み込みニューラルネットワークを用いることで、単純さと効率性を維持する。
• ピxls単位の再構成を目的としたMSEと、フレーム間の運動一貫性を目的としたDDRを組み合わせた損失関数で、モデルをエンドツーエンドに学習する。
• 予測フレームをフィードバックすることで、より長い未来のシーケンスを生成する反復的自己回帰推論を適用する。

実験結果

リサーチクエスチョン

• RQ1 並列化可能なアテンション機構が、性能を維持または向上させながら、再帰的ユニットに代わって時空間予測に使用可能か？
• RQ2 時系列アテンションを静的（フレーム内）および動的（フレーム間）コンponentに分解することで、長距離依存関係のモデリングが向上するか？
• RQ3 微分発散正則化により、ピxlsレベルの再構成を超えたフレーム間の運動ダイナミクスを捉えることで、一般化性能が向上するか？
• RQ4 再帰的ベースラインと比較して、提案手法の学習速度および推論効率はどのようにスケーリングするか？
• RQ5 あるデータセットで学習したモデルが、異なるデータセット（例：KITTIからCaltech Pedestrian）でテストされた場合に、一般化性能を示せるか？

主な発見

• 提案されたTAUベースのモデルは、KTH（10→20フレーム）で0.911のSSIMと34.13のPSNRを達成し、SimVP や E3D-LSTM といった先行手法を上回る最先端の性能を示した。
• Moving MNISTデータセットでは、MSEが35.0に収束するまでにわずか50エポックで学習が完了し、CrevNet（1エポックあたり30分）やPhyDNet（1エポックあたり7分）と比較して著しく高速で、1台のV100で1エポックあたり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 Pedestrianでテストした場合でも、強い性能を維持した。これは、強力な特徴抽出能力を示している。
• 反復的自己回帰推論により、40フレームにわたる長い未来のシーケンスを生成でき、品質の著しい劣化なく一貫性を保った。