使用电路探测揭示Transformer中的中间变量

Uncovering Intermediate Variables in Transformers using Circuit Probing

使用电路探测揭示Transformer中的中间变量

https://arxiv.org/pdf/2311.04354

摘要

神经网络模型在各种复杂任务上取得了高性能，但它们实现的算法却难以解释。通常需要假设网络计算中涉及的中间变量，以便理解这些算法。例如，语言模型在生成句子时是否依赖特定的句法属性？然而，现有的分析工具使得测试此类假设变得困难。我们提出了一种新的分析技术——电路探测（circuit probing），它能够自动发现计算假设中间变量的低级电路。这使得可以通过对模型参数的针对性消融进行因果分析。

我们将这种方法应用于在简单算术任务上训练的模型，证明了它在以下三个方面的有效性：（1）解读模型学习到的算法，（2）揭示模型内的模块化结构，以及（3）跟踪训练过程中电路的发展。在这三项实验中，我们证明了电路探测结合并扩展了现有方法的能力，为各种分析提供了一种统一的方法。最后，我们在现实世界的用例中展示了电路探测：揭示了GPT2-Small和Medium中负责主谓一致和反身代词的电路。

1 引言

Transformer模型是现代机器学习的主要工具，推动了从自然语言处理（NLP）（Devlin等人，2018；Radford等人，2019；Brown等人，2020）到计算机视觉（Dosovitskiy等人，2020）再到强化学习（Chen等人，2021）等不同子领域的突破。尽管取得了成功，但人们对它们学习实现的算法知之甚少。这个核心问题激发了一系列分析和可解释性研究，试图“打开黑箱”（Rogers等人，2021；Elhage等人，2021；Belikov，2022）。尽管付出了相当大的努力，这些模型仍然几乎完全不透明。

解释一个在复杂任务上取得成功的模型所固有的一个挑战是，研究人员通常没有他们试图揭示的算法的完整图景。然而，他们可能能够提出构成此类算法的高级因果变量。例如，人们可能会直觉地认为，计算句子中主语名词的句法数量可能对语言建模有用（Chomsky，1965；Linzen等人，2016）。这个变量必须在解决语言建模任务的算法中具有因果作用，因为它由于一致性规则（即主动词的句法数量必须与主语的句法数量匹配）而限制了句子的其他部分。然而，它也为影响下一个标记预测的其他变量留下了无限的可能性。虽然我们专注于语言建模，但这种讨论更广泛地适用于神经网络应用的任何复杂领域，从视觉（Dosovitskiy等人，2020）到天文学（Ćiprijanović等人，2020）到蛋白质折叠（Jumper等人，2021）。

我们提出电路探测（circuit probing）以便于研究Transformers（Vaswani等人，2017）中的中间变量。电路探测在模型权重上引入了一个可训练的二进制掩码，该掩码被优化以揭示（如果存在的话）计算高级中间变量的电路。这种技术结合了标准探测和因果分析方法的最佳特点，使研究人员能够（1）测试模型是否表示高级中间变量，（2）测试它们是否在模型行为中具有因果作用（而不仅仅是从模型表示中解码出来），以及（3）揭示计算这些变量的特定模型权重子集。我们首先使用简单的算术任务来展示电路探测相对于现有方法的优势，表明它比现有方法更忠实于底层模型（即，它只提供支持模型实际表示的因果变量的证据），并可靠地揭示了在模型行为中具有因果作用的电路。然后，我们使用电路探测来分析GPT2-Small和Medium（Radford等人，2019）上的两个句法现象，揭示了负责主谓一致和反身代词一致性的特定电路。

2 电路探测

在尝试解释执行某项任务的模型时，通常需要假设模型正在计算某个中间变量。电路探测试图（1）测量该中间变量是否被计算，以及（2）识别负责计算它的模型组件（即电路）。

与之前的研究（Conmy等人，2023；Cao等人，2021）类似，我们尝试通过优化冻结模型权重上的二进制掩码来揭示模型组件。最近的研究表明，神经网络通常在子网络层面上表现出结构（Csordás等人，2020；Lepori等人，2023b；Hod等人，2021）。鉴于此，我们尝试在单个注意力和MLP（多层感知机）块内揭示电路。这些层对残差流产生加性更新（Elhage等人，2021）。直观上，计算中间变量的电路应该产生根据该变量划分的输出。

