技術要約: QQWorld

問題提起

潜在世界モデル（Latent world models）は、コンパクトな表現空間内で将来の状態を予測することにより、効率的なプランニングを可能にする。LeWorldModel (LeWM) などの最近のアーキテクチャにおける重要な設計上の選択肢は、学習された潜在変数の周辺分布を等方的なガウス分布へと正則化することである。これは通常、正規性を強制するための微分可能なペナルティとして、Epps–Pulley (EP) テストを用いて達成される。

しかし、著者らは、EP正則化を行っているにもかかわらず、LeWMが顕著なヘビーテイル（重い裾）を持つ潜在分布を学習してしまう現象を観察している。これは、極端な潜在値が学習されたダイナミクスを、表現が不十分な領域へと押し出す可能性があり、マルチステップのロールアウトにおける誤差を増幅させ、ターゲットとなるガウス事前分布との不一致を増大させるため、問題となる。

本論文は、EPテストの最適化幾何学における根本的な限界を特定している：

1. 勾配消失: EP統計量は、ガウスカーネルを用いた二乗最大平均偏差（MMD）と等価である。著者らは、この目的関数によって誘導される補正勾配が、サンプル（外れ値）において急速に（超指数関数的に）減衰することを証明している。
2. 効果的な制御の欠如: 潜在座標がガウスカーネルの相互作用スケールを超えると、EP正則化による復元力は無視できるものとなる。その結果、ヘビーテイルの偏差は十分に制御されず、モデルが非正規性に対して明示的にペナルティを課しているにもかかわらず、それらが存続し成長することを許してしまう。

手法

この勾配消失の問題に対処するため、著者らはEP正則化をQuantile–Quantile (QQ) マッチング目的関数に置き換えたQQWorldを提案する。

コアメカニズム: QQ正則化

モーメントのマッチングやカーネルベースの不一致を用いる代わりに、QQWorldは投影された潜在サンプルを、ランク一致させたガウス分位点に直接整合させる。

• 目的関数: 投影されたサンプルの集合 $X = \{x_n\}$ に対して、損失関数は、サンプルの順序統計量 ($\hat{x}_n$) と、対応する理論的なガウス分位点 ($q_n = \Phi^{-1}(\frac{n-0.5}{N})$) との間の二乗差の総和として定義される：
  $$L_{QQ}(X) = \sum_{n=1}^N (\hat{x}_n - q_n)^2$$
• 勾配の挙動: サンプル $x_n$ に対するQQ損失の勾配は $2(x_n - q_{\rho(n)})$ であり、ここで $\rho(n)$ はランクである。EPの勾配とは異なり、この勾配の大きさはターゲットとなる分位点からの偏差に応じて線形に増加する。これにより、テイルのサンプルに対して強力で対称的な補正力が働き、ヘビーテイルを効果的に抑制することが保証される。
• 理論的保証: 著者らは、QQ損失をゼロに駆動することがEP損失をゼロに駆動することも証明している（一方向の制御）。しかし、その逆は成立しない。小さなEP損失は、極端に大きなテイルの偏差を伴って共存し得るが、QQマッチングはより厳格な分布の整合性を強制する。

クロスバッチQQ (Cross-Batch QQ)

ランク推定の精度を高めるために必要な大きなバッチサイズに伴うメモリ制約に対処するため、著者らはCross-Batch QQを導入している。

• 戦略: この手法は、前 $K$ 回のイテレーションからデタッチ（detach）された特徴量をFIFOキューに含めることで、実効的なランキングプールを拡大する。
• メカニズム: ランクはサイズ $M = N(K+1)$ のプール集合に対して計算されるが、勾配は現在のバッチサイズ $N$ に対してのみバックプロパゲーションされる。
• バイアスと分散のトレードオフ: 著者らは、キューの長さ $K$ を増やすことでランク推定の分散は減少するが（小規模なバッチにおいて有益）、潜在分布が時間とともに大きく変化する場合、「表現の陳腐化（representation-staleness）」というバイアスが生じるというトレードオフを特性化している。

