技術概要：World Machine：時系列のための生成的世界モデル化に向けて

1. 問題定義

現在の生成 AI モデル、特に Transformer ベースのモデルは、環境の予測的理解と制御可能なシミュレーションを可能にする真の「世界モデル」の構築において、重大な課題に直面しています。従来の Transformer は、コンテキスト長に対して二次的な計算コストとメモリコストを抱えており、長い時系列のスケーラビリティを制限しています。さらに、既存のアプローチは多くの場合、手作りのルールに依存するか、観測データを生成する内部の潜在状態を明示的にモデル化することなく、感覚ストリーム間の相関関係のみに焦点を当てています。以下のようなアーキテクチャが必要です：

• 観測データ量やコンテキストの多様性への適応。
• 感覚的な相関関係だけでなく、内部の世界状態（潜在表現）の明示的なモデル化。
• 完全な履歴に依存することなく、切り詰められたコンテキスト（例：単一の直前の状態）からの予測を可能にすることで推論コストを削減。
• 異種データおよび複数の感覚チャネル（多モーダルである可能性を含む）の処理。

2. 手法

2.1. World Machine アーキテクチャ

World Machine は、**潜在世界状態（$ws$）**上で動作するデコーダ専用 Transformer に基づく、生成式・多モーダル・時系列モデル化アーキテクチャです。

• 潜在世界状態（$ws$）： これらはシミュレートされた世界の現在の構成を符号化する実数ベクトルであり、標準的な Transformer の埋め込みと類似していますが、固定されたシミュレーション周波数で時間とともに進化します。
• 感覚チャネル： モデルは特定のチャネルを介して入力および出力感覚データを処理します。入力チャネルは事前エンコーダを通過し、モデルは $ws$ をデコードして出力感覚データを生成します。
• コア： シミュレーションエンジンは Transformer ブロックで構成されます。これは、感覚データに基づいて潜在状態をスケーリングおよびシフトさせることで条件付ける感覚ブロック（特に adaLn-Zero 型）を利用します。あるステップで感覚データが欠落している場合、ブロックはこれらの操作をスキップします。
• 位置符号化： モデルは正弦波符号化の代わりにALiBi（Attention with Linear Biases）を採用しています。これにより、無制限のシーケンスサイズが可能となり、位置情報が潜在状態の埋め込みを乱すことなく、直接アテンション行列に注入されます。
• 状態活性化： モデルの安定性を確保するため、状態出力に tanh 活性化が適用されます。

2.2. 訓練プロトコル：状態発見

主要な課題は、コアが真実のラベルとして提供されない潜在状態（$ws$）上で動作することです。これを解決するため、著者は状態発見プロセスを提案します：

• 初期化： 潜在状態は一様分布で初期化されます。
• 反復更新： 訓練中、コアからの予測状態（$\hat{ws}$）は時間的にシフト（「左シフト」演算子を使用）され、次のエポックの入力としてフィードバックされます。各シーケンスの最初の状態は、ヌルベクトルにリセットされます。
• 最適化： モデルは、出力次元全体にわたる損失の加重和を使用して、次の時間ステップの感覚データ（次のトークン予測に相当）を予測するように訓練されます。これにより、$ws$ が感覚データを再構築するために必要な情報を符号化することが強制されます。

2.3. 強化された訓練プロトコル

表現の質と堅牢性を向上させるため、訓練プロトコルにはいくつかの具体的なステップが含まれます：

• 感覚マスク： 訓練中に感覚データをランダムに隠蔽し、データが欠落している場合にモデルが潜在状態に依存することを強制します。
• シーケンス分割： シーケンスをランダムなセグメントに分割します。モデルは先行データなしでシーケンス中間の状態を処理する必要があります。「高速進行」メカニズムは、あるセグメントの最後の状態を次のセグメントの最初の状態にシフトして、情報勾配を伝播させます。
• ノイズ付加： $ws$ および/または感覚次元にガウスノイズを追加します。
• 短時間想起： 現在のステップに対する過去または未来のデータを表す合成チャネルを作成し、時間的コンテキストを改善します。
• ローカルモード： セグメントに対してアテンション機構をランダムに無効化し、モデルが現在の $ws$ と直近の感覚データのみ（RNN の隠れ状態処理に類似）に依存することを強制します。

2.4. 評価フレームワーク

著者は、ランダムな初期状態とインパルス/方形波入力を持つ減衰物理系を表す合成 1 次元時系列であるToy1Dデータセットを導入しました。評価タスクには以下が含まれます：