例如，如果一个电路正在计算主语名词的句法数量，那么该电路应该产生属于两个等价类的输出，对应于单数主语和复数主语。因此，我们优化一个二进制掩码，以便——如果变量在特定层内由模型计算——该层的输出根据该变量进行聚类。对此的自然优化目标是软最近邻损失（见附录D），这是一种对比损失，它最小化来自同一类别的输出之间的（余弦）距离，并最大化来自不同类别的输出之间的（余弦）距离（Salakhutdinov & Hinton，2007；Frosst等人，2019）。图1中展示了电路探测。

在这项工作中，我们在MLP和注意力块的神经元层面（即构成Transformer中MLP和注意力块的线性变换矩阵的列）上进行掩码。我们使用连续稀疏化（Savarese等人，2020），这是一种在训练过程中将软掩码退火成离散掩码的剪枝技术，来学习二进制掩码（见附录C）。我们还使用l0正则化来鼓励稀疏的二进制掩码。如果这个过程成功，它将在模型组件内产生一个计算假设中间变量的稀疏电路。见算法1的伪代码。

电路探测评估：我们通过两种方式评估电路探测：（1）我们在发现的电路产生的输出向量上训练一个1-最近邻分类器，然后在保留的数据上测试这个分类器。如果电路探测成功找到了计算特定中间变量的电路，那么它的输出向量将根据该变量的可能标签进行划分，我们期望这个分类器能够取得高性能。我们在这里采用了一种相当保守的策略，随机采样每个标签的1个输出向量来训练最近邻分类器。见算法2的伪代码。（2）我们消融发现的电路（即反转学习的二进制掩码）并分析模型行为的变化。

因果关系：这种消融相当于提出了一个反事实问题：“如果模型在特定块中不计算特定中间变量z，模型的输出将如何变化？”，假设电路仅在计算z时具有因果作用。我们注意到，这个假设不太可能完全正确，因为神经元通常是多义性的（Elhage等人，2022）。然而，（1）我们试图通过使用L0正则化来鼓励最大稀疏子网络来减轻这一点，以及（2）先前的研究表明，稀疏电路可能出人意料地是单义性的（Hamblin等人，2022）。

3 基线比较

电路探测实现了两个不同的目标：（1）它允许人们测量电路根据中间变量对输入进行划分的程度（即它提供了探测准确性测量），（2）它允许通过消融电路来进行因果分析。先前的方法只实现了这些目标中的一个，而不是两者。因此，我们对每种类型的分析使用不同的基线。在分析探测准确性时，我们将电路探测与“普通探测”（以下简称“探测”）和对比探测进行比较。探测涉及学习一个分类器，从模型激活中解码关于假设中间变量的信息（Tenney等人，2019；Hewitt & Manning，2019；Ettinger，2020；Li等人，2022；Nanda等人，2023）。先前的研究表明，探测常常错误地描述了模型执行的底层计算（Hewitt & Liang，2019；Zhang & Bowman，2018；Voita & Titov，2020）。我们通过实验证明，电路探测比线性或非线性探测更忠实于模型的计算（见第4.2节、4.3节）。

对比探测是电路探测的一种消融方法，我们使用类似的对比目标训练线性探测。我们在附录R中展示了结果，并发现这种方法大多失败。

在进行因果分析时，我们将与三种现有方法进行比较：因果抽象分析、健忘探测和反事实嵌入。因果抽象分析干预Transformer层产生的激活向量，将中间变量定位到特定的向量子空间（Geiger等人，2021；2023；Wu等人，2023）。然而，它需要假设一个完整的因果图——一个高级描述，说明输入如何映射到预测，包括中间变量之间的所有交互。这对于我们希望应用神经网络的大多数现实世界任务来说是不可能的（例如，我们将在第4.4节中对语言建模应用电路探测）。我们通过实验证明了无界分布式对齐搜索（无界DAS；Wu等人，2023），一种最先进的因果抽象分析技术，的效用，并表明电路探测得出了相同的结果（见第4.1节、4.2节、4.3节）。健忘探测试图抹去关于中间变量的线性可解码信息，然后观察这种抹除对下游行为的影响（Elazar等人，2021）。如果这种行为发生了变化，那么就意味着原始网络内存在中间变量。尽管已经采用了多种技术来抹去线性可解码信息（Ravfogel等人，2020；2022；Shao等人，2023），我们在我们的健忘探测实验中使用了最先进的LEACE方法（Belrose等人，2023）。我们通过实验证明，电路探测比健忘探测更忠实于模型的计算（见第4.2节和4.3节）。最后，反事实嵌入（Tucker等人，2021）被引入以使用特定探测进行因果分析。我们发现，在我们的实验中，反事实嵌入相当没有信息量（见第4.1节和4.2节）。

