世界モデルの内部表現の解剖 — Transformer は環境の「地図」を獲得するか
本研究は、5x5 GridWorld の遷移ダイナミクスを学習した小規模 Transformer 世界モデル(d_model=64, 2層)の内部表現を、Mechanistic Interpretability の手法群を用いて網羅的に解剖したものである。入力空間が 25状態 x 4行動 = 100通りと完全に列挙可能であることを活かし、Attention パターン解析・PCA による中間表現の幾何学分析・Linear Probing・ニューロンレベルの機能分析・因果的介入(Causal Intervention)の5つの分析を実施した。
主要な結果
- 空間構造の内部再現: PCA 可視化により、全層でグリッドの空間的トポロジーが中間表現に保存されていることを確認(隣接/非隣接セルの距離比 ~2.9)
- 遷移ルールの段階的獲得: Linear Probing で、入力段階では次状態予測 R2=0.81 だが、層2では R2=0.999 に到達。Transformer が層を通じて遷移関数を計算していることを実証
- 機能特化ニューロンの発見: 「壁際ニューロン」「方向ニューロン」「ゴール距離ニューロン」など、特定の環境概念に対応するニューロンを同定。因果的介入により、これらのニューロンが実際に予測に因果的に寄与していることを確認