2. 研究手法

2.1 世界モデルのアーキテクチャ

項目 備考
モデル TransformerWorldModel Encoder-only
d_model 64 埋め込み次元
ヘッド数 4 16次元/ヘッド
層数 2 TransformerEncoderLayer
入力 [state_token, action_token] 2トークン固定長
出力 次状態 + 報酬 state_token から予測
パラメータ数 108,931 全パラメータ

入力は状態(正規化座標 [y, x] ∈ [0,1]^2)を Linear 層で、行動(0~3の整数)を Embedding 層でそれぞれ d_model 次元に射影し、2トークン列として Transformer に入力する。state_token の最終出力から次状態(Sigmoid 活性化で [0,1] に制約)と報酬を予測する。

2.2 環境

項目 備考
環境 GridWorldEnv 5x5 グリッド
状態空間 25 状態 (y,x) 各0~4
行動空間 4 行動 上(0), 下(1), 左(2), 右(3)
報酬 +1.0 (ゴール), -0.01 (その他) ゴール=(4,4)
壁の処理 クリッピング 壁に向かうと位置不変

2.3 分析手法

5つの分析を以下の手順で実施した:

  1. A. Attention パターン分析: 全100入力に対する各層・各ヘッドの attention weight を取得し、全体平均・壁際/内側比較・行動別に可視化
  2. B. 中間表現の幾何学: 各層の state_token 出力ベクトル(100x64次元)に PCA を適用し2次元に射影。隣接セルと非隣接セルの距離比で空間構造の保存度を定量化
  3. C. Linear Probing: 各層の中間表現から、座標・壁接触・マンハッタン距離・次状態を線形モデルで読み出し、R2スコア/精度で評価
  4. D. ニューロンレベル分析: 64個の各ニューロンについて、壁際/内側の活性差、行動選択性、ゴール距離との相関を計算し、機能特化ニューロンを同定
  5. E. 因果的介入: 各ニューロンをゼロに設定し、予測変化量(MSE)で因果的重要度を測定。Activation Patching(層間のトークンシャッフル)も実施

2.4 評価指標

  • R2スコア: 線形回帰の決定係数。1.0で完全な線形読み出しが可能
  • 分類精度(Accuracy): 壁接触の2値分類における正答率
  • 距離比(非隣接/隣接): PCA空間における隣接セル間距離と非隣接セル間距離の比。>1.5で空間構造ありと判定
  • 因果的影響度(MSE): ニューロン除去後の予測変化量

2.5 実験環境

  • ソフトウェア: Python, PyTorch (CUDA), NumPy, Matplotlib
  • ハードウェア: CUDA対応GPU
  • 学習設定: 8000ステップのランダム行動データ、50エポック、バッチサイズ128、学習率1e-3
  • 学習結果: 最終 Loss=0.0001(収束済み)