[論文レビュー] Emergent Linear Representations in World Models of Self-Supervised Sequence Models

Neel Nanda, Andrew Lee+1 more|arXiv (Cornell University)|Sep 2, 2023

Neural Networks and Reservoir ComputingComputer Science被引用数 3

ひとこと要約

この論文は、自己教師あり系列モデルOthelloGPTが、絶対色ではなく相対色—'Mine'、'Yours'、'Empty'—の線形符号化によって盤面状態を表現していることを示している。これらの線形方向をプローブすることで、単純なベクトル演算を用いてモデルの内部状態を正確に制御でき、従来の勾配ベースの手法よりも解釈可能で効果的な干渉が可能になる。

ABSTRACT

How do sequence models represent their decision-making process? Prior work suggests that Othello-playing neural network learned nonlinear models of the board state (Li et al., 2023). In this work, we provide evidence of a closely related linear representation of the board. In particular, we show that probing for "my colour" vs. "opponent's colour" may be a simple yet powerful way to interpret the model's internal state. This precise understanding of the internal representations allows us to control the model's behaviour with simple vector arithmetic. Linear representations enable significant interpretability progress, which we demonstrate with further exploration of how the world model is computed.

研究の動機と目的

• 自己教師あり系列モデル（OthelloGPTなど）が内部世界状態をどのように表現しているかを調査すること、特にその表現が線形であるかどうかを明らかにすること。
• 過去の研究が非線形な盤面状態符号化を主張しているのに対し、本研究では線形表現が得られているという矛盾を解消すること。
• 内部活性化に対するベクトル演算を用いることで、より単純で解釈可能な方法でモデル行動を制御する手法を開発すること。
• モデルが合法手を計算し、空きマスや裏返しマスを検出するメカニズムについての解明的知見を提供すること。
• 世界モデルにおける線形表現の出現と、解釈可能性および制御へのその影響を探索すること。

提案手法

• 著者たちは、盤面状態の表現を絶対色（Black, White, Empty）ではなく、プレイヤーの視点に依存する相対色（Mine, Yours, Empty）として再定式化し、ターン順に変化する。
• 各トークン位置における相対色状態を予測する線形プローブを訓練し、Mine/Yours/Empty分類でほぼ完璧な精度（例：timestep 4で98.3%）を達成した。
• 内部活性化を、Mine, Yours, Empty の各方向に投影することで、モデルの内部信念状態を直接操作するベクトル演算を適用した。
• 編集された盤面状態を模倣するように活性化を編集し、その結果、モデルが修正された信念状態に整合した合法手を生成することを確認することで、手法の妥当性を検証した。
• 他の線形特徴、例えば各ターンにおけるマスの裏返し状態についてもプローブを行い、残留ストリームのダイナミクスを分析することで、合法手予測の複数の回路を検出した。
• 線形プローブと非線形プローブの性能を比較し、絶対色ラベルでは非線形プローブが高精度を達成する一方で、相対色に対する線形プローブがほぼ優れた性能を示した。

Figure 1: The emergent world models of OthelloGPT are linearly represented. We find that the board states are encoded relative to the current player’s colour ( Mine vs. Yours ) as opposed to absolute colours ( Black vs. White ).

実験結果

リサーチクエスチョン

• RQ1 OthelloGPTにおけるオセロ盤面状態の内部表現は線形構造を示しており、その形式は何か？
• RQ2 活性化空間における線形方向に対するベクトル演算を用いることで、モデルの内部世界状態と行動を因果的に変更できるか？
• RQ3 なぜ自己教師あり世界モデルに線形表現が出現するのか？ また、以前に報告された非線形表現とはどのように関係しているのか？
• RQ4 線形プローブを用いることで、モデルが合法手を計算し、空きマスや裏返しマスを検出するメカニズムについてどのような解釈的知見が得られるか？
• RQ5 合法手予測のための複数の異なる回路がモデルに存在するか？ それらは線形特徴表現とどのように関係しているか？

主な発見

• モデルは盤面状態を絶対色（Black, White, Empty）ではなく、現在のプレイヤーの視点に基づく相対的表現—'Mine'、'Yours'、'Empty'—で符号化しており、これにより、timestep 4 で98.3%の精度を達成する高精度な線形プローブが可能になった。
• Mine/Yours/Empty の方向への線形投影により、内部世界モデルの再構築がほぼ完璧に達成され、絶対色ラベルに対する非線形プローブを上回る性能を示した。
• 内部活性化に対する単純なベクトル演算により、モデル行動の効果的かつ解釈可能な制御が可能になり、例えば編集された盤面状態に適合した合法手のプレイを誘導できる。
• 勾配計算を必要としないため、勾配ベースの編集技術とは対照的に、より単純で透明性の高い代替手法を提供する。
• モデルは特にエンドゲームにおいて、合法手予測のための複数の回路を示しており、マスの裏返しイベントの線形表現も検出可能である。
• これらの発見は、線形表現が従来の仮定よりも根本的である可能性を示唆しており、残留接続が層間で特徴を繰り返し精練可能であることを示している。

Figure 2: Intervening methodology: we intervene by adding either Empty , Mine , or Yours directions into each layer of the residual stream. Red squares in each board indicate the tiles that have been intervened, teal tiles indicate new legal moves post-intervention that the model predicts.