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つの分析を以下の手順で実施した:
- A. Attention パターン分析: 全100入力に対する各層・各ヘッドの attention weight を取得し、全体平均・壁際/内側比較・行動別に可視化
- B. 中間表現の幾何学: 各層の state_token 出力ベクトル(100x64次元)に PCA を適用し2次元に射影。隣接セルと非隣接セルの距離比で空間構造の保存度を定量化
- C. Linear Probing: 各層の中間表現から、座標・壁接触・マンハッタン距離・次状態を線形モデルで読み出し、R2スコア/精度で評価
- D. ニューロンレベル分析: 64個の各ニューロンについて、壁際/内側の活性差、行動選択性、ゴール距離との相関を計算し、機能特化ニューロンを同定
- 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(収束済み)