LLMにゲームを操作させるなら、普通は操作を文字列として生成させる。例えばSnakeなら、ゲームの状態をプロンプトとして与えて、
{"action":"LEFT"}
のような回答を生成させる。それをプログラム側でパースし、実際のゲーム操作へ変換する。処理としては、
ゲーム状態
↓
LLMへ入力
↓
prefill
↓
decode
↓
"LEFT" を生成
↓
parse
↓
コントローラ操作
となる。
LLMエージェントとしては、ごく普通の構成だ。ただ、ゲームのコントローラを動かしたいだけだと考えると、少し回りくどい。LLM内部ではすでにゲーム状態を処理し、
次は左へ行く
という判断につながる何らかの内部状態を作っているはずだ。それをいったん、
"LEFT"
という言語へ変換し、外側のプログラムでもう一度、
LEFTボタン
へ戻している。つまり、
内部表現
↓
言語
↓
操作
という経路を通っている。では、この真ん中の「言語」は本当に必要なのだろうか。今回試したかったのはそこだった。
decodeを通さずに操作を取り出したい
目標は単純で、
LLMに操作名を生成させず、prefill後の内部状態から直接コントローラ操作を取り出す
ことである。つまり、
ゲーム状態
↓
LLM prefill
↓
hidden state
↓
LEFT / RIGHT / UP / DOWN
としたい。これができれば、LLMは一文字も生成しなくてよい。
ゲームの状態を見せる。
Transformerを通す。
途中の内部状態を読む。
そのままボタンを押す。
LLMを「テキストを生成する装置」としてではなく、
入力から行動に必要な内部表現を作る巨大な計算器
として使うことになる。
decodeを省きたい理由は、速度だけではない
もちろんdecodeを省略できれば、その分の計算コストは減る。しかし今回興味があるのは、単純な高速化だけではない。ゲームには時間の流れがある。例えばSnakeなら、
t=0
Snake: (3, 4)
Food: (8, 4)
t=1
Snake: (4, 4)
Food: (8, 4)
t=2
Snake: (5, 4)
Food: (8, 4)
t=3
...
というように状態が連続して変化する。これを毎tick独立した問い合わせとしてLLMへ投げるのではなく、ゲームの進行そのものをコンテキストへ追加し続けることもできる。
State(t=0)
↓
prefill
State(t=1) を追加
↓
incremental prefill
State(t=2) を追加
↓
incremental prefill
State(t=3) を追加
↓
...
KV Cacheを維持したまま、ゲームの時系列をLLMへ読み込ませ続ける。ここで一つ仮説がある。
ゲームの状況が時間とともに変化すれば、それに対応してLLMの隠れ層も変化していき、その変化の中に「今どの操作をすべきか」という情報が現れるのではないか。
イメージとしては、
game state:
S0 → S1 → S2 → S3 → S4
hidden state:
h0 → h1 → h2 → h3 → h4
という対応である。
そして、
h2 では RIGHT
h4 では LEFT
のように、その時点の内部状態から操作を読めるかもしれない。さらに言えば、絶対的なhidden stateだけではなく、
h(t) - h(t-1)
のような変化量に操作情報が出る可能性もある。この「時系列に沿って変化する内部状態」を使いたいというのが、decodeをできるだけ避けたいもう一つの理由だった。
decodeしてrollbackする方法もある
もちろん、ゲームの時系列をprefillし続けながら操作を取り出す方法は他にもある。例えば、ある時点で推論状態をcheckpointして、
State t
↓
prefill
↓
checkpoint
↓
数token decode
↓
{"action":"LEFT"}
↓
actionを取得
↓
rollback
とすることもできる。自前の推論エンジンを持っていれば、このような処理は比較的自由にできる。ゲームの履歴やKV Cacheは維持しつつ、操作が必要な瞬間だけ一時的にdecodeする。そして操作を取得したら、decode前の状態へ戻る。これはかなり現実的な方法だと思う。LM Headも既存のものをそのまま使える。モデルを追加学習する必要もない。
ただし、毎tickこれを行うとなると、それなりにコストが高い。
actionが欲しいだけなのに文字列を作る?
ゲームが1秒に何回も進行する場合、
checkpoint
↓
decode
↓
parse
↓
rollback
を何度も繰り返すことになる。しかも必要なのはせいぜい、
LEFT
という数bit相当の情報である。それを取得するために、
{
"action":
"LEFT"
}
という文字列を何tokenも生成するのは、どう考えても贅沢だ。もちろんプロンプトを工夫して、
L
の1 tokenだけ出させることもできる。しかしそれでも、
Transformerのdecode step
LM Head
全語彙へのlogits計算
sampling
KVの更新
rollback
などは残る。欲しいものが5種類程度のactionなら、もっと直接取り出せないだろうか。
decodeをどこまで削れるか
実際、actionを取得する方法には何段階かある。
1. JSONを最後まで生成する
2. "LEFT" のような操作名だけ生成する
3. L/R/U/Dなど1 tokenだけ生成する
4. LM Headのlogitsを直接読む
5. hidden stateからactionを直接分類する
6. hidden stateの時間変化からactionを読む
下へ行くほどLanguage Modelとしてのdecodeから離れていく。
今回試しているのは5番目である。
LM Headすら通さず、Transformer途中のhidden stateから操作を取り出せるか。
LLMの「脳」に電極を刺してみる
今回はGemma 4 E2Bを llama.cpp で動かした。そして通常なら外から見ることのない計算グラフ途中のtensorを観測した。取得したのは、
llama_cpp_graph_tensor:l_out-17
というtensorで、1536次元のvectorが得られる。概念的には、
ゲーム状態
↓
Transformer
↓
Layer ...
↓
hidden state ─────→ ここを読む
↓
Layer ...
↓
LM Head
↓
token
という形になる。
モデル本体は変更していない。
weightを書き換えていない。
LoRAも入れていない。
新しいheadも追加していない。
既存モデルを普通に動かし、その途中を横から観測しているだけである。
自分の中では、これはLLMの「脳」に電極を刺すようなイメージに近い。新しい回路を埋め込むのではなく、すでに流れている信号を外から読む。
1536次元からコントローラのボタンを読む
当然ながら、1536個のfloatを眺めても何をしたいのかは分からない。そこで外側に小さなdecoderを置く。今回は多項ロジスティック回帰を使った。つまり、
1536次元 hidden state
↓
linear probe
↓
LEFT / RIGHT / UP / DOWN / NONE
という非常に単純な分類器である。意図的にlinear probeを使っている。複雑なニューラルネットを外側に置いてしまうと、
そのdecoder自身がゲームを理解しているのでは？
という問題が出てくる。今回は、
LLM内部に、すでに操作を読み出せる情報が存在しているか
を見たい。なのでdecoderはなるべく弱くした。
教師データはLLM自身から作る
では、hidden stateに何を正解として学習させるのか。今回は人間がSnakeの最適手を教えたわけではない。学習データを作る段階だけは、通常通りLLMをdecodeする。例えばLLMが、
{"action":"RIGHT"}
と出したとする。
そのとき、decode前に取得しておいたhidden stateに、
RIGHT
というラベルを付ける。構造としては、
┌── hidden state ──→ 学習データ
ゲーム状態 → LLM
└── decode → RIGHT ─→ 教師ラベル
となる。つまり学習しているのは、
このゲーム状態ではRIGHTが正しい
というpolicyではない。そうではなく、
このLLMがこの内部状態から最終的にどの操作をdecodeしようとしているか
を予測している。最初は通常decodeを教師として使い、その後でdecodeを取り外せないか試すわけだ。
「次tokenを読んでいるだけ」ではないのか
hidden stateからRIGHTが当たったとしても、
単に次に
"RIGHT"tokenが出ることを予測しているだけでは？
という疑問は当然ある。それならactionを読んでいるというより、next-token predictionを横から盗み見ているだけになる。そこで、同じ操作について複数の表記を用意した。
例えば右移動なら、
D
Right
右
左移動なら、
A
Left
左
とする。
ゲーム内部ではすべて共通actionへ正規化する。
D
Right
右
↓
RIGHT
こうすることで、特定の文字列だけをlinear probeが拾う可能性を多少減らせる。ただし、これだけでは十分ではない。よりちゃんと検証するなら、
WASD + Englishで学習
↓
日本語表記だけで評価
のようなVocabulary Holdoutも必要になる。もしこれでもactionを読めれば、単なるtoken予測より一段抽象化された情報を拾っている可能性が高くなる。
本気で作るなら専用Action Headの方が正しいと思う
ここまで「モデルを改造しない」ことにこだわってきたが、本当に性能を求めるなら、おそらく今回の方法が最終形ではない。
自然なのは、どこかのhidden layerから専用のaction headを分岐させることだと思う。
Transformer
│
├── LM Head → language
│
└── Action Head → controller
例えば、
hidden state
↓
small MLP
↓
action logits
とする。
本体をfreezeしてheadだけ学習してもよい。LoRAなどで内部表現そのものをaction向けに調整する方法もある。実用的なVLAやpolicy modelを作るなら、たぶんこちらが本筋になる。ただ、それをやってしまうと、
actionを出すようにモデルを改造したからactionが出た
という話になる。今回まず知りたかったのは違う。既存LLMを一切変更しなくても、内部にはすでにactionとして利用できる情報が流れているのか。だから今回は外からhidden stateを観測するだけにした。
まず電極で読んでみて、それで面白そうなら次に専用の回路を付ければよい。
まずSnakeで試す
最初に使ったのはSnakeだった。ゲームの盤面、Snakeの位置、Foodの位置、現在の進行方向などをテキストでLLMへ与える。通常通りdecodeして操作を取得すると同時に、l_out-17 のhidden stateも保存する。そこからlinear probeを学習した。未使用seedで評価すると、一致率は、
81.25%
だった。ただし、この値だけを見ると誤解しやすい。Snakeでは直進、つまり NONE がかなり多かった。多数派クラスを選び続けても高い精度になってしまうため、81.25%そのものを強い結果とは言えない。
この段階で確認できたのは、
decode前のhidden stateに、最終的な操作と相関する信号が含まれている
というところまでだ。次はactionの分布を均等にしたdatasetで評価する必要がある。
ブロック崩しでも同じ場所を読んでみる
Snakeだけでは、そのゲーム特有の特徴を拾っている可能性がある。
そこでブロック崩し(Breakout)でも試した。
操作は、
LEFT
RIGHT
NONE
の3種類。使うモデルは同じGemma 4 E2B。観測位置も同じ l_out-17。
そこから同じ1536次元vectorを取得する。
Breakout専用のlinear probeを学習して未使用episodeで評価すると、
70.0%
だった。
多数派baselineは66.7%。
大きな差ではないが、少なくともSnake以外のゲームでも同じ場所から操作に関連する信号は取れた。
SnakeとBreakoutで同じdecoderを使えるか
次に試したのは、ゲームをまたいでactionが共有できるかである。
SnakeにもLEFTがある。
BreakoutにもLEFTがある。
そこで両者の操作を、
LEFT
RIGHT
UP
DOWN
NONE
という共通actionへ正規化した。probeにはgame IDを教えない。入力は純粋な1536次元hidden stateだけ。
そこから1つのlinear probeでactionを分類する。結果は、
| 評価対象 | Accuracy |
|---|---|
| Snake + Breakout | 77.27% |
| Snake | 80.0% |
| Breakout | 70.0% |
だった。少なくともSnakeとBreakoutについては、
同一モデルの同一観測地点から、一つのlinear decoderである程度共通の操作を取り出せた。
では未知のゲームにも使えるのか
ここまで来ると、
LLM内部にはLEFTやRIGHTのような汎用的なaction表現があるのではないか
と期待したくなる。そこでTetrisを追加した。
Snake + Breakoutだけで学習したprobeを、そのままTetrisへ適用する。再学習はしない。
ゲーム単位ではzero-shotになる。結果は、
7.69%
だった。多数派baselineよりも悪い。
「LEFTニューロン」を見つけたわけではない
この結果を見る限り、
この方向 = LEFT
この方向 = RIGHT
という普遍的な単純軸がhidden spaceに存在しているのを見つけた、とはいえなそうだ。
hidden stateには、
ゲーム状態
プロンプト構造
出力形式
選択可能なaction
次token情報
ゲーム固有の判断
などが混ざっている。SnakeとBreakoutでは、その中から共通して利用できる方向をlinear probeが拾えた。しかしTetrisでは対応関係が崩れた
さらに先にはVLMがある
今はゲーム状態をテキストとして与えている。しかし、この考え方はVLMでもできるはずだ。最終的にはゲーム画面そのものを入力して、
screen / video frame
↓
VLM
↓
hidden representation
↓
action
としたい。普通にVLMを使えば、
画面
↓
VLM
↓
"The ball is moving left, so move the paddle left."
↓
parse
↓
LEFT
のような構成になる。しかし必要なのがcontroller actionだけなら、この文章は本来不要である。VLMが内部で、
自機の位置
ボールの位置
敵の位置
移動方向
危険
次の行動
などを表現できているなら、その途中からactionだけ抜けばよい。究極的には、
video
↓
VLM
↓
latent dynamics
↓
controller
で動かせるかもしれない。ゲームであれば、これはかなりVLAに近い構造になる。
今回の実験で分かったこと
今回やったのは、まだかなり小さな実験である。
Gemma 4 E2Bをそのまま動かし、llama.cpp の途中から1536次元hidden stateを取得した。
モデル自体は変更せず、外側に単純なlinear probeを付けた。
SnakeやBreakoutでは、LLM自身のdecode結果と相関するactionをある程度読み取れた。
SnakeとBreakoutを混ぜた共通decoderもある程度機能した。
一方で、そのdecoderを未知のTetrisへzero-shotで持っていくと完全に失敗した。
なので現時点では、
LLM内部に普遍的なLEFT/RIGHT表現を発見した
とは到底言えない。
それでも、
LLMが言語として答えを出す前のhidden stateに、外部操作へ利用できる情報が存在している
という感触は得られた。
0 件のコメント:
コメントを投稿