OpenAI 团队揭秘权重稀疏 Transformer 的可解释电路

理解大语言模型如何运作,一直是人工智能领域的核心难题。

过去几年里,神经网络的能力突飞猛进,但我们对它们内部工作机制的理解却远远落后。机械可解释性研究试图逆向工程神经网络,完全理解它们内部实现的算法。然而,理解普通Transformer模型面临一个关键障碍:神经元的激活模式难以预测,不对应人类可理解的概念。

研究者们认为,这个问题的根源可能在于"叠加"现象——密集模型实际上是在近似一个更大的、未纠缠的稀疏网络的计算。以往的方法试图通过学习一个能让激活值看起来稀疏的基来解决这个问题,然后在这个基中理解模型的计算。但这些方法获得的人类可理解电路,实际上是通过抽象掉那些只被部分理解的复杂计算得来的,因此结果电路可能反映的是所选择的抽象方式,而不完全是模型的真实机制。

OpenAI团队提出了一个全新的范式,能够产生更简单、更通用、甚至在最底层抽象级别都能完全理解的电路。他们的方法是训练权重中绝大多数为零的Transformer模型——即权重的L0范数很小。这种约束极大地简化了模型计算。由于每个神经元只能从少数几个残差通道读取或写入,模型被迫不将概念表示分散到多个残差通道,也不使用超过严格需要的神经元数量来表示单个概念。

团队通过隔离能执行每个任务的最小电路来验证模型学习了不同任务的解耦电路。在这些电路中,他们发现神经元激活通常对应简单概念,比如"单引号后面的token"或"列表嵌套深度",而权重编码的概念之间的连接往往很直观。作为相对严格的验证,研究者还证明这些解耦电路对于模型在这些任务上的行为是必要且充分的:保留电路中的少数神经元并平均消融其他所有神经元可以保持任务损失,而删除电路中的少数节点会严重损害任务损失。

虽然权重稀疏训练对可解释性有重大好处,但它有一个关键缺点:需要从头训练新模型,这些模型的训练和部署效率极低,不太可能达到前沿能力水平。研究者还展示了使用"桥接"技术的初步结果,通过在每层使用桥接将权重稀疏模型的表示与目标密集模型的表示耦合,从而让稀疏模型可以作为原始密集模型的可解释替代品。

权重稀疏性如何提升可解释性

研究团队首先测量了稀疏性是否能让模型学习更小的电路。他们在一个密集模型和一个具有相当预训练损失的稀疏模型中,为每个任务计算最小电路,并对所有任务的电路大小求平均值。结果显示,剪枝权重稀疏模型产生的电路比具有相当预训练损失的密集模型小约16倍。他们还能够以更多边的代价构建任意精确的电路,这表明简单行为的电路在权重稀疏模型中比密集模型更加解耦和可定位。

为了进一步验证电路的忠实性,研究者还表明这些电路不仅是充分的,而且是必要的。当他们消融属于电路的一小部分节点,保留网络的其余部分时,性能会严重受损。

研究还发现,这种方法随模型规模改善。当增加模型的隐藏维度(保持层数不变)时,可解释性-能力前沿得到改善。改变L0沿着前沿移动,在能力和可解释性之间进行权衡。如果保持L0固定并使用更大的模型,能力和可解释性都会提高。具有相同L0的较大模型表达能力更强,每个神经元或残差通道的非零权重更少。

可解释性的另一个方面是特征质量,这在经验上似乎与激活稀疏性密切相关。研究发现,增加权重稀疏性自然会增加残差流激活的稀疏性。

定性电路研究的深入分析

研究团队的最终目标是完全从机制上理解模型的计算。虽然他们在任务上找到了小电路,并用平均消融验证了这些电路忠实地实现了模型的底层计算,但电路大小只是可解释性的代理指标。为了验证模型实现了人类可理解的算法,他们手动解释了两个模型中三个不同任务的剪枝电路(因其明显的简单性而选择),每个电路大约花费了研究者一天的工作时间。

稀疏电路的一个独特前景是,我们可以期望在没有任何特定任务数据或假设的情况下解释它们。乐观地说,在将节点与自然概念联系起来之后,我们可以通过直接检查它们的边来提取相关电路。

闭合字符串的电路解析