主な貢献

1. EPの限界の特定: 本論文は、EP正則化の補正力がテイルのサンプルに対して消失することを形式的に分析・実証し、なぜLeWMが明示的な正規性ペナルティを与えているにもかかわらずヘビーテイルを抑制できないのかを説明した。
2. QQWorldの提案: 新しいハイパーパラメータを導入することなく、分位点・分位点マッチング目的関数を利用する、LeWMのドロップイン・リプレースメントとしてのQQWorldを提案した。
3. メモリ効率の高い学習: 正確な分位点推定を行うために大きな実効バッチサイズを使用しながら、バックプロパゲーションのバッチサイズとランキングプールのサイズを切り離すことで、GPUメモリ使用量を低く抑えるCross-Batch QQを開発した。
4. 理論的分析: Cross-Batch QQにおけるバイアスと分散のトレードオフの特性化、およびQQとEPの目的関数の間の、一方向の制御関係の証明を行った。

実験結果

著者らは、Two-Room、PushT、Reacher、OGBench-Cubeの4つの制御環境においてQQWorldを評価した。

• プランニング性能: QQWorldは平均プランニング成功率**85.08%**を達成し、LeWMのベースライン（79.75%）を5.33ポイント上回った。また、4つの環境すべてにおいて、LeWMおよび他の最近のバリアント（Sub-JEPA、SD-JEPA、SMWM）を一貫して上回った。
• 分布の整合性:
  • QQ RMSE: 実験的なガウス分位点と理論的な分位点の平均誤差を、0.157（LeWM）から0.121（QQWorld）へと減少させた。
  • テイル確率: サンプルがカイ二乗分布の90パーセンタイルを超える確率（径方向のテイルの重さの指標）を、0.315から0.123へと減少させ、ガウス基準の0.10に近づけた。
  • 統計テスト: LeWMはEP目的関数で直接訓練されているにもかかわらず、QQWorldはLeWMと比較してコルモゴロフ–スミルノフ（KS）統計量を15.8%、EP統計量を31.4%減少させた。
• 物理状態のプロービング: PushT環境において、QQWorldはエージェントの位置、ブロックの位置、ブロックの角度といった物理状態のデコードにおいてLeWMと同等またはわずかに優れた性能を示した。これは、強力な正則化がタスクに関連する情報のエンコーディングを劣化させていないことを示している。
• クロスバッチの効率性: 小さなバッチサイズ（$N=32$）とキューの長さ（$K+1=3$）を用いた場合でも、QQWorldは83.50%の成功率を達成し、バッチサイズ128で訓練されたLeWM（79.75%）を上回りつつ、GPUメモリ使用量を73%削減した。

意義と主張

本論文は、統計的な分布不一致の測定におけるテストの有効性が、必ずしも学習目的関数としての有効性を保証するわけではないと主張している。具体的には、以下の点を挙げている：

• 勾配の質が重要: 分布を区別するだけでなく、正則化目的関数は最適化プロセス全体を通じて、情報量が多く、挙動の良好な勾配を提供しなければならない。EPテストはテイルのサンプルに対してはこの基準を満たさないが、QQ目的関数はこれを満たす。
• テイル制御の重要性: 潜在空間におけるヘビーテイルの効果的な抑制は、世界モデルにおける安定したマルチステップ・プランニングに不可欠である。
• スケーラビリティ: Cross-Batch QQ戦略は、物理的なバッチサイズを増やすことがメモリの観点から困難な、大規模な世界モデルを訓練するためのスケーラブルなメカニズムを提供する。

結論として、QQWorldは、分布の不一致を測定するための統計的テストの最適化幾何学に直接対処することで、既存の潜在世界モデルに対する、シンプルかつハイパーパラメータフリーな改善を提供するものである。