Efficient Automated Circuit Discovery in Transformers using Contextual Decomposition
作者: Aliyah R. Hsu, Georgia Zhou, Yeshwanth Cherapanamjeri, Yaxuan Huang, Anobel Y. Odisho, Peter R. Carroll, Bin Yu
分类: cs.AI, cs.CL, cs.LG
发布日期: 2024-07-01 (更新: 2025-03-02)
💡 一句话要点
提出基于上下文分解的CD-T方法，高效自动发现Transformer中的电路。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 电路发现 可解释性 上下文分解 机制解释 大型语言模型 自动化
📋 核心要点
- 现有的自动电路发现方法依赖于激活修补或其近似，存在运行缓慢、近似误差以及需要非零梯度等问题。
- 论文提出上下文分解方法CD-T，通过数学公式隔离模型特征贡献，递归计算节点贡献并剪枝，实现高效电路发现。
- 实验表明，CD-T在电路恢复和模型行为复制方面优于现有方法，同时显著降低了运行时间。
📝 摘要（中文）
本文提出了一种用于Transformer的上下文分解方法(CD-T)，旨在构建大型语言模型中可解释的电路。CD-T能够生成任意抽象级别的电路，并且首次能够高效地生成细粒度的电路，例如特定序列位置的注意力头。CD-T包含一组数学公式，用于隔离模型特征的贡献。通过递归地计算模型计算图中所有节点的贡献，然后进行剪枝，与最先进的基线相比，CD-T能够将电路发现的运行时间从数小时缩短到数秒。在三个标准电路评估数据集（间接对象识别、大于比较和文档字符串补全）上，CD-T优于ACDC和EAP，以较低的运行时间平均恢复了97% ROC AUC的手动电路。此外，通过证明我们的电路比高达原始模型大小60%的随机电路更可靠80%，我们提供了CD-T电路的可靠性并非偶然的证据。最后，我们表明，对于所有任务，CD-T电路能够使用比基线更少的节点完美地复制原始模型的行为（可靠性=1）。我们的结果强调了CD-T在高效自动机制可解释性方面的巨大前景，为深入了解大型语言模型的工作原理铺平了道路。
🔬 方法详解
问题定义：现有自动机制解释研究依赖激活修补等方法识别模型中的电路，但这些方法存在运行时间长、近似误差大，以及对特定指标（如非零梯度）的依赖等问题。因此，如何高效、准确地发现大型语言模型中的电路是一个关键挑战。
核心思路：论文的核心思路是利用上下文分解（Contextual Decomposition）来隔离模型中各个特征的贡献，从而构建可解释的电路。通过递归地计算模型计算图中所有节点的贡献，并进行剪枝，可以有效地提取出对特定任务至关重要的子图，即电路。这种方法旨在克服现有方法的局限性，提高电路发现的效率和准确性。
技术框架：CD-T的技术框架主要包含以下几个阶段：1) 前向传播：对输入数据进行前向传播，得到模型中所有节点的激活值。2) 上下文分解：利用CD-T的数学公式，递归地计算每个节点对最终输出的贡献。3) 电路构建：基于节点贡献度，构建模型电路图。4) 电路剪枝：根据贡献度阈值，移除不重要的节点和连接，得到精简的电路。
关键创新：CD-T的关键创新在于其上下文分解方法，它提供了一种精确且高效的方式来量化模型中各个特征的贡献。与现有的激活修补等方法相比，CD-T不需要近似计算，并且能够处理任意抽象级别的电路，包括细粒度的注意力头。此外，CD-T的运行时间显著低于现有方法，使其能够应用于大型语言模型。
关键设计：CD-T的关键设计包括：1) 贡献度计算公式：论文中定义了一组数学公式，用于计算每个节点对最终输出的贡献。这些公式基于链式法则，递归地计算每个节点的贡献。2) 剪枝策略：论文采用基于贡献度阈值的剪枝策略，移除贡献度低于阈值的节点和连接。阈值的选择会影响电路的复杂度和性能，需要根据具体任务进行调整。3) 递归计算：CD-T采用递归的方式计算节点贡献，保证了计算的准确性和效率。
🖼️ 关键图片
📊 实验亮点
实验结果表明，CD-T在三个标准电路评估数据集上优于ACDC和EAP，平均恢复了97% ROC AUC的手动电路，同时显著降低了运行时间。此外，CD-T电路比随机电路更可靠80%，并且能够使用比基线更少的节点完美地复制原始模型的行为（可靠性=1）。这些结果表明CD-T是一种高效且可靠的自动电路发现方法。
🎯 应用场景
该研究成果可应用于大型语言模型的可解释性分析，帮助研究人员理解模型的内部工作机制，例如模型如何进行推理、知识存储在哪里等。此外，该方法还可以用于模型压缩和优化，通过识别和移除不重要的电路，降低模型的计算复杂度和存储需求。未来，该技术有望推动可信赖人工智能的发展，提高模型的可解释性和可靠性。
📄 摘要（原文）
Automated mechanistic interpretation research has attracted great interest due to its potential to scale explanations of neural network internals to large models. Existing automated circuit discovery work relies on activation patching or its approximations to identify subgraphs in models for specific tasks (circuits). They often suffer from slow runtime, approximation errors, and specific requirements of metrics, such as non-zero gradients. In this work, we introduce contextual decomposition for transformers (CD-T) to build interpretable circuits in large language models. CD-T can produce circuits of arbitrary level of abstraction, and is the first able to produce circuits as fine-grained as attention heads at specific sequence positions efficiently. CD-T consists of a set of mathematical equations to isolate contribution of model features. Through recursively computing contribution of all nodes in a computational graph of a model using CD-T followed by pruning, we are able to reduce circuit discovery runtime from hours to seconds compared to state-of-the-art baselines. On three standard circuit evaluation datasets (indirect object identification, greater-than comparisons, and docstring completion), we demonstrate that CD-T outperforms ACDC and EAP by better recovering the manual circuits with an average of 97% ROC AUC under low runtimes. In addition, we provide evidence that faithfulness of CD-T circuits is not due to random chance by showing our circuits are 80% more faithful than random circuits of up to 60% of the original model size. Finally, we show CD-T circuits are able to perfectly replicate original models' behavior (faithfulness $ = 1$) using fewer nodes than the baselines for all tasks. Our results underscore the great promise of CD-T for efficient automated mechanistic interpretability, paving the way for new insights into the workings of large language models.