【DL輪読会】Looped World Models

DEEP LEARNING JP [DL Papers] Looped World Models Yusei Koen, Matsuo Lab http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 題名 Looped World Models 著者 Hongyuan Adam Lu, Z.L., Victor Wei （FaceMind Research Asia） リンク https://arxiv.org/pdf/2606.18208 2

• https://arxiv.org/pdf/2606.18208

背景：World Models • World Modelは、現在の状態 st から行動 at をとった時に 次の状態 st+1 がどうなるかを予測するモデル • 既存のWorld Modelの抱えるトレードオフ – 長期予測を行うと誤差が蓄積していく – 大きいモデルを使用すると誤差は小さくなるが推論コスト増加 • 既存のWorld Modelは、全ての遷移に同じ計算コストを課す – 「あまり状態が変わらない遷移」と「状態が複雑に変化する遷移」に 同じリソースを割くのは非効率的 3

背景：Looped Transformer • 言語モデルの研究においてLooped Transformerが提 案されてきた – Latent stateを複数回同じtransformer blockに通す – 近年はLLMにLoop構造を取り入れることで パラメータ効率が向上することが確認されている • これらのLooped Transformerの研究は、 言語モデルの研究においてのみ評価されてきた Transformer Transformer Loop Transformer • World ModelにLooped Transformerを使えないか？ – Loop構造を取り入れることで、遷移の難易度によって計算量 を可変にできる – 固定の深さのモデルが抱えるトレードオフを解決できるのでは 4

提案手法：Looped World Model • 環境遷移予測のモデルに初めてLooped Transformerを適用 • Loop構造を適用することで、 少ないパラメータサイズで大規模モデル と同等の性能を実現 • LoopのEarly Exit機能により、 遷移の難易度によって動的にLoop回数を 変更することが可能 この論文は具体的な実装、実験の説明、 それぞれの要素のアブレーションが かなり不足しており、よく分からないところ が多い。 5

提案手法：Looped World Model • 基本的な forward の流れ 1. 外界の時刻 k の観測 ok と行動 ak をencodeする 2. 前の時刻 k-1 の隠れ状態と現在時刻の入力をLooped Transformerに入力する • Looped Transformerは T 回ループして処理を行う 3. 出力された隠れ状態から予測ヘッドを用いて、次の時刻の観測、報酬、タスクの終端を予測 6

提案手法：Looped World Model • 具体的な実装 1. Prelude • Transformer Blockの入力にLayerNormをかけることで、入力信号が過度に大きくなることを防ぐ 2. Recurrent Block • • Transformer LayerをT回Loopする 初期の隠れ状態を として、t = 0, 1, …, T-1回の更新式を previous hidden state • • • input Transformer の出力 と定義する。 AとBは dxd 次元の学習可能な行列で隠れ状態と入力をどれだけ次のループに注入するかを学習する この定式化自体は先行研究（Parcae: Scaling Laws For Stable Looped Language Models）に 倣っている （e と h はd次元ベクトル？？Transformerの入力なので本当はsequence？） 7

提案手法：Looped World Model • 具体的な実装 3. Spectral Stability Constraint • • 隠れ状態の係数行列を、各要素が0から1の対角行列となるように制約 隠れ状態が発散しないようにする previous hidden state input Transformer の出力 4. Coda • • 最終隠れ状態をヘッドに入れる前に特定の次元に射映 Codaの出力が次の環境時刻の最初のhidden stateとして用いられる 5. Cross-Timestep State Propagation • 環境時刻 k でのループの終端時刻 t=Tでの隠れ状態が、環境時刻 k+1 のループの初期時刻 t=0 の隠れ状 態として用いられる 8

提案手法：Looped World Model • 学習方法 – 学習時にはSequenceごとにループ回数をランダムに変える • 学習可能パラメータを持つポアソン分布からサンプルする – World Model Loss • 各環境時刻での予測誤差で学習する • Loopのbackpropはtruncated backpropを用いる（Loop回数の半分のステップまで） 9

提案手法：Looped World Model • 学習方法 – Entropy-Regularised Adaptive Depth • 各ループ時刻で隠れ状態からearly exitの確率を出力する線形層のGateを学習する • Gateの出力が、0か1に張り付かないようにエントロピー正則化を加える • これを含めて最終的な学習損失は、 となる。 Gateの学習信号は自分の認識だとこのEntropy正則化のみなので、どのようにEarly Exitを学習 しているのかはよく分からなかった 10

