[論文レビュー] SwinLSTM:Improving Spatiotemporal Prediction Accuracy using Swin Transformer and LSTM

Song Tang, Chuang Li+2 more|arXiv (Cornell University)|Aug 19, 2023

Human Mobility and Location-Based Analysis被引用数 4

ひとこと要約

本稿では、Swin Transformerブロックを用いてConvLSTMの畳み込み演算を置き換えることで、時空間データにおけるグローバルな空間的依存関係をより良く捉える、新しい再帰セルSwinLSTMを提案する。Swin Transformerの自己注意メカニズムを簡素化されたLSTMと統合することで、特別な訓練テクニックを用いずに、Moving MNIST、TaxiBJ、Human3.6m、KTHデータセットにおいて、予測精度が顕著に向上し、SOTAの性能を達成した。

ABSTRACT

Integrating CNNs and RNNs to capture spatiotemporal dependencies is a prevalent strategy for spatiotemporal prediction tasks. However, the property of CNNs to learn local spatial information decreases their efficiency in capturing spatiotemporal dependencies, thereby limiting their prediction accuracy. In this paper, we propose a new recurrent cell, SwinLSTM, which integrates Swin Transformer blocks and the simplified LSTM, an extension that replaces the convolutional structure in ConvLSTM with the self-attention mechanism. Furthermore, we construct a network with SwinLSTM cell as the core for spatiotemporal prediction. Without using unique tricks, SwinLSTM outperforms state-of-the-art methods on Moving MNIST, Human3.6m, TaxiBJ, and KTH datasets. In particular, it exhibits a significant improvement in prediction accuracy compared to ConvLSTM. Our competitive experimental results demonstrate that learning global spatial dependencies is more advantageous for models to capture spatiotemporal dependencies. We hope that SwinLSTM can serve as a solid baseline to promote the advancement of spatiotemporal prediction accuracy. The codes are publicly available at https://github.com/SongTang-x/SwinLSTM.

研究の動機と目的

• 畳み込みニューラルネットワーク（CNN）ベースのモデルが局所的な受容野を持つため、長距離の空間的依存関係を捉えるのが困難であるという制限を解消すること。
• LSTM内の畳み込み演算を自己注意メカニズムに置き換えることで、時空間予測の精度を向上させること。
• 時系列データにおける空間的・時間的ダイナミクスを効果的にモデル化できる、新しい再帰セルSwinLSTMを開発すること。
• 今後の時空間予測研究のための強力で汎用性の高いベースラインを確立すること。

提案手法

• SwinLSTMは、ConvLSTM内の畳み込み演算をSwin Transformerブロックに置き換えることで、2次元特徴マップにおけるグローバルな空間的自己注意をモデル化する。
• 入力フレームは重複のないパッチに分割され、学習可能なトークンに埋め込まれ、注意計算が行われる。
• 時間的記憶を維持するために簡素化されたLSTMユニットを用い、入力ゲート、忘れゲート、出力ゲートはSwin Transformerブロックの出力と連携して動作するように調整されている。
• Swin Transformerブロックは、シフトされたウィンドウ自己注意を採用することで、グローバルな依存関係を効率的に計算しつつ、計算効率を維持する。
• ネットワークアーキテクチャでは、エンコーダーとデコーダーステージ間で特徴マップをダウンサンプリングおよびアップサンプリングするために、パッチマージングおよびパッチ拡張レイヤーを用いる。
• 再構成ヘッドは、転置畳み込み、バイリニア補間、または線形投影を用い、最終的な隠れ状態から予測未来フレームを生成する。

実験結果

リサーチクエスチョン

• RQ1 LSTM内の畳み込み演算を自己注意メカニズムに置き換えることで、時空間予測精度が向上するか？
• RQ2 Swin Transformerを用いてグローバルな空間的依存関係をモデル化することで、再帰モデルの時空間タスクにおける性能にどのような影響を与えるか？
• RQ3 多様なデータセットにわたる一般化を最適化するためのパッチサイズとSwin Transformerブロック数の最適な設定は何か？
• RQ4 提案されたSwinLSTMアーキテクチャは、さまざまな時空間予測ベンチマークで一般化性を示すか？
• RQ5 SwinLSTMの隠れ状態およびセル状態は、空間的運動や軌道情報をどのように符号化しているか？

主な発見

• Moving MNISTデータセットでは、SwinLSTMが最も優れた性能を示し、テストMSEが0.390、SSIMが0.980に達し、ConvLSTMや他のSOTA手法を上回った。
• Human3.6mデータセットでは、SwinLSTMがテストMSE 33.2（10で除算済み）およびSSIM 0.913を達成し、ConvLSTMを顕著に上回った。
• アブレーションスタディの結果、TaxiBJおよびHuman3.6mでは、転置畳み込みによる再構成が、バイリニア補間や線形投影を上回る最良の結果を示した。
• 最適なパッチサイズはデータセットによって異なる：Human3.6mではパッチサイズ4が最良の性能を示し、TaxiBJではパッチサイズ2が最良だった。
• Swin Transformerブロック数（STB）は性能に非単調な影響を及ぼし、TaxiBJおよびHuman3.6mの両方で18ブロックが最良の結果をもたらした。
• 特徴マップの可視化により、セル状態が数字の軌道を記憶しているのに対し、隠れ状態は現在の位置を符号化しており、より深いSwinブロックは段階的にグローバルな空間相関を学習していることが確認された。