技術要約：SUNTA – 驚愕（Surprise）に基づくチャンク分割を用いた階層的ビデオ予測

問題提起

階層的状態空間モデル（HSSM）は、シーケンスを時間的なチャンク（塊）に分解することで、複数のタイムスケールにわたるダイナミクスのモデリングを可能にし、長期的な予測を実現するための有望なフレームワークを提供します。しかし、HSSMの有効性は、チャンクの境界がどのように決定されるかに決定的に依存しています。既存のアプローチは通常、以下のいずれかに依存しています：

1. 固定長チャンキング： 決定論的な間隔であり、データの固有の時系列構造と一致しないことがよくあります。
2. 類似性ベースの検出： 観測の変化（例：視覚的な変化）から境界を推論します。この手法は、顕著な視覚的変化がないまま重大な意味論的遷移が発生した場合や、潜在的な動的変化がないにもかかわらず表面的な視覚的変化が発生した場合に機能しません。

著者らは、チャンキングは内部世界モデルにおける**予測誤差（驚愕/Surprise）**によって駆動されるべきであると主張しています。なぜなら、これらの誤差は、現在のコンテキストが不十分であり、より高次の抽象化が必要であることを直接的に示すからです。しかし、驚愕に基づくチャンキングをHSSMに統合することには、2つの重要な課題があります：

• 階層的崩壊（Hierarchical Collapse）： エンドツーエンドの学習において、高次モデルが改善されるにつれて、低次モデルをより効果的に条件付けるようになるため、低次モデルの予測誤差が減少します。これにより、チャンク境界を定義するために必要な驚愕信号そのものが消失し、階層が崩壊してしまいます。
• オープンループ生成における驚愕の欠如： 驚愕信号は通常、予測を外部の観測と比較することを必要とします。オープンループのロールアウト（生成）中には観測が利用できないため、チャンクの切り替えをトリガーするために必要なボトムアップの驚愕を計算することが不可能です。

手法：SUNTA

著者らは、デカップリングされた学習戦略とトップダウンの驚愕メカニズムを通じて、これらのギャップに対処するために設計された2レベルのHSSMであるSurprise-based Nested Temporal Abstraction (SUNTA) を提案します。

1. モデルアーキテクチャ

SUNTAは、異なるタイムスケールで動作する2つのレベルで構成されています：

• レベル1（低次レベル）： 再帰的状態空間モデル（RSSM）に基づいています。生の観測値 ($x_t$) を処理して、確率的な潜在状態 ($s^{(1)}_t$) を推論します。これは、次の観測と潜在状態を予測することを学習します。
• レベル2（高次レベル）： 学習された低次の潜在シーケンスを「観測」として扱います。境界 $m_t$ によって定義されたチャンク内で動作します。チャンク内では、高次の状態 ($s^{(2)}_t$) は一定に保たれ、低次のシーケンスを単一の表現へと圧縮します。これには、高次の状態から低次の潜在シーケンスを再構成するトップダウン・デコーダが含まれます。

2. 主要なメカニズム

驚愕に基づくチャンキング

チャンクの境界は、低次モデルにおける予測誤差（ベイズ的驚愕）の大きさによって決定され、事後分布と事前分布の間のカルバック・ライブラー（KL）ダイバージェンスとして定量化されます。

• 推論（認識）： $DKL(q(s^{(1)}_t | \dots) \parallel p(s^{(1)}_t | \dots))$ が閾値 $\tau_{inf}$ を超えたときに境界がトリガーされます。
• 生成（オープンループ）： 「驚愕の欠如」問題を解決するために、SUNTAはトップダウンの驚愕信号を導入しています。外部の観測と比較する代わりに、モデルは（トップダウン・デコーダによって生成された）進化する低次のロールアウトを、現在の高次のコンテキストと比較します。生成された低次の軌跡と高次のコンテキストとの不一致が閾値 $\tau_{gen}$ を超えたときに、境界がトリガーされます。

デカップリングされた学習（レベルごとの最適化）

階層的崩壊を防ぐために、SUNTAはエンドツーエンドの最適化ではなく、逐次的な学習戦略を採用しています：