提案手法：Looped World Model • Deferred Decodingの学習 – Deferred Decoding: 複数ステップ潜在空間でrolloutして終端 状態のみdecodeする – 最初の観測O0のみ入力して、各ステップでは 行動のみ入力して、潜在空間で予測を行う – 終端時刻でのみdecodeしてその誤差で学習 – 中間の潜在表現が変な方向に発散しないように 正則化を加える 1. h0 hK 潜在空間で予測 Latent Consistency Loss – – 2. OK 各時刻の潜在表現が入力からencodeした潜在表現と一致するように学習 普通のLatent World Model的なLoss Spectral contraction budget – 潜在表現の移動距離に制約をつける – 最終的な損失 O0 h0 hK 移動距離に上限をつける Spectral contraction budget 11

提案手法：Looped World Model • Deferred Decodingの学習 – 学習中用いるKはカリキュラム学習的に増やしていく • 最初から長期の予測を行うのは難しいため OK – このDeferred Decodingと先ほどの1stepごとのWorld Model 学習がどのように組み合わされているのかは論文内で 述べられていない。 • 最初はWorld Model学習をして、その後にDeferred Decodingの学習を している？ h0 hK 潜在空間で予測 O0 12

実験設定：ベンチマーク • 使用しているベンチマーク 1. ScienceWorld 2. ALFWorld • 本来はLLM Agentを評価するベンチマーク – タスクがあって、それを実行できるかどうかを 評価する • 本研究では、ある状態からある行動を取った時 にどういう状態になるかを予測させるタスクと して使用 – メインは5stepの行動を入力して、その後の状態を 当てられるかを評価 – アブレーションとして行動のstep数を変化させた時の 予測性能を評価（後述） ScienceWorldの場合のタスクのイメージ （論文内に記述が無いので推測） Input: Initial observation: You are in the kitchen. You see a stove, a metal pot, and a cup of water. The stove is off. Actions: 1. take metal pot 2. pour water into metal pot 3. put metal pot on stove 4. turn on stove 5. wait Target: The metal pot is on the stove. The stove is on. The water in the metal pot is boiling. このモデルがどのように文章を学習・生成しているのかは述べられていない。 隠れ状態に条件づけてTransformer headで自己回帰的に生成している？ 13

実験設定：評価指標 1. EM: Exact Match – 生成された文章と正解文章の完全一致 2. Token-F1 – 生成された文章と正解文章のトークン単位でのF1score 3. BLEU-4 – 生成された文章と正解文章のn-gram単位での一致度 4. Entity – （おそらく）正解文章の環境固有の単語が生成された文章にどの程度含まれているか 14

実験結果①：LLMとの比較 • 約 1B パラメータ の LoopWM を claude-opus-4-6max / gemini-3-flash / qwen-3.5-flash と比較 • ScienceWorld – 全指標で提案手法の性能が高くなっている • ALFWorld – 他のモデルとほぼ同等の性能 100倍小さいパラメータ数で大規模モデルと同等の性能 を実現。 → LoopWMがどのようなデータで学習されているかなど の記載が本文中に無いため、フェアに比較できているの か微妙 各評価指標のモデルごとの比較 （あまりにも論文のTableが 見にくかったためTableの数値を元に作成） 15

実験結果②：Deferred Decoding • ScienceWorldにおいて、stepごとに予測性能を 評価 – 上のTable：LoopWMの性能 – 下のTable：Geminiの性能で正規化したLoopWMの性能 • どの評価指標においても、stepが増えても性能 は劣化していない • どのstepにおいてもGeminiよりも高い性能 • Stepが進むほどGeminiとの差が広がる傾向 → 誤差蓄積が抑えられている 16

まとめ・所感 まとめ • 初めてWorld ModelにLooped Transformer を適用した研究 所感 • 画像モダリティでの実験がない – World ModelにLooped Transformerを適用 したと言いつつ評価は言語タスク • Loop構造によるパラメータ効率の向上 → 100倍少ないパラメータ数でLLMと同等 性能 • 比較対象モデルがLLMのみ • Early ExitによるAdaptiveなLoopによって、 計算効率を向上 • アブレーションの不足 – LoopなしのTransformerとの比較が無いので 良くなっているのか分からない – Early ExitやLoopの詳細な実験が 行われていない 17