研究从一个简单任务开始:字符串以单引号或双引号开头,必须以相应类型的引号闭合。他们发现这个任务的电路使用两个步骤,涉及一个MLP层中的两个神经元和一个注意力头(使用一个QK通道和一个V通道)。

在第一步,最早的MLP层(第0层MLP)将双引号和单引号的嵌入组合成一个"引号检测器"神经元(通道985,在两种引号上都为正)和一个"引号类型分类器"神经元(通道460,在双引号上为正,在单引号上为负)。在第二步,第10层注意力头使用"引号检测器"作为键(通道1),使用"引号类型分类器"作为值(通道0)。由于最后一个token具有恒定的正值查询,注意力头输出成功闭合了字符串。

研究者还检查了找到的节点在整个预训练分布上是否具有相同的解释。他们发现,即使在预训练分布上,某些(但不是全部)节点也是完全单语义的。总的来说,所描述的四个组件(两个MLP神经元、一个QK通道和一个V通道)有41条边将它们连接到网络的其余部分,其中9条被该电路使用——基于此,研究者乐观地认为,即使没有特定任务数据,也有可能理解这个电路。

计数嵌套深度的电路

接下来,研究者检查了一个更复杂的任务:探测模型是否能够适当地用单括号或双括号闭合平面列表或嵌套列表。他们手动提取了这个任务的最小电路,相当确信模型使用以下三个步骤在任务分布上计数列表的嵌套深度:

嵌入步骤中,最终token的token嵌入被忽略,左括号token的嵌入写入残差通道759、826和1711,这些通道成为左括号token的"括号检测器"。

计数步骤中,给定这些括号检测器激活,模型使用第2层中的单个值通道计数括号。模型将嵌入通道求和到一个值通道(第125头通道12),该通道充当每个左括号token处激活的"左括号检测器"。注意力头125具有几乎为零的查询和整个序列上的恒定键,这意味着softmax注意力操作相当于对上下文求平均。然后头125将得到的平均左括号检测器值写入残差通道1249,残差通道1249用其幅度编码"列表深度"。

阈值化步骤中,为了确定是否要输出双括号,需要将列表深度阈值化为二进制输出。第4层中的第二个注意力头(头80)通过使用强注意力sink来完成此操作,使用列表深度作为查询通道激活(通道4)。注意力softmax充当阈值;平面列表和列表外的token具有远小于sink的查询·键值,而嵌套列表具有远大于sink的查询·键值。因此,头80仅在嵌套列表上(查询足够大的地方)向残差流通道1079("嵌套列表闭合")输出正值,并输出双括号。

对括号计数的机制理解可以用来准确预测模型在未见过的相关输入上的表现。该电路依赖于对上下文中先前看到的token的简单平均,在列表之前存在"干扰"不匹配左括号的情况下会失败。使用对抗性代码注释,研究者成功欺骗模型在平面列表上预测双括号完成。此外,他们知道注意力电路使用残差通道1249的激活幅度来表示不同的括号嵌套深度,并且还对上下文长度取这个特征的平均值。因此可以通过使列表非常长来"稀释"上下文,从而使(未剪枝的)模型在嵌套列表上错误地预测单括号。这种效果与残差通道1249的激活幅度完全相关。

这个电路总共使用6个通道,有283条边将它们连接到网络的其余部分。另外11个第3层注意力通道和48个第7层MLP神经元被从上述描述中省略,它们贡献了额外的1217条边。如果没有特定任务数据,可能很难在追踪这个电路方面取得进展。

变量类型追踪电路

即使描述模型行为的电路没有完全解耦成少量单独可解释的激活,研究者的模型仍然倾向于学习部分可解释的计算图。通过检查电路,他们相信模型使用以下两步算法来追踪名为"current"的变量是集合还是字符串。

首先,给定输入"current = set()"或"current = """,第4层中关注最近token的头(头73)通过值通道8将"current"的嵌入复制到"set()"或""""token中。然后,当模型需要稍后在序列中回忆变量"current"的值时,第6层中的头(头26)使用"current"的嵌入作为查询和键激活。因此,该头将完成任务的信息从"set()"或""""token复制到最终token。

所描述的电路使用连续两层中两个注意力头的4个QK通道和3个V通道。这些通道有100条边将它们连接到网络的其余部分。如果没有特定任务数据,可能很难理解这个电路。

