[論文レビュー] Convolutional State Space Models for Long-Range Spatiotemporal Modeling

Jimmy T. H. Smith, Shalini De Mello+3 more|arXiv (Cornell University)|Oct 30, 2023

Geophysics and Gravity Measurements被引用数 13

ひとこと要約

ConvSSMsとConvS5を導入する、畳み込み状態空間アプローチで、長距離時空モデリングを並列化し、1ステップあたりの推論が線形時間で行われ、高速な自己回帰生成を実現し、長距離タスクでTransformersとConvLSTMを上回る。

ABSTRACT

Effectively modeling long spatiotemporal sequences is challenging due to the need to model complex spatial correlations and long-range temporal dependencies simultaneously. ConvLSTMs attempt to address this by updating tensor-valued states with recurrent neural networks, but their sequential computation makes them slow to train. In contrast, Transformers can process an entire spatiotemporal sequence, compressed into tokens, in parallel. However, the cost of attention scales quadratically in length, limiting their scalability to longer sequences. Here, we address the challenges of prior methods and introduce convolutional state space models (ConvSSM) that combine the tensor modeling ideas of ConvLSTM with the long sequence modeling approaches of state space methods such as S4 and S5. First, we demonstrate how parallel scans can be applied to convolutional recurrences to achieve subquadratic parallelization and fast autoregressive generation. We then establish an equivalence between the dynamics of ConvSSMs and SSMs, which motivates parameterization and initialization strategies for modeling long-range dependencies. The result is ConvS5, an efficient ConvSSM variant for long-range spatiotemporal modeling. ConvS5 significantly outperforms Transformers and ConvLSTM on a long horizon Moving-MNIST experiment while training 3X faster than ConvLSTM and generating samples 400X faster than Transformers. In addition, ConvS5 matches or exceeds the performance of state-of-the-art methods on challenging DMLab, Minecraft and Habitat prediction benchmarks and enables new directions for modeling long spatiotemporal sequences.

研究の動機と目的

• 長距離時空シーケンスのスケーラブルなモデリングの必要性を動機づけ、局所的な空間構造および長期的な時間依存性を捉える。
• テンソル状態と線形、連続時間ダイナミクスを組み合わせて効率的な並列化を可能にするConvSSMsを開発する。
• 長距離依存性をモデル化するために、ダイアゴナル化されたSSM風ダイナミクスと並列スキャンを用いてConvSSMsをパラメータ化するConvS5を提案する。
• 長期的なデータセットとベンチマークでTransformerとConvLSTMに対する実証的性能向上を示す。

提案手法

• テンソル値の状態と線形畳み込みダイナミクスを持つConvSSMsを定義する。
• 二分結合代数演算子と並列スキャンを用いて sequence 全体にわたる ConvSSM 再帰を並列化する方法を示す。
• ConvSSMsと状態空間モデル（SSMs）との等価性を確立し、HiPPO風の初期化と離散化を活用する。
• ConvSSMsを対角化されたSSM風ダイナミクスでパラメータ化し、長距離モデリングのために並列スキャンを適用してConvS5を提案する。
• 非線形活性化を持つ複数のConvS5層を積み重ね、訓練を現実的に維持しつつ非線形時空ダイナミクスを捕らえる。
• Moving-MNIST、DMLab、Minecraft、Habitatの長距離ベンチマークで実証的証拠を提供し、自己回帰生成を高速化する。

Figure 1 : ConvRNNs [ 17 , 18 ] (left) model spatiotemporal sequences using tensor-valued states, $\mathbfcal{X}_{k}$ , and a nonlinear RNN update, $\mathbf{G}()$ , that uses convolutions instead of matrix-vector multiplications. A position-wise nonlinear function, $\mathbf{h}()$ , transforms the st

実験結果

リサーチクエスチョン

• RQ1 ConvSSMsは長距離時空シーケンスのサブ二乗時間複雑度の並列化された訓練と推論を提供できるか。
• RQ2 ConvSSMsをどのようにパラメータ化・初期化して長距離依存性を効果的に捉えるか。
• RQ3 ConvS5は長距離のビデオ予測ベンチマークでTransformerおよびConvLSTMのベースラインを上回るか。
• RQ4 ConvS5とTransformerおよびConvRNNの比較で計算コスト（速度、メモリ）におけるトレードオフはどうなるか。
• RQ5 ConvS5は品質を保ちながら複雑な3D環境ベンチマーク（DMLab、Minecraft、Habitat）へスケールできるか。

主な発見

• ConvS5はサブ平方次数の複雑度と高速な自己回帰生成を伴う並列シーケンス処理を可能にする。
• ConvS5はConvLSTMより訓練が速く、長距離Moving-MNISTでTransformerよりサンプル生成がはるかに速い。
• ConvS5はDMLab、Minecraft、Habitatの長距離ビデオ予測ベンチマークで最先端と同等またはそれを上回る。
• ConvS5はSSM風パラメータ化とHiPPOベースの初期化を備えたConvSSM構造の恩恵を受ける。
• このアプローチはフレームのトークンベース圧縮に依存せずに高性能を達成する。

Figure 2 : The dynamics of a ConvSSM with pointwise state kernel (top) can be equivalently viewed as the dynamics of an SSM (bottom). See Proposition 3 . Each ConvSSM state pixel evolves according to an SSM state update with shared state matrix, $\mathbf{A}_{\mathrm{SSM}}$ , and input matrix, $\math