4 实验

我们展示了四个不同的实验，以说明电路探测可以帮助解决的问题的广泛性。实验1、2和3研究了在简单算术任务上训练的玩具模型，其中完整的因果图很容易构建。这些实验既（1）细化或再现了先前工作的结果，（2）将电路探测与现有的分析方法进行了比较。在评估电路探测解码中间变量的能力时，我们将其与线性和非线性探测进行了比较，并在进行因果分析时，将其与健忘探测、因果抽象分析和反事实嵌入进行了比较。实验4将电路探测应用于语言模型，以证明该方法可以扩展到更现实的模型和设置。

4.1 实验1：解读神经网络算法

目标：可解释性研究的一个核心目标是描述模型实现的算法（Olah，2022）。当我们能够在两个假设的替代方案之间进行裁决时，这一崇高目标就变得更加可行。我们展示了所有探测方法和大多数因果分析方法在描述在简单算术任务上训练的模型实现的算法时产生了一致的结果。

4.2 实验2：中间变量的模块化

目标：我们现在应用电路探测来分析Transformer模型的内部组织，这是最近几项研究的主题（Lepori等人，2023b；Csordás等人，2020；Hod等人，2021；Mittal等人，2022）。我们展示了电路探测可以用来描述计算是否以模块化和可重用的方式在Transformer内实现，而其他方法未能揭示这种结构。

因果分析：因果抽象分析与电路探测和线性探测一致（见附录K），因果地涉及每个任务的自由和共享变量。另一方面，从非线性探测器生成的反事实嵌入再次未能产生任何变量具有因果性的证据（见附录L）。我们得出结论，反事实嵌入在我们的训练非线性探测器中更多地表现为对抗性示例，而不是模型的有意义的反事实输入。接下来，我们进行因果分析以了解模型是否表现出模块化。我们预期，消融计算任一任务共享变量的电路应该破坏两个任务的性能。另一方面，我们预期，消融任务1的自由变量应该破坏任务1的性能，同时对任务2的性能影响较小（反之亦然）。在附录M中，我们分析了电路探测返回的两个自由变量电路的形态，发现电路仅在一个注意力块内的张量中有所不同。因此，我们只消融该张量内的电路权重。从表1中，我们看到消融计算任一任务共享变量的电路破坏了两个任务的性能。另一方面，消融任务1的自由变量破坏了任务1的性能，同时保持了任务2的一些性能。我们观察到任务2的自由变量也有相同的趋势。我们使用健忘探测进行类似的分析，发现健忘探测没有揭示这种内部结构。我们在表1中展示了在注意力块之后的残差流上进行健忘探测的结果。总的来说，我们的结果表明，电路探测在描述模型如何构建其计算方面优于竞争技术。

4.3 实验3：电路探测作为进度测量

目标：电路探测使我们能够在中间变量的层面上深入了解Transformer的训练动态。最近的研究表明，模型可能在过拟合训练数据后很长时间突然学会泛化（Power等人，2022）。尽管这种从过拟合到泛化的相当不连续的转换（通常称为“顿悟”），Nanda等人（2022）揭示了计算可泛化算法的电路是在训练过程中连续形成的。虽然他们的工作需要逆向工程整个算法以获得这种对电路形成的洞察，但我们使用电路探测（只需要我们假设一个高级中间变量）在略有不同的环境中重现了他们的发现。

4.4 实验4：语言模型中的电路探测

目标：先前的实验集中在可以指定完整因果图的玩具任务上。然而，我们开发可解释性工具的兴趣在于分析在实际中使用的模型，这些模型用于无法构建因果图的任务。在这里，我们使用电路探测来研究预训练的GPT2-Small和GPT2-Medium如何执行语言建模。特别是，我们调查了两种依赖于句法数的语言学依赖关系：主谓一致和反身代词。

主谓一致性是指英语中句子的主语必须与句子的主要动词在句法数上保持一致的现象。例如：“The keys are on the table”是符合语法的，而“The keys is on the table”则不符合语法。我们假设在预测句子的主要动词时，会计算代表主语名词句法数的中间变量。反身代词确保反身代词与其指代对象在数上保持一致。例如：“the consultants injured themselves”是符合语法的，而“the consultants injured herself”则不符合语法。我们假设在预测反身代词时，会计算代表指代对象句法数的中间变量。

