現代の言語モデル(このチャットの背後にあるAIのようなもの)を、巨大で複雑なオーケストラだと想像してみてください。このオーケストラの中には、数千人のミュージシャン(「アテンション・ヘッド」と呼ばれます)がいて、それぞれ異なる楽器を演奏しています。各ミュージシャンは、他のミュージシャンや楽譜(入力テキスト)に耳を傾け、いつ大きく音を出し、いつ音を小さくし、どの音を叩くべきかを判断します。
長い間、科学者たちは各ミュージシャンが何をしているのかを理解しようとしてきましたが、その説明は曖昧なものでした。彼らは「このミュージシャンは名詞を好むようだ」とか、「あのミュージシャンは文の冒頭に注目している」といった言い方をすることがありました。これらは、交響曲を「幸せな感じがする」とか「悲しい感じがする」と表現するようなものです。それは役に立ちますが、音楽が具体的にどのように作られているのかまでは教えてくれません。
新しいアプローチ:楽譜を書くこと
この論文は、オーケストラを理解するための異なる方法を提案しています。単に音楽を記述するのではなく、研究者たちは各ミュージシャンが何をすべきかを正確に指示する、実際の「楽譜」(実行可能なコード)を書こうと試みました。
その手順は以下の通りです。
- ミュージシャンを観察する: まず、彼らはAIモデルが何千もの文章を読み取る様子を観察し、各「ミュージシャン」(アテンション・ヘッド)がどのように反応したかを記録しました。どの単語に対して各ヘッドが注意を払ったのかを記録したのです。
- ゴーストライターに依頼する: 彼らはこれらの記録を取り、別の非常に賢いAI(「プログラム合成」エージェント)に、それらの反応を模倣できるシンプルなPythonプログラムを書かせました。これは、ゴーストライターにミュージシャンの演奏を見て、彼らが従っていた正確なルールを書き留めてもらうようなものです。例えば、「もし前の単語がピリオドなら、次の単語を見よ。もし単語が名前なら、動詞を見よ」といった具合です。
- ゴーストライターをテストする: ゴーストライターが作成したコードを取り出し、新しい文章に対して実行して、元のAIの挙動と一致するかどうかを確認しました。彼らは、コードが本物にいかに近いかを確認するために、スコアリングシステム(類似性テストのようなもの)を使用しました。
- 大きな入れ替え: これが最もエキサイティングな部分です。彼らは、元の複雑なニューラル「ミュージシャン」をオーケストラから取り出し、先ほど書いたシンプルな、人間が読める形式のコードと入れ替えました。彼らは、オーケストラが依然として同じように演奏し続けるかどうかを確認したかったのです。
判明したこと
結果は驚くほど成功でした。
- 多くのミュージシャンは単純なルールに従っている: 彼らは、AIのアテンション・ヘッドの多くが、実際には非常に論理的で記号的なルールに従っていることを発見しました。例えば、あるヘッドは単に「文の最初の単語を見る」ようにプログラムされており、また別のヘッドは「カンマを探す」ことや「新しい文が始まる単語を見つける」ことに特化していました。
- オーケストラは演奏を続ける: 複雑なニューラル・ミュージシャンの最大25%から40%を、これらのシンプルなコード・スクリプトに置き換えても、オーケストラはほとんどリズムを外しませんでした。モデルの質問への回答能力や物語の理解力は、ほぼ全く変わらないままでした。
- 大きなモデルほど解読しやすい: 興味深いことに、モデルが大きく高度になるほど(Llama-3Bのように)、これらの単純なコード・ルールを見つけるのが容易になりました。モデルが大きくなるにつれて、より組織化され、各ミュージシャンが非常に具体的で理解しやすい役割を持つようになるようです。
まとめ
この論文は、私たちは必ずしもAIを謎めいた「ブラックボックス」として扱う必要はないということを示しています。モデルの仕組みの大部分において、複雑で説明不可能な数学を、シンプルで人間が読めるコンピュータ・コードに置き換えることができるのです。
それは、複雑な手品が実は魔法ではなく、誰にでも書き写せる特定の指示のセットであると気づくようなものです。研究者たちは、現代のAIの多くの部分において、「魔法」を「指示」へと入れ替えても、ショーを台無しにすることなく成立することを証明したのです。
問題の定義
深層ニューラルネットワークが行う計算をアルゴリズム的な観点から理解することは、機械学習における重要な課題であり続けている。解釈可能性に関する従来のアプローチは、主に、人間が定義した概念に対するプローブ(探索)を用いるトップダウン的な手法、あるいは入力と出力の要約を用いるボトムアップ的な手法のいずれかを通じて、ニューロンや特徴に意味を割り当てることに依存してきた。これらの手法は強力ではあるものの、ニューラル計算の完全かつ形式的な記述を提供するには至らないことが多い。さらに、中間コンポーネントに自然言語による説明をラベル付けすることは、解釈が曖昧になったり、定式化が困難になったりする結果を招く。著者らは、数十億パラメータを持つモデルの複雑さと自然言語による説明との間の溝を埋めるための、人間にとって読みやすく、かつ形式的に検証可能な媒体が必要であると主張している。
手法
本論文は、アテンション・ヘッドの計算を説明するために、**プログラム合成(program synthesis)**に基づいたフレームワークを提案する。核心となる目的は、入力テキストシーケンス(X)からアテンション行列(A)へと直接写像できる実行可能なPythonプログラム(π)を見つけ出し、特定のアテンション・ヘッドのロジックを近似することである。
フレーム本フレームワークは、以下の4つの明確なステップで進行する:
- アテンション・マップの抽出: 著者らは、ターゲットモデルから、訓練シーケンス(構造的な単純さからTinyStoriesデータセットを使用)にわたる正解のアテンション活性化行列を抽出する。その後、顕著なトークンペア間の相互作用に焦点を絞るため、上位2.5%のアテンション重みを抽出するようにこれらのマップをフィルタリングする。
- プログラム合成: 補助的な大規模言語モデル(LLM)が、合成エージェントとして機能し、抽出されたアテンション・パターンと入力シーケンスをプロンプトとして与えられる。このエージェントは、これらのパターンを再現する実行可能なPython関数を生成する任務を負う。合成エージェントは、言語処理や数値処理のためのNumPy、spaCy、NLTKなどのライブラリを利用できる。
- 洗練とランキング: 候補となるプログラムは、構文の妥当性が検証され、イェンセン・シャノン距離(JSD)を用いて実際のアテンション・パターンと比較・スコアリングされる。上位の候補は、合成エージェントが実際のパターンと予測されたパターンの差異を対比させて修正版のプログラムを作成する、フィードバック条件付きの洗練ループを受ける。
- 選択: 各ヘッドに対する最適なプロキシ・プログラムは、保持されたデータシーケンスにおける類似性(Intersection over Union、すなわちIoUで測定)を最大化することによって選択される。
主な貢献
- プログラムによる解釈可能性: 著者らは、ニューラルなアテンション・ヘッドを、自然言語による記述に代わる、記号的で実行可能なコードに置き換える手法を導入し、形式的に検証可能な代替手段を提示している。
- 置換による因果的検証: 相関関係のみを観察する従来の研究とは異なり、このアプローチは、学習されたニューラル・コンポーネントを合成されたプログラムで直接置換することによって、説明の妥当性を検証する。
- スケーラビリティ: 本手法は、BERT-Base、GPT-2-Small、TinyLlama-1.1B、Llama-3Bという4つの異なるトランスフォーマー・アーキテクチャに適用されており、異なるモデルサイズやアーキテクチャ(双方向および因果的)に対する本手法の適用可能性を示している。
結果
- 高い整合性: テストされたすべてのモデルにおいて、かなりの割合のアテンション・ヘッドが、高い精度を持つ実行可能なプログラムによって近似可能である。最も優れたプログラムは、GPT-2で69%、TinyLlamaで74%、Llama-3Bで79%の平均IoUスコアを達成している。
- モデルの規模とアーキテクチャ: 著者らは、自己回帰型(デコーダー)モデルのアテンション・ヘッドは、BERTのような双方向(エンコーダー)モデルよりも適合させやすい傾向があることを観察している。また、適合の質はモデルのスケールとともに向上しており、これは大規模なモデルほど、そのヘッドがより高度な機能的分化を示している可能性を示唆している。
- 因果的忠実度: アテンション・ヘッドを合成されたプログラムで置換した際、モデルは機能的な性能を維持する。具体的には、最大25%のヘッドを置換しても、パープレキシティの増加はわずか16%に抑えられる。
- ダウンストリーム性能: アテンション・ヘッドの30〜40%を、最も適合するプログラムで置換しても、ダウンストリームの質問回答ベンチマーク(HellaSwag、PIQA、SciQ、ARC-Easy、Social IQA、COPAを含む)における性能は著しく低下しない。
- 相関関係: プログラムのIoU類似度と、ヘッド置換時のパープレキシティ増加との間には、強い負の相関(スペアマンの r>0.9)が存在する。これは、IoUが機能的に忠実なプログラムを選択するための信頼できるプロキシであることを示している。
意義と主張
本論文は、最先端の言語モデルにおいてさえ、かなりの部分のアテンション・パターンが記号的な用語で理解可能であることを主張している。主な意義は、学習済みのニューラル・コンポーネントが、実行可能なコードを介した記号的なサロゲート(代理物)によって、モデルの挙動を大きく変えることなく置換可能であることを示した点にある。これは、現代のLM駆動型のプログラム合成手法が、ディープネットワークがいかに機能するかという広範な問いに迫るための、実行可能な枠組みを提供していることを示唆している。
著者らは、本研究を、モデルの挙動が重みの操作ではなく、ロジックを通じて直接推論、修正、テスト可能となる、トランスフォーマーの完全な記号的特性化への一歩として位置づけている。しかし、論文は限界についても謙虚であり、(多くのヘッドが依然として低いIoUスコアしか達成していないことから)すべてのモデルのアテンション・ヘッドが現在完全に特性化されているわけではないこと、また、現在のプログラム・ライブラリは効果的ではあるものの、まだ網羅的ではないことを認めている。本研究は、トランスフォーマー・アーキテクチャの完全な記号的蒸留を目指す将来の研究のための基礎を築くものである。
毎週最高の computer science 論文をお届け。
スタンフォード、ケンブリッジ、フランス科学アカデミーの研究者に信頼されています。
受信トレイを確認して登録を完了してください。
問題が発生しました。もう一度お試しください。
スパムなし、いつでも解除可能。
週刊ダイジェスト — 最新の研究をわかりやすく。登録