使用桥接从现有模型提取电路

到目前为止,所有结果都来自从头训练的权重稀疏模型。然而,如果也能使用这些方法来理解已经训练好的模型的行为,那将是很有价值的。密集模型在计算上比稀疏模型高效得多,如果能确信稀疏模型具有与密集模型机制上类似的电路,那将很有价值。

研究者在这个方向上进行了初步探索,通过训练一个权重稀疏模型,其计算与密集模型在同一层的计算相对应。然后,权重稀疏模型激活的"可解释"扰动可以映射到密集模型激活的相应扰动,以实现期望的行为变化。

使用剪枝,他们识别可以在特定任务上表现良好的最小稀疏电路。他们手动从该电路中选择一个节点,该节点似乎(1)基于消融对任务很重要,以及(2)编码某个感兴趣的特征。他们在稀疏模型中扰动这些节点,并使用桥接将扰动线性映射到原始密集模型。

结果表明,这个过程允许构建密集模型激活的扰动,与改变感兴趣的特征一致。在字符串闭合任务中,研究者扰动了桥接稀疏模型最终注意力层输入处的一个充当"引号类型分类器"的残差通道。他们使用双引号字符串提示模型,并引导该通道使其激活类似于单引号字符串,以构建可解释的扰动。应用桥接到这个扰动后,这导致密集模型输出单引号的概率急剧增加,这与编辑存储在引号token中的引号类型的模型表示一致。

在另一个任务中,研究者操作了最终MLP层输入处的一个通道,该通道的激活在以"if"、"while"或"except"开头的行中高度为负(这些行都应以冒号结尾)。他们使用以"return True"结尾的代码提示模型,并将该通道引导到其反事实的"while True"激活。应用桥接到这个扰动后,这增加了密集模型输出冒号而不是换行符的概率,这与部分编辑密集模型关于当前行是否应基于第一个token以冒号结尾的表示一致。

方法的局限性与未来工作

研究团队认为他们的方法还有许多改进空间。

在计算效率方面,稀疏模型需要比具有相当能力的密集模型多100-1000倍的训练和推理计算。优化和系统都有相当大的改进空间。例如,当前的优化过程导致许多死神经元,这降低了训练效率。团队对更好的重新初始化技术很感兴趣。他们还相信当前方法没有最大限度地有效探索不同的稀疏性掩码。系统改进方面也有空间,特别是通过使用稀疏内核,尽管从优化角度来看这并非易事。此外,团队对权重稀疏的专家混合模型(MoE)很感兴趣,无论是出于系统还是可解释性的原因。

在多义特征方面,研究者的电路,尤其是对于更复杂的任务,并不完全由单义节点或边组成;通常概念仍然分散在一些也执行其他任务的节点上(尽管程度远小于密集模型)。一种可能性是,模型使用少量节点组之间的少量叠加可能从根本上是有利的,而不是完全单义的节点。也可能将稀疏模型的宽度扩展到接近典型SAE(其隐藏维度在数百万而不是数千)可以解决这个问题。

在非二值特征方面,研究者的特征并非都是可二值化的,这意味着它们有时在幅度上携带信息,而不仅仅是开或关。如果一个特征不能被离散化,这意味着要完全解释它,我们不仅需要解释它的活动模式,还需要解释它的幅度。

在定义忠实性方面,平均消融不是忠实性的完美衡量标准。最终,需要某种因果清洗的变体才能对电路的忠实性完全有信心。

在定义可解释性方面,研究者使用的"可解释性"概念(具有紧凑的特定任务电路)并不能完全捕获可解释性的直观概念。定性研究指向了人类可理解性的更强概念,可以尝试在改进的可解释性指标中编码。

更广泛地说,权重稀疏性可能不是完全解开叠加所需的唯一归纳偏置,可能需要进一步的概念进展。例如,专家稀疏性不仅可以通过提高效率来帮助,还可以通过提供重要的归纳偏置。研究发现,权重稀疏性允许使用更大的激活L0而不会崩溃为多义特征。类似地,可以将专家稀疏性视为允许使用更大的权重L0而不会崩溃为叠加中的计算,因为不会在每次前向传递中使用每个参数。