任务：我们使用Marvin & Linzen（2018）的模板生成句子前缀，这些前缀的延续可能是一个主要动词（研究主谓一致性时）或一个反身代词（研究反身代词时）。见附录S中的例子前缀。我们在包含一个干扰项的句子上运行电路探测，这个干扰项是一个非主语或非指代名词（例如“the keys to the cabinet are on the table”中的“cabinet”）。对于每种现象，我们在句子前缀的最后一个标记上运行电路探测，以揭示计算主语名词或指代对象句法数的电路。然后我们消融发现的电路，并评估模型继续保留句子前缀语法能力。具体来说，我们评估模型是否更有可能预测与主语/指代对象句法数一致或不一致的标记。对于主谓一致性，我们检查“is”和“are”的logits——如果当主语是单数时（例如“The officer...”），“is”的logit高于“are”的logit，那么我们认为模型在该句子前缀上成功了。对于反身代词，我们运行两次分析，一次比较“herself”和“themselves”的logits，另一次比较“himself”和“themselves”的logits。我们在包含一个干扰名词的IID句子和包含两个干扰名词的OOD数据集上评估模型（见附录S中的例子）。我们假设相同的电路计算两种句子结构的相关语言依赖性。即使我们从这个分析中恢复了积极的结果，也有可能我们只是在破坏整个模型，而不是消融一个专门的电路。作为对照，我们从电路中的神经元的补集集合中随机采样5个随机子网络，并重新运行消融分析。随机采样的子网络总是包含与我们的电路相同数量的神经元。

探测：见附录T.1和U中对电路探测准确性的调查。通常，我们发现大多数注意力层可以计算正确的句法数，但MLP只有在GPT2-Small和Medium的中间层才开始表现良好。见附录Y中所有GPT2设置的线性探测准确性。

因果分析：对于两种句法依赖性，我们预期消融发现的电路将使模型在区分主语/指代对象的句法数方面变得更差。我们预期消融随机子网络不应该损害模型在任何数据集上的性能。我们在正文中展示了GPT2-Small的结果，并在附录U中展示了GPT2-Medium的结果。对于这两种现象，我们发现依赖性是在第6层的注意力块中计算的。消融电路探测返回的电路显著降低了IID和OOD数据集上两种现象的性能，而消融随机子网络则不影响模型性能（见图4）。其他块没有表现出这种特征模式（见附录T.2）。这与先前的工作一致，表明句法依赖性在Transformer的中间层中表示（Tenney等人，2019；Vig & Belinkov，2019）。见附录V中的电路重叠分析，附录W中的定性结果，以及附录X中的结果，证明模型中较早计算单个标记的句法数。

5 讨论

相关工作：电路探测与近期在机制解释性方面的努力相关——这是一个蓬勃发展的领域，试图逆向工程神经网络算法。通过大量的人工努力，研究人员已经揭示了玩具模型（Olsson等人，2022；Nanda等人，2022；Chughtai等人，2023）以及更现实模型（Wang等人，2022；Hanna等人，2023；Merullo等人，2023）正在实现的算法。更广泛地说，已经有很多工作分析了语言模型的句法（Linzen & Baroni，2021；Goldberg，2019；Tenney等人，2018；McCoy等人，2018）和语义能力（Pavlick，2022；2023；Yu & Ettinger，2020；Hupkes等人，2020；Dziri等人，2023）。电路探测还与试图将模型分解为功能性子网络的工作相关（Csordás等人，2020；Hamblin等人，2022；Lepori等人，2023b；Zhang等人，2021；Panigrahi等人，2023；Hod等人，2021；Cao等人，2021）。电路探测的成功进一步证明了子网络是分析模型的有用视角。

结论：我们介绍了电路探测，这是一种揭示计算高级中间变量的低级电路的新方法。通过四个实验，我们已经展示了如何深入了解模型正在实现的底层算法，这些算法在模型中的结构，以及它们在整个训练过程中如何发展。电路探测结合并扩展了现有方法的能力，并在多个设置中超越了它们。然而，目前尚不清楚在一个给定块内多个电路如何组合以创建对残差流的一个加性更新，因此无法替换单个变量以执行反事实干预。未来的工作可能会寻求理解电路如何相互组合以实现这一目的。

原文链接： https://arxiv.org/pdf/2311.04354