ICML'25 | 识别语言模型关键组件!Information Bottleneck Circuit方法解析!
作者|卞天,阿里巴巴达摩院实习生
摘要
近年来,回路发现(Circuit Discovery)作为解释语言模型特定决策的研究方向备受业界关注,它旨在从语言模型中找到负责解决特定任务的计算子图。然而,现有的大多数研究忽视了这些回路的整体性,并且需要为不同任务设计特定的损害输入,不仅不准确且效率低下。
为此,我们提出了一种基于信息瓶颈(Information Bottleneck, IB)原理的回路发现方法——IBCircuit,旨在全面识别信息丰富的回路。与传统的因果干预方法相比,IBCircuit 是一个优化框架,能够整体性地发现回路,且可以应用于任何给定任务,无需繁琐的损害输入设计。在间接宾语识别(Indirect Object Identification, IOI)和“大于”任务(GreaterThan)中,IBCircuit在关键节点组件和边组件方面识别出更可靠且更小的回路,表现优于相关工作。
论文地址:https://openreview.net/pdf?id=APElRzkmGY

现存问题及挑战
当前,语言模型的回路发现面临以下几个主要挑战:
-
整体性缺失:大多数现有方法将回路组件视为独立实体,忽略了它们在模型中的整体计算流。这种处理方式破坏了模型的整体性,难以准确反映回路的真实行为。
-
任务特定设计:现有方法通常需要为不同任务设计特定的损害输入(如activation patching方法)。这种方法不仅繁琐,而且在面对多任务时效率低下,难以扩展。
-
计算复杂度:随着模型规模的增大,现有方法在回路发现中的计算成本显著增加,导致其在大型语言模型上的应用受到限制。

当前工作价值介绍
针对上述挑战,我们提出了IBCircuit方法,其主要优势包括:
-
信息瓶颈原理:IBCircuit 基于信息瓶颈原理,通过最大化回路对特定任务的相关信息,同时最小化其对模型整体信息的依赖,从而识别出最具信息量且不冗余的回路。
-
整体优化框架:不同于传统方法的逐个组件干预,IBCircuit 通过引入可控的高斯噪声和可学习的IB权重,实现对模型各组件信息流的全面优化。这样的方法不仅保留了回路的整体性,还大幅提升了效率。
-
任务无关性:IBCircuit 可以应用于任何任务,无需为每个任务设计特定的损害输入。这一特性使得IBCircuit在多任务环境下具有更强的适应性和扩展性。

方法概述
-
信息瓶颈原理
信息瓶颈(Information Bottleneck, IB)旨在从输入数据中提取出与目标任务相关的最小信息表示。具体而言,IB通过最大化输入数据与目标任务之间的互信息,同时最小化输入数据与表示之间的互信息,实现信息的有效压缩与提取。
-
IBCircuit框架
IBCircuit利用信息瓶颈原理,通过以下步骤实现回路发现:
-
噪声注入与权重学习:在模型的各个组件(如注意力头和MLP)中引入可控的高斯噪声,并通过学习IB权重来调控信息流。IB权重决定了每个组件的信息保留程度,从而识别出关键组件。
-
优化目标:IBCircuit的目标函数包含两部分。一是最大化回路输出与目标任务输出之间的互信息,确保回路对任务的高相关性;二是最小化回路与模型整体之间的互信息,确保回路的简洁性和最小冗余。
-
回路形成:通过离散化学习到的IB权重,选择那些信息保留程度较高的组件,最终形成一个信奉最小冗余且高度相关的回路。


实验与结果
-
实验设置
我们在两个任务上评估了IBCircuit的表现:
-
间接宾语识别(IOI):该任务要求模型识别句子中的间接宾语,例如在句子“John 给了 Mary 一个苹果”中,识别“Mary”作为间接宾语。
-
大于任务(GreaterThan):该任务要求模型在给定的年份数据中预测一个大于特定值的年份,例如在“战争持续了从1741年到17年”中,预测后一个年份大于41。
-
实验结果
-
基于先前工作的验证:如下图所示,IBCircuit在IOI任务中显著优于其他基线方法,尤其是在回路识别的准确性和简洁性方面。然而,在GreaterThan任务中,IBCircuit在节点数量较多时表现略逊于ACDC,可能由于任务输出的多样性增加了学习难度。

-
消融实验:如上图(a)所示,结合KL损失和互信息损失,IBCircuit在保持任务表现的同时,有效减少了回路中的冗余信息。缺少任一部分的模型(IBCircuit-woMI和IBCircuit-onlyMI)在性能上均有所下降,验证了两者的协同重要性。
-
可信度与简洁性:IBCircuit在保持高任务表现的同时,能够识别出更少的关键组件,显示出其在可信度和简洁性方面的优势。例如在IOI任务上,IBCircuit在较低组件数量时仍能保持较高的Logit差异和较低的KL Div,证明了其优秀的回路识别能力。


现有工作与未来工作的关联
IBCircuit 在回路发现领域的研究中,填补了现有方法在整体性和任务无关性上的不足。同时,IBCircuit 的成功应用也为未来的研究指明了方向:
-
更复杂任务的回路发现:未来可以尝试将IBCircuit应用于更多复杂的自然语言处理任务,进一步验证其通用性和有效性。
-
跨模型的回路对比:通过比较不同模型中的回路结构,IBCircuit 可以帮助我们更好地理解不同模型在处理相同任务时的内部机制差异。
-
实时回路调控:结合在线学习和实时调控技术,IBCircuit 有望实现对模型回路的动态调整,提升模型在多变任务环境下的适应能力。

结论
在复杂的语言模型中,理解和识别关键计算子图是实现模型透明化和可解释化的重要一步。本文提出的IBCircuit方法,基于信息瓶颈原理,提供了一种高效、整体且通用的回路发现框架。在IOI和GreaterThan任务上的实验结果表明,IBCircuit在识别高效且可靠的回路方面,显著优于现有方法。未来,我们期待IBCircuit能在更多任务和更大规模的模型中得到应用,为语言模型的解释性研究开辟新的路径。
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐



所有评论(0)