【論文紹介】TDB: トランスフォーマーデバッガを使ったGPTの内部メカニズムの解析例

Transformerベースの言語モデルの内部動作を確認するためのツールを用いて各層の役割について解析する取り組みを行っている以下論文（Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small。うまく訳せませんが、「GPT-2 smallにおける間接目的語識別の解明：（学習から）自然発生するタスク処理の内部メカニズム」という感じでしょうか…）がとても面白かったのでまとめてみました。

わたしの理解不足、誤解も多々あると思いますので、興味を持たれた方はぜひ原文にあたってみてください。

Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small

概要

• LLMの内部構造に基づいたふるまいを理解するため、GPT-2 smallモデルを対象として、間接目的語識別（IOI）という自然言語タスクをどのように実行するかについて、トランスフォーマーの各層の寄与や動作メカニズムの解析手法の提案、および解析結果を報告。

• 解析手法として、モデルのある層の動作を変化（例えば、特定のアテンションヘッドの出力を固定）して、その変更がモデルの出力にどのような影響を与えるかを評価する、「パスパッチング」「ノックアウト」というリバースエンジニアリング方法を用いた。

• 今回のGPT-2 smallモデルで行った間接目的語識別（IOI）タスクの解析では、GPT-2 smallの１層あたり12のアテンションヘッド x 12層で合計144のヘッドのうち、このタスクに寄与している26ヘッドを特定した。

• 各々のヘッドは、７つの機能を行うグループに分類でき、それらのコンポーネントが協調して１つのタスクを達成するメカニズム（Circuits：回路）に整理することができた。

回路は文中の名前（この場合は間接目的語と主語）の情報を、最終的なトークン予測に適切にルーティングするための様々なカテゴリーの注意ヘッドを使用して構成される

この論文で解析をおこなった、間接目的語識別（IOI）タスクについて

間接目的語識別（IOI）タスクは、「When Mary and John went to the store, John gave a drink to」の文を完成させるために「Mary」のような間接目的語を正しく予測するタスクです。タスクの実現には、つぎのような動作が想定されます。

1. 文中のすべての名前（この場合は「Mary」と「John」）を識別する。

2. 重複する名前（この例では「John」）を除去する。

3. 残った名前（この例では「Mary」）を出力する。

各アテンションヘッドの役割

今回のGPT-2 smallモデルで行った間接目的語識別（IOI）タスクの解析で、アテンションヘッド毎の役割が整理されました。

3.1 Name Mover Heads（名前移動ヘッド）、およびNegative Name Mover Heads（負の名前移動ヘッド）

間接目的語（IO）トークンに優位に高い注意確率を示すことで、モデルが間接目的語を正確に予測するのに寄与する。これに対して、Negative Name Mover Headsは逆の役割を果たす。このことで、正と負の両方の情報をバランス良く考慮していることが示唆されます。

3.2 S-Inhibition Heads（主語抑制ヘッド）

Name Mover Headsが間接目的語（IO）に対して適切に注意を払うために、特定のトークン（この文脈では「S」、つまり主語）を「抑制」する役割をする。

3.4 Backup Name Mover Heads（バックアップ名前移動ヘッド）

主要な名前移動ヘッド（Name Mover Heads）をすべて「ノックアウト」した後でも、モデルが間接目的語識別（IOI）タスクをある程度正確に実行し続けることが発見されました。バックアップ名前移動ヘッドは、主要な名前移動ヘッドがノックアウトされた際に、名前トークンを適切な位置に「移動」させる役割を補完します。つまり、これらのヘッドは、モデルの予測の堅牢性と適応性を高めるために働いていることが示されました。

回路内の各コンポーネントが回路全体の機能にどの程度寄与しているかを可視化。スコアが高いコンポーネントは、回路のパフォーマンスにとってより重要であることを示しています。Attention headの添え字(〇, △)は、〇層目の△個目のアテンションヘッドを示す。

# 確かにバックアップ名前移動ヘッドは、特定の層に集中するのではなく、モデルの中盤から終盤にかけて薄く分散していて、安定性と堅牢性に役立ってそうですね。

TDB: Transformer Debugger

OpenAIのGitHubリポジトリに、transformer-debuggerが公開されています。実行にはAzure OpenAI APIを使う必要があるので、個人ユーザーの私は試せていませんが、デモのビデオも用意されていますので、それを眺めるだけでも面白かったです。

感想

LLMをリバースエンジニアリング手法を使って解析する手法は、複雑な言語モデルの個別の振る舞いを理解したり解釈するうえで強力なツールになりそうです。
今後は、たとえばLLMの動作を保証・認証することが求められた場合にアセスメントする具体的な手法アプローチになるかもしれませんね。

以上、最後までお読みいただきありがとうございました。