[论文解读] 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) 与 Head Importance Score for Pruning (HISP)。
• 提供定量指标来评估电路提取的成功，并通过已知电路任务的实验进行验证。

提出的方法

• 定义一个三步的机械解释性工作流程：选择行为、定义计算图粒度、对激活进行补丁以提取电路。
• 引入 Automatic Circuit DisCovery (ACDC) 通过从输出开始向输入移动的激活补丁来修剪边，使用基于 KL divergence 的度量。
• 在同一任务导向框架下调整 SP 与 HISP 技术以恢复子网络。
• 使用完整模型和子图输出之间的 KL divergence 作为主要评估度量，在提示及其损坏版本的数据集上求平均。
• 提供 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 哪些指标最好地量化 recovered circuits 的质量，以及它们与电路稀疏性之间的权衡？
• RQ4 发现的电路是否对不同的数据损坏和任务定义具有鲁棒性？
• RQ5 自动化在多大程度上可以扩展到更大模型和更复杂的行为？

主要发现

• ACDC 可以通过选择稀疏边集来恢复一个 GPT-2 Small Greater-Than 类任务的电路（在测试的电路中演示为恢复了 5/5 个组件类型）。
• 在一个具体实验中，ACDC 将 GPT-2 Small 检查的 32,000 条边中来自 68 条边进行裁剪，与此前人工识别的边相匹配。
• ACDC 在跨多个 transformer 任务中在识别电路方面与基于梯度下降的方法具有竞争力的表现，尽管鲁棒性和边的完整性因任务和度量而异。
• 与 SP 和 HISP 相比，ACDC 在若干任务（IOI、Greater-Than、tracr-reverse）上通常表现良好，但对腐扰分布和超参数敏感。
• 在独立的 induction-task 评估中，ACDC 往往在 KL divergence 与电路大小之间提供更好的权衡（边更少）相比替代方法，尽管所有方法在完全自动化工作流的步骤 3 上仍存在局限。
• 零激活补丁可完全恢复 toy-model 的电路，展示了自动电路发现在某些条件下推广的潜力。

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