Transformers Use Causal World Models in Maze-Solving Tasks

要約

解釈可能性に関する最近の研究では、さまざまなドメインにわたるタスクで訓練されたトランスフォーマー モデルの内部動作が調査されており、多くの場合、これらのネットワークが驚くほど構造化された表現を自然に開発していることが発見されています。
このような表現がタスク ドメインの構造を包括的に反映している場合、それらは一般に「ワールド モデル」(WM) と呼ばれます。
この研究では、迷路課題で訓練された変圧器でそのような WM を発見します。
特に、スパース オートエンコーダ (SAE) を採用し、アテンション パターンを分析することにより、WM の構築を検証し、回路分析と SAE 特徴ベースの分析の間の一貫性を実証します。
私たちは、孤立した特徴に介入してその因果的役割を確認し、そうすることで特定のタイプの介入間の非対称性を見つけます。
驚くべきことに、入力トークン シーケンスでアクティブな特徴を指定しようとするとモデルが失敗する場合でも、モデルはトレーニング中に確認するよりも多くのアクティブな特徴に関して推論できることがわかりました。
さらに、位置エンコーディングを変えると、モデルの残差ストリーム内で WM がエンコードされる方法が変わる可能性があることが観察されています。
おもちゃの領域におけるこれらのWMの因果的役割を分析することで、トランスフォーマーによって獲得された表現における創発構造の理解に向けて進歩し、より解釈可能で制御可能なAIシステムの開発につながることを期待しています。

要約(オリジナル)

Recent studies in interpretability have explored the inner workings of transformer models trained on tasks across various domains, often discovering that these networks naturally develop surprisingly structured representations. When such representations comprehensively reflect the task domain’s structure, they are commonly referred to as “World Models” (WMs). In this work, we discover such WMs in transformers trained on maze tasks. In particular, by employing Sparse Autoencoders (SAEs) and analysing attention patterns, we examine the construction of WMs and demonstrate consistency between the circuit analysis and the SAE feature-based analysis. We intervene upon the isolated features to confirm their causal role and, in doing so, find asymmetries between certain types of interventions. Surprisingly, we find that models are able to reason with respect to a greater number of active features than they see during training, even if attempting to specify these in the input token sequence would lead the model to fail. Futhermore, we observe that varying positional encodings can alter how WMs are encoded in a model’s residual stream. By analyzing the causal role of these WMs in a toy domain we hope to make progress toward an understanding of emergent structure in the representations acquired by Transformers, leading to the development of more interpretable and controllable AI systems.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google