[論文レビュー] Towards Automated Circuit Discovery for Mechanistic Interpretability

Arthur Conmy, Augustine N. Mavor-Parker+3 more|arXiv (Cornell University)|Apr 28, 2023

Explainable Artificial Intelligence (XAI)被引用数 31

ひとこと要約

この論文は機構解釈性のワークフローを体系化し、Automatic Circuit DisCovery (ACDC) を導入して、モデルの挙動を実装するサブグラフ（回路）を自動で識別し、既存の方法と比較し、ROC/AUC解析を用いて検証する。

ABSTRACT

Through considerable effort and intuition, several recent works have reverse-engineered nontrivial behaviors of transformer models. This paper systematizes the mechanistic interpretability process they followed. First, researchers choose a metric and dataset that elicit the desired model behavior. Then, they apply activation patching to find which abstract neural network units are involved in the behavior. By varying the dataset, metric, and units under investigation, researchers can understand the functionality of each component. We automate one of the process' steps: to identify the circuit that implements the specified behavior in the model's computational graph. We propose several algorithms and reproduce previous interpretability results to validate them. For example, the ACDC algorithm rediscovered 5/5 of the component types in a circuit in GPT-2 Small that computes the Greater-Than operation. ACDC selected 68 of the 32,000 edges in GPT-2 Small, all of which were manually found by previous work. Our code is available at https://github.com/ArthurConmy/Automatic-Circuit-Discovery.

研究の動機と目的

• 言語モデルにおける回路を識別するための機構的解釈性で用いられる標準的なワークフローを体系化する。
• ACDCアルゴリズムを用いてエッジレベルの回路抽出ステップを自動化する。
• 回路発見のために Subnetwork Probing (SP) および Pruning のための Head Importance Score (HISP) を適用・比較する。
• 既知の回路タスクでの実験を通じて、エクセクション成功を評価する定量的指標を提供し、検証する。

提案手法

• 三段階の機構的解釈性ワークフローを定義する：挙動を選択し、計算グラフの粒度を定義し、回路を抽出するために活性化をパッチする。
• KLダイバージェンスベースの指標を用い、出力から入力へ向かって活性化パッチを通じてエッジを剪定する Automatic Circuit DisCovery (ACDC) を導入する。
• 同じタスク指向のフレームワークの下でサブネットワークを回復するために SP および HISP の手法を適用する。
• 完全モデルとサブグラフ出力間の KLダイバージェンスを主要評価指標として使用し、プロンプトとそれらの破損版のデータセットで平均化する。
• ACDCのオープンソース実装をコミュニティ利用のため提供する。

Figure 1 : Automatically discovering circuits with ACDC. Left: a computational graph for GPT-2 Small, with a recovered circuit for the IOI task highlighted in red. Only edges between adjacent layers are shown. Right: the recovered circuit with labelled nodes. All heads recovered were identified as p

実験結果

リサーチクエスチョン

• RQ1 自動パッチングは、トランスフォーマーモデルにおいて特定の挙動を実装するサブグラフ（回路）を識別できるか？
• RQ2 既存の手法（SPおよびHISP）と比較して、自動的な回路発見は既知の回路や挙動を回復する点でどうか。
• RQ3 回復された回路の質を最もよく定量化する指標は何か、回路のスパースさとどのようにトレードオフするか。
• RQ4 発見された回路は異なるデータの破損やタスク定義に対してロバストか。
• RQ5 自動化は大規模モデルやより複雑な挙動へどの程度スケールできるか？

主な発見

• ACDCは、GPT-2 Small Greater-Than風のタスクに対して、稀少なエッジのセットを選択することで回路を回復できる（テストされた回路で5/5の成分タイプを回復したことが示される）。
• ACDCは、特定の実験でGPT-2 Smallの32,000エッジ中68のエッジを剪定し、手作業で以前に特定されたエッジと一致。
• ACDCは、多くのトランスフォーマータスクで勾配降下法ベースの手法と競合する性能を示しているが、堅牢性とエッジの完全性はタスクおよび指標によって異なる。
• SPおよびHISPと比較して、ACDCは複数のタスク（IOI、Greater-Than、tracr-reverse）で概ね良好に機能するが、破損分布とハイパーパラメータに敏感。
• 単独の誘導タスク評価では、ACDCはKLダイバージェンスと回路サイズの間でより良いトレードオフを示す傾向があり（エッジ数が少ない）、他の方法よりも優れているが、すべての方法がワークフローのステップ3を完全に自動化するには限界がある。
• ゼロ活性化パッチングは toy-model の回路を完全に回復できる可能性があることを示しており、特定の条件下で自動回路発見が一般化できる可能性を示す。

(a) Choose computational graph, task, and threshold $\tau$ .