1. レベル1の学習： 生の観測に対して証拠下限（ELBO）を最大化するように低次RSSMを最適化します。
2. レベル1の凍結： 低次の潜在シーケンスとチャンク境界を推論します。
3. レベル2の学習： 推論された低次の潜在変数を観測として扱い、高次モデルを学習します。
   これにより、高次モデルの改善が、境界検出に使用される低次の驚愕信号を抑制してしまうことを防ぎます。

時間的パターン補完（Temporal Pattern Completion: TPC）

高次エンコーダが学習時には完全なチャンクを見る一方で、推論時には部分的なチャンクしか見ることができないという、学習とテストのミスマッチに対処するため、SUNTAは正則化項を導入しています。この損失は、部分的なチャンクのプレフィックス（接頭辞）から推論された高次表現が、完全なチャンクから推論された表現と一致するように促すと同時に、負の（無関係な）コンテキストからは分離を維持するように促します。これにより、生成中の不完全なコンテキストに対する堅牢性が向上します。

主な貢献

1. 驚愕駆動型チャンキング： 視覚的な類似性ではなく、内部の予測誤差を使用して時間的抽象化の境界を決定することを提案し、それが内在的なダイナミクスとより良く一致することを主張しました。
2. デカップリングされた学習戦略： 驚愕信号の抑制を防ぐレベルごとの最適化アプローチを導入し、階層的崩壊の問題を解決しました。
3. 生成のためのトップダウンの驚愕： 外部の観測なしに、高次のコンテキストと低次のロールアウトの間の不一致を測定することで、オープンループ生成中のチャンク境界を検出するメカニックを開発しました。
4. 時間的パターン補完： 高次表現の学習時（完全なコンテキスト）と推論時（部分的なコンテキスト）の間のギャップを埋めるための正則化技術を提供しました。

実験結果

SUNTAは、Bouncing Ball（合成2D）、Serial Nine Rooms（遅延依存性のある2Dナビゲーション）、3D Maze（一人称視点ナビゲーション）の3つのデータセットで評価されました。

• 長期予測： SUNTAはすべてのベースライン（RSSM、VTA、LOVE、VPR）を大幅に上回りました。
  • Bouncing Ballデータセットにおいて、SUNTAは250タイムステップにわたって正確な予測を維持しましたが、すべてのベースラインは最初の10タイムステップ以内に劣化しました。
  • 3D MazeおよびSerial Nine Roomsにおいて、SUNTAは階層的および非階層的なベースラインと比較して、優れたSSIMと低いMSEを達成しました。
• 境界検出： SUNTAは、グラウンドトゥルースの時間的境界を特定する上で、非常に高いF1スコア（Bouncing Ballで0.984、3D Mazeで0.992）を達成し、類似性ベースの手法（VPR）や他のHSSM（VTA、LOVE）を大幅に上回りました。
• アブレーション研究：
  • TPC正則化を取り除くと、長期的な予測性能が低下しました。
  • 固定タイムスケール・バリアント（CW-SUNTA）は、フラットなRSSMと同程度の性能しか示さず、階層化の利点は階層構造そのものではなく、適応的な驚愕ベースの境界から来るものであることを証明しました。
  • 階層的ベースライン（VTA、LOVE）は、しばしば非階層的なRSSMよりも低い性能を示しました。これは、不適切な階層化が有害になり得ることを示唆しています。

意義と主張

本論文は、SUNTAが環境の固有の時間的構造と階層的抽象化を一致させることが極めて重要であることを示していると主張しています。著者らは以下のように述べています：

• 階層的モデリングの主な利点は、アーキテクチャ自体にあるのではなく、予測誤差が急増するシーケンスを動的にセグメント化することで、高次のダイナミクスの学習を簡素化できる能力にあります。
• 驚愕ベースのチャンキングは、視覚的な類似性よりも、意味論的な遷移が視覚的に顕著でない環境においても、時間的抽象化のためのより信頼できるメカニズムとなります。
• 提案されたデカップリングされた学習とトップダウンの驚愕メカニズムは、既存のHSSMが短期間を超えて維持できない能力である、世界モデルにおける安定した長期予測を可能にします。

著者らは、限界についても言及しており、驚愕のスパイクが存在する環境への依存、現在のシミュレーションされた視覚データへの焦点、および2レベルの階層への限定された評価を挙げています。今後の課題として、より深い階層の探索や意思決定への統合を挙げています。