在改进剪枝方面,当前的剪枝方法侧重于剪枝节点(出于计算简单性),但直接剪枝边会是理想的。此外,研究发现剪枝算法通常不能完全消除所有可剪枝的节点,需要额外的手动剪枝才能获得干净的电路。

在超越小模型和简单任务的规模化方面,研究者从解释在简单任务上狭义训练的模型开始,不确定方法会如何规模化。即使在乐观的假设下,由于解释的粒度,在解释更有能力模型中的复杂行为时,该技术也会产生极其庞大和复杂的电路。可能需要依靠自动化可解释性来理解这样的电路。悲观地说,也有可能更有能力的语言模型以从根本上违背简单描述的方式执行复杂任务,这将限制雄心勃勃的机械可解释性作为一个整体的前景。

由于基本约束,非结构化权重稀疏神经网络不太可能接近密集网络的效率。因此,使用这种方法完全解释前沿模型,或从头训练可解释的前沿模型将是不可行的。需要克服这个障碍来帮助改进对前沿模型的理解。研究者对两个主要方向感到兴奋。

首先,可以扩展方法来创建一系列可解释的模型生物体,达到GPT-3的能力水平。似乎有可能Transformer学习在不同规模的稀疏和密集模型中都出现的通用电路模式。如果是这样,那么研究这些模型生物体的电路模式将让我们了解在前沿模型中寻找什么模式,并帮助我们更好地定位研究。特别是,如果创建的模型生物体的计算通过桥接与密集模型耦合,比较它们的计算可能对研究密集模型中的叠加和干扰权重等现象很有价值。

其次,尽管理解前沿模型在其整个预训练分布上是极其昂贵的,但可以通过寻求理解更少来节省计算。特别是,可以在狭窄但重要的任务分布(例如欺骗、拒绝、目标寻求)上训练稀疏桥接模型。虽然这可能无法实现对前沿模型行为的雄心勃勃的逆向工程,但这可能是安全案例的宝贵工具。

研究者还对使用稀疏电路来支持自动化可解释性感到兴奋。稀疏电路,就像字典学习方法一样,为理解模型计算提供了新的原语——一种表达计算更简单的新语言。他们怀疑自动化可解释性在这些原语上存在瓶颈,使稀疏电路成为自动化的自然补充。

论文地址：

https://cdn.openai.com/pdf/41df8f28-d4ef-43e9-aed2-823f9393e470/circuit-sparsity-paper.pdf

本文来自至顶AI实验室，一个专注于探索生成式AI前沿技术及其应用的实验室。致力于推动生成式AI在各个领域的创新与突破，挖掘其潜在的应用场景，为企业和个人提供切实可行的解决方案。

Q&A

Q1:什么是权重稀疏Transformer模型?

A:权重稀疏Transformer是一种特殊训练的神经网络模型,其绝大多数权重参数被设置为零,即权重的L0范数很小。例如,最稀疏的模型大约每1000个权重中只有1个非零值。这种约束极大地简化了模型计算,因为每个神经元只能从少数几个残差通道读取或写入,迫使模型不将概念表示分散到多个通道,也不使用超过严格需要的神经元数量来表示单个概念。

Q2:权重稀疏模型如何提升可解释性?

A:权重稀疏模型通过限制神经元的连接数量,使得模型学习到的任务特定电路比密集模型小约16倍。在这些电路中,神经元激活通常对应简单概念,如"单引号后面的token"或"列表嵌套深度",权重编码的概念之间的连接往往很直观。研究团队通过剪枝技术隔离出执行特定任务的最小电路,并用平均消融验证这些电路的必要性和充分性,从而实现了前所未有的人类可理解程度。

Q3:权重稀疏训练方法有什么局限性?

A:权重稀疏模型的主要局限在于计算效率极低,需要比具有相当能力的密集模型多100-1000倍的训练和推理计算,且不太可能达到前沿模型的能力水平。此外,对于更复杂的任务,电路并不完全由单义节点组成,某些特征仍然分散在多个节点上。研究还发现,一些特征不能被二值化,意味着需要解释其激活幅度而不仅仅是开关状态。尽管存在这些限制,研究团队相信通过系统优化、改进剪枝算法和结合桥接技术可以逐步改善这些问题。