• 浅い予測（Prediction Shallow）： 最後の符号化状態のみ（感覚データなし）を使用して、シーケンスの未来 50% を推論します。これはコンテキストの切り詰め能力をテストします。
• 状態使用（Use State）： 感覚入力なしで、以前に符号化された状態から出力を生成します。
• 感覚マスク（Mask Sensory）： ランダムな割合で感覚データがマスクされたシーケンスを生成します。
• 影響テスト： 特定のハイパーパラメータおよびプロトコルステップが性能と訓練時間に与える影響を定量化するための統計的手法（ウィルコクソンの符号順位検定を使用）。

3. 主要な貢献

1. 新規アーキテクチャ： 経験を高次マルコフ過程として扱い、感覚的な相関関係だけでなく内部の潜在状態（$ws$）を明示的にモデル化する、統合されたエンドツーエンドアーキテクチャである World Machine の導入。
2. 状態発見プロトコル： 真実の状態ラベルなしで潜在状態の符号化と進化を学習することを可能にする訓練メカニズム。これにより、ヌルベクトルからの自己回帰的推論が可能になります。
3. コンテキスト切り詰め能力： 切り詰められたコンテキスト（特に「浅い予測」タスク）での推論を実行する能力の実証。これは、標準的な Transformer に伴う二次的なコストを削減する可能性のある道筋を提供します。
4. 包括的な影響分析： 性能とコストへの具体的な貢献を決定するために、影響テストフレームワークを用いた訓練プロトコル構成要素（マスク、シーケンス分割、ローカルモードなど）の体系的評価。
5. オープンソースベースライン： 最近の大規模世界モデルの不可視性に対処するため、アーキテクチャ、訓練プロトコル、および合成データセットの公開。

4. 結果

Toy1D データセットを用いて、さまざまなモデル構成と訓練プロトコルで実験が行われました。

• 訓練可能性： 損失曲線の減少が示すように、状態発見技術はモデルの訓練を成功裏に可能にしました。
• タスク性能：
  • モデルは「通常」の自己回帰的推論で最も良好なパフォーマンスを発揮しました。
  • ベースモデル（状態発見のみを使用）は、感覚データの欠如を必要とするタスクを実行できませんでした。
  • 完全プロトコル（マスク、シーケンス分割、ノイズなどを組み合わせたもの）は、特に浅い予測タスクにおいて顕著な改善を示しましたが、これは依然として最も困難なタスクでした。
  • 定性的な結果は、特に感覚データが利用可能な場合に整合性のある予測を示しましたが、高周波数のシーケンスは依然として困難でした。
• プロトコルの影響：
  • 安定性： tanh 状態活性化の使用は決定的でした。これを削除し（MSE 正則化に置き換え）、除外すると発散率が約 20% になりましたが、保持することで発散を 1% 未満に削減しました。
  • 性能： 「浅い予測」タスクは最も大きなばらつきを示しました。最良の実証的構成（シーケンス分割、特定のブロック構成、ノイズ付加、ローカルモードを組み合わせる）は、1,000 エポック後に0.3435の浅い予測 MSE を達成しました。
  • コスト/ベネフィット： ローカルモード変数は、高価なアテンション機構をバイパスするため、性能と訓練時間の両方にプラスの影響を示しました。
  • 相関： 「状態使用」と「浅い予測」タスクは、他のタスクとの相関が低く、これらは世界モデル化能力のためのユニークかつ必要な評価指標であることを示唆しています。

5. 意義と主張

本論文は、予測的理解と制御可能なシミュレーションのバランスを取る生成的世界モデル化への基礎的なステップとして World Machine を位置付けています。

• 効率性： 単一の潜在状態（浅い予測）から未来の状態を推論する能力は、現在の Transformer の二次スケーリングの問題を緩和する可能性を示唆しており、長期的な計画にとってより費用対効果の高いアプローチを提供します。
• 適応性： このアーキテクチャは、欠落データや多様な感覚入力を処理するように設計されており、完全な観測が常に可能ではない複雑で動的な環境に適しています。
• 透明性： 完全に再現可能なオープンソースベースラインを公開することで、著者は大規模な専有システムにおいてしばしば不可視である世界モデルのメカニズムを解明することを目指しています。
• 将来の可能性： 著者は、このアプローチが多モーダル問題（例：視覚慣性オドメトリ）や大規模言語モデル（LLM）に適用され、推論コストを削減し、品質とコストのトレードオフに対する制御を改善する可能性を指摘しています。また、実世界の多モーダルデータセットでのテストや、データ生成および環境コストに関する倫理的懸念への対応といった、今後の研究の必要性も認識しています。

本論文は、合成データにおける現在の結果は有望である一方、アーキテクチャは訓練可能であり、整合性のある予測を生成する能力を有しており、提案された状態発見および訓練プロトコルの実現可能性を検証していると結論付けています。