论文讲解——《Circuit Tracing: Revealing Computational Graphs in Language Models》¶
Author: Gemini 2.5 Pro, mtdickens
深入语言模型内部:用“电路追踪”揭示计算图¶
大家好!如果你和我一样,对大型语言模型(LLM)比如 GPT 系列、Claude 或者 Llama 的强大能力感到惊叹,那你可能也思考过一个问题:这些模型内部到底是怎么运作的?它们如何根据我们输入的提示(Prompt)一步步生成回答?
LLM 通常被视为“黑箱”,我们知道输入输出,但中间过程却像一团迷雾。理解这个过程至关重要,这不仅关乎好奇心,更关乎模型的安全性、可靠性以及未来改进的方向。今天,我想和大家聊聊一篇名为《Circuit Tracing: Revealing Computational Graphs in Language Models》的论文,它提出了一种很有意思的方法,试图揭开 LLM 在处理单个提示时“思维过程”的神秘面纱。
阅读本文,你需要:
- 了解 Transformer 模型的基本结构(Attention、MLP、Residual Stream、Embedding、Logits 等)。
- 了解模型训练的基本概念(损失函数、梯度下降等)。
- 不需要任何可解释性(Interpretability)的预备知识。
本文目标:
- 解释为什么直接看模型的神经元(Neurons)很难理解其工作原理。
- 介绍论文提出的核心方法:使用“跨层转码器”(Cross-Layer Transcoder, CLT)构建“替代模型”(Replacement Model)。
- 讲解如何为特定提示生成“归因图”(Attribution Graph),可视化计算流程。
- 简述如何验证这些图的有效性以及方法的局限性。
挑战:为什么我们不能直接看神经元?¶
你可能会想,神经网络不就是由一层层的神经元组成的吗?我们直接看每个神经元在做什么不就行了?
这个想法很自然,但在实践中遇到了巨大困难。主要原因有两个:
- 多义性(Polysemanticity):单个神经元往往同时负责多种毫不相关的概念或功能。想象一下,你工具箱里的一把扳手,它既能拧螺丝,又能当锤子用,还能用来开瓶盖。这个神经元就像这把扳手,当你看到它被激活时,你很难确定它这次到底在干什么。
- 叠加(Superposition):模型想要表示的概念(比如“猫”、“编程”、“法国首都”)可能远比它拥有的神经元数量多得多。为了节省空间,模型会把多个概念的表示“叠加”到同一组神经元上。就像你试图把 100 件物品硬塞进一个只有 50 个格子的储物柜,每个格子都得放好几样东西。这使得单个神经元的激活模式更加难以解读。
这两个现象导致我们很难把神经元当作模型计算的基本“语义单元”。我们需要找到更好的“积木块”来理解模型的计算过程。
解决方案 Part 1:寻找更好的积木块 - 特征(Features)¶
论文提出,我们不应该直接看神经元,而是要找到更“纯粹”的表示单元,研究人员称之为特征(Features)。理想情况下,一个特征应该只代表一个相对单一、可解释的概念(比如“检测到代码中的 for 循环”、“识别出这是一个问句”、“与‘狗’相关的概念”)。这种单一语义的特性被称为单义性(Monosemanticity)。
如何找到这些特征呢?研究人员借鉴并改进了一种叫做稀疏编码(Sparse Coding)的技术。核心思想是:用一个更大的、稀疏激活的特征集合来表示模型内部的激活向量。
- 更大:特征的数量远超原始神经元的数量。
- 稀疏激活:对于任何一个输入,只有极少数特征是活跃的(激活值非零)。
这就像用一套非常庞大但高度特化的工具(特征)来取代原来那堆多功能但模糊的工具(神经元)。
相关技术:稀疏自编码器(SAE)与转码器(Transcoder)
你可能听说过稀疏自编码器(Sparse Autoencoder, SAE)。它试图学习一个 bottleneck 层(特征层),使得从这个层重构出的向量与原始输入向量尽可能相似,同时要求 bottleneck 层的激活是稀疏的。
稀疏自编码器的数学原理
稀疏自编码器是一种自编码器,其核心思想是在传统的自编码器的基础上,增加一个稀疏性约束。该约束迫使隐藏层的激活值尽量接近于0,从而学习到更有意义的特征。以下是其数学原理的详细解释:
1. 自编码器基础
自编码器的目标是学习一个恒等函数,即尽可能地将输入数据复制到输出端。它由两部分组成:
-
编码器 (Encoder): 将输入数据 \(\mathbf{x} \in \mathbb{R}^{d}\) 映射到隐藏层表示 \(\mathbf{h} \in \mathbb{R}^{k}\)。
- 数学表示: \(\mathbf{h} = f(\mathbf{W}\mathbf{x} + \mathbf{b})\), 其中 \(\mathbf{W} \in \mathbb{R}^{k \times d}\) 是编码器的权重矩阵,\(\mathbf{b} \in \mathbb{R}^{k}\) 是偏置向量,\(f(\cdot)\) 是激活函数(例如 sigmoid 或 ReLU)。
-
解码器 (Decoder): 将隐藏层表示 \(\mathbf{h}\) 映射回重构的输入数据 \(\mathbf{\hat{x}} \in \mathbb{R}^{d}\)。
- 数学表示: \(\mathbf{\hat{x}} = g(\mathbf{W'}\mathbf{h} + \mathbf{b'})\), 其中 \(\mathbf{W'} \in \mathbb{R}^{d \times k}\) 是解码器的权重矩阵,\(\mathbf{b'} \in \mathbb{R}^{d}\) 是偏置向量,\(g(\cdot)\) 是激活函数。 通常, \(\mathbf{W'}\) 设置为 \(\mathbf{W}^T\)。
-
损失函数: 衡量原始输入 \(\mathbf{x}\) 和重构输入 \(\mathbf{\hat{x}}\) 之间的差异。常见的损失函数包括:
- 均方误差 (MSE): \(L(\mathbf{x}, \mathbf{\hat{x}}) = \frac{1}{2} ||\mathbf{x} - \mathbf{\hat{x}}||^2\)
- 交叉熵损失: 用于输入数据是概率分布的情况。
2. 稀疏性约束
稀疏自编码器引入了稀疏性约束,目的是使隐藏层的大部分神经元处于非激活状态。 这有助于学习到更抽象和有意义的特征表示。
-
平均激活度: 对于隐藏层第 \(j\) 个神经元,其平均激活度 \(\hat{\rho_j}\) 定义为在所有训练样本上的激活值平均:
- \(\hat{\rho_j} = \frac{1}{m} \sum_{i=1}^{m} [a_j^{(i)}]\),其中 \(m\) 是训练样本的数量,\(a_j^{(i)}\) 是第 \(i\) 个样本在隐藏层第 \(j\) 个神经元的激活值。
-
目标稀疏度: 我们设定一个目标稀疏度 \(\rho\),通常是一个接近于0的小值(例如 \(\rho = 0.05\))。
-
Kullback-Leibler (KL) 散度惩罚: 为了使 \(\hat{\rho_j}\) 接近 \(\rho\),我们使用 KL 散度作为惩罚项。 KL 散度衡量两个概率分布之间的差异。
- KL 散度: \(KL(\rho || \hat{\rho_j}) = \rho \log \frac{\rho}{\hat{\rho_j}} + (1 - \rho) \log \frac{1 - \rho}{1 - \hat{\rho_j}}\)
3. 完整损失函数
稀疏自编码器的完整损失函数包括重构损失和稀疏性惩罚项:
-
\(L_{total} = \frac{1}{m} \sum_{i=1}^{m} L(\mathbf{x}^{(i)}, \mathbf{\hat{x}}^{(i)}) + \lambda \sum_{j=1}^{k} KL(\rho || \hat{\rho_j})\)
- 其中 \(L(\mathbf{x}^{(i)}, \mathbf{\hat{x}}^{(i)})\) 是单个样本的重构损失,例如 MSE。
- \(\lambda\) 是稀疏性惩罚项的权重,用于平衡重构误差和稀疏性约束的重要性。
- \(\sum_{j=1}^{k} KL(\rho || \hat{\rho_j})\) 是所有隐藏层神经元的 KL 散度之和,衡量整个隐藏层的稀疏程度。
4. 训练过程
训练稀疏自编码器的目标是最小化总损失函数 \(L_{total}\)。 这通常使用梯度下降算法或其他优化算法来实现。训练过程涉及以下步骤:
- 前向传播: 输入数据 \(\mathbf{x}\) 通过编码器得到隐藏层表示 \(\mathbf{h}\),再通过解码器得到重构的输入 \(\mathbf{\hat{x}}\)。
- 计算平均激活度: 计算每个隐藏层神经元的平均激活度 \(\hat{\rho_j}\)。
- 计算损失: 计算重构损失和 KL 散度惩罚项,得到总损失 \(L_{total}\)。
- 反向传播: 计算损失函数对权重 \(\mathbf{W}, \mathbf{W'}\), 偏置 \(\mathbf{b}, \mathbf{b'}\) 的梯度。
- 更新参数: 使用梯度下降或其他优化算法更新权重和偏置。
总结
稀疏自编码器通过引入稀疏性约束,能够学习到输入数据中更本质的特征表示。 这种特征表示对后续的分类、聚类等任务非常有帮助。 其数学原理的核心在于最小化重构损失的同时,迫使隐藏层神经元的大部分处于非激活状态,从而提取更具有信息量的特征。 通过调节 \(\rho\) 和 \(\lambda\) 参数,可以控制稀疏性和重构精度的平衡。
而这篇论文使用的是一种叫做转码器(Transcoder)的变体。与 SAE 重构自身输入不同,Transcoder 试图用稀疏特征来预测模型某个组件(比如 MLP 层)的输出,其输入是该组件的输入。
关键创新:跨层转码器(Cross-Layer Transcoder, CLT)
这篇论文更进一步,提出了一种跨层转码器(Cross-Layer Transcoder, CLT)。它的结构很特别:
-
读取(Encoding):CLT 的第
l
层特征,会读取(Encode)原始模型第l
层残差流(Residual Stream) 的激活x^l
。 $$ \mathbf{a^{\ell}} = \text{JumpReLU}!\left(W_{enc}^{\ell} \mathbf{x^{\ell}}\right) $$x^l
:原始模型第l
层的残差流激活向量。W_enc^l
:第l
层 CLT 特征的编码器权重矩阵。JumpReLU
:一种非线性激活函数,有助于稀疏性。a^l
:第l
层 CLT 特征的激活向量(稀疏的)。
-
写入(Decoding):第
l'
层的单个特征a^{l'}_i
,会尝试去重构(Decode)所有后续层(l >= l'
)的 MLP 输出y^l
。它对每一层l
都有不同的解码权重W_{dec}^{l' -> l}
。 $$ \mathbf{\hat{y}^{\ell}} = \sum_{\ell'=1}^{\ell} W_{dec}^{\ell' \to \ell} \mathbf{a^{\ell'}} $$y^l
:原始模型第l
层 MLP 的真实输出。a^{l'}
:第l'
层 CLT 特征的激活向量。W_{dec}^{l' -> l}
:从第l'
层特征解码到第l
层 MLP 输出的权重矩阵。ŷ^l
:CLT 对第l
层 MLP 输出的重构。注意,这个重构是由所有l' <= l
的 CLT 特征共同完成的。
CLT 的训练目标:
最小化两部分损失:
- 重构损失(MSE Loss):让 CLT 的重构输出
ŷ^l
尽可能接近原始 MLP 的真实输出y^l
。 $$ L_\text{MSE} = \sum_{\ell=1}^L |\mathbf{\hat{y}^{\ell}} - \mathbf{y^{\ell}} |^2 $$ - 稀疏损失(Sparsity Loss):惩罚特征的激活,鼓励只有少数特征被激活。
$$
L_\text{sparsity} = \lambda\sum_{\ell=1}^L \sum_{i=1}^N \textrm{tanh}(c \cdot |\mathbf{W_{dec, i}^{\ell}}| \cdot a^{\ell}_i)
$$
λ
,c
是超参数。N
是每层特征数。W_{dec, i}^l
是特征i
的所有解码向量拼接起来。这个惩罚项鼓励激活值a_i^l
小,特别是当解码权重W_{dec, i}^l
比较大的时候(即这个特征比较“重要”时)。
CLT 的优势:
- 跨层视角:一个特征可以捕捉跨越多个 MLP 层的计算模式,比如一个概念在早期层被识别,并在后续多层中持续影响计算。这可能让计算图更简洁。
- 替代 MLP:CLT 被训练来模拟 MLP 的功能,这为后面构建“替代模型”打下了基础。
解决方案 Part 2:构建可解释的代理 - 替代模型¶
有了训练好的 CLT,我们就可以构建一个替代模型(Replacement Model)。很简单:把原始 Transformer 模型中的每一个 MLP 层都换成对应的 CLT 计算。
具体来说,当模型计算到第 l
层时:
- 用该层的残差流
x^l
计算出第l
层 CLT 特征的激活a^l
。 - 计算 CLT 对第
l
层 MLP 输出的重构ŷ^l
(这会用到l
层及之前所有层的 CLT 特征a^{l'}, l'<=l
)。 - 用这个重构值
ŷ^l
替换掉 原始 MLP 的输出,然后继续模型的后续计算(比如加上 Attention 层的输出,进入下一层)。
这个替代模型完全由 Attention 层和 CLT 特征(以及 Embedding/Unembedding)组成。论文发现,对于一个 18 层的模型,规模最大的 CLT 构建的替代模型,其预测下一个词的准确率能达到原始模型的 50% 左右。这说明 CLT 确实捕捉到了 MLP 的大部分功能。
更进一步:局部替代模型(Local Replacement Model)
50% 的准确率还不够。我们希望分析模型在处理某个特定提示 P 时的行为,并且希望我们的分析模型能完美复现原始模型在该提示 P 上的输出。
于是,论文提出了局部替代模型(Local Replacement Model),它针对特定提示 P 进行构建:
- 仍然用 CLT 替换 MLP。
- 冻结 Attention 模式:计算提示 P 通过原始模型时的 Attention 权重,并在局部替代模型中固定使用这些权重,不再重新计算。
- 冻结 Normalization 分母:同样,Layer Normalization 中的均值和方差(或者说其分母)也使用原始模型在提示 P 上的值,并固定下来。
- 添加误差项:计算原始模型在提示 P 上每一层 MLP 的真实输出
y^l
和 CLT 的重构ŷ^l
之间的差值error^l = y^l - ŷ^l
。在局部替代模型的每一层,将这个特定于提示 P 的误差向量error^l
加回到 CLT 的输出上。
经过这番操作,局部替代模型在处理提示 P 时,其内部所有激活值和最终的 Logits 输出将严格等于原始模型!
关键优势:线性化!
最重要的一点来了:由于我们冻结了 Attention 模式和 Normalization 分母这两个非线性操作,并且 CLT 本身的设计(用特征激活 a^l
线性组合解码权重 W_dec
)使得 MLP 的非线性被“绕过”了,局部替代模型中,从一个特征的激活到另一个下游特征的“前激活值”(Pre-activation,即输入非线性函数之前的值)的影响变成了线性的!
这意味着,对于这个特定的提示 P,我们可以清晰地计算出:特征 A 的激活变化 1 个单位,会对特征 B 的前激活值产生多少影响。这就是下一步构建“归因图”的基础。
解决方案 Part 3:可视化计算 - 归因图(Attribution Graphs)¶
现在我们有了针对特定提示 P 的、行为与原模型一致、且内部交互线性的“局部替代模型”。我们可以把它展开成一个巨大的计算图,这就是归因图(Attribution Graph)。
归因图的节点(Nodes):
- 输入节点(Input Nodes):提示 P 中的每个 Token 的 Embedding 向量。
- 中间节点(Intermediate Nodes):在处理提示 P 时,所有被激活的 CLT 特征(注意:每个特征节点都绑定了特定的 Token 位置)。
- 误差节点(Error Nodes):之前计算的、特定于提示 P 的 MLP 重构误差
error^l
(每个误差节点也绑定了层l
和 Token 位置)。这些节点只有输出,没有输入。 - 输出节点(Output Nodes):模型预测的下一个 Token 的 Logits(通常只显示概率最高的几个 Token)。
归因图的边(Edges):
边代表节点之间的直接线性影响。一条从源节点 s
(比如一个特征 a_s
)到目标节点 t
(比如另一个特征 a_t
或一个 Logit)的边的权重 A_{s->t}
被定义为:
a_s
:源节点s
的激活值(如果是输入 Token 或误差节点,可以认为是 1)。w_{s->t}
:从源节点s
到目标节点t
的虚拟权重(Virtual Weight)。这个权重代表了在局部替代模型(冻结了 Attention 等之后)中,s
的单位激活对t
的前激活值产生的线性影响。
如何计算虚拟权重 w_{s->t}
?
这个权重 w_{s->t}
实际上是源节点 s
的输出向量(比如特征 s
的解码向量 W_{dec}^{s->l}
,或 Token 的 Embedding)通过局部替代模型中的所有线性路径(包括残差连接、被冻结的 Attention Head 的 OV 矩阵)传播到目标节点 t
的输入端(比如特征 t
的编码向量 W_{enc}^t
,或 Logit 的梯度)的总效果。
论文中给出了通过反向传播(Backpropagation)配合梯度停止(Stop-gradients)来高效计算这些边权重的具体方法。关键在于,通过在所有非线性组件(CLT 特征激活函数除外)上停止梯度,反向传播计算出的正好是这种线性影响。
归因图的特性:
- 可解释性:理想情况下,CLT 特征是单义的,使得节点易于理解。
- 线性归因:目标节点(特征或 Logit)的前激活值,精确等于所有指向它的边的权重之和。这提供了一个清晰的、可分解的贡献说明。
- 因果性(部分):图中的路径(在冻结 Attention 的假设下)代表了信息流动的因果链条。
处理复杂性:剪枝(Pruning)与可视化
即使特征是稀疏的,一个长提示的归因图也可能有成千上万的节点和边,难以直接阅读。论文采用了剪枝(Pruning)策略:
- 计算每个节点对最终输出 Logits 的间接总影响(考虑所有路径)。
- 保留影响最大的节点(比如保留贡献了 80% 影响的节点)。
- 在保留的节点之间,再根据边的直接影响大小进行一轮剪枝。
这样可以得到一个更小、更易于理解的核心计算图。
同时,论文还开发了一个交互式可视化界面,让研究者可以探索这个图,点击节点查看特征详情(比如哪些文本片段能激活它),高亮路径,甚至将功能相似的特征手动组合成超节点(Supernodes),进一步简化视图。
案例研究速览:
- 缩写(NDAG):图清晰显示了模型如何从 "National", "Digital", "Analytics", "Group" 中提取首字母 "N", "D", "A", "G",并通过特定特征(如 "say D_", "say DA_", "say DAG")组合起来,最终输出 "DAG"。
- 事实回忆(Michael Jordan):图展示了两条主要通路:一条从 "plays the sport of" 触发“说出一种运动”相关的特征;另一条从 "Michael Jordan" 触发“篮球”相关的特征。两条通路汇合,最终使模型选择了 "basketball"。
- 加法(36+59=95):图揭示了模型并非简单地按位相加,而是并行地使用了多种“启发式”策略:一条通路处理个位数(6+9=15,取 5);另一条通路处理大致的量级(约 30 + 约 60 ≈ 90);最终这些通路的结果(个位是 5,量级约 90)结合起来,指向了 "95"。
验证:我们看到的图是真的吗?¶
归因图是基于(局部)替代模型的,它和原始模型虽然在特定提示上输出一致,但内部机制可能存在差异。因此,验证至关重要。
论文采用干预(Intervention)实验来验证:
- 选择一个特征(或超节点):根据归因图,这个特征应该对某个下游特征或最终 Logit 有影响。
- 在原始模型中进行干预:修改这个特征的激活值(比如强制设为 0,或者乘以 -1)。注意,这里需要使用论文中讨论的特定干预技术(如“约束性 Patching”),以避免跨层效应的重复计算。干预时,通常也保持 Attention 模式冻结。
- 观察结果:看下游特征的激活或最终 Logits 是否如归因图预测的那样发生了变化。
论文展示了多个案例,干预实验的结果通常与归因图的预测在定性上是一致的。例如,在缩写例子中,抑制(乘以-1)代表 "Digital" 的特征,确实会抑制下游的 "say D_" 和 "say DA_" 特征,并降低最终输出 "DAG" 的概率。
超越单个提示:全局权重(Global Weights)¶
归因图只描述了模型处理一个特定提示的过程。我们更希望理解通用的计算模式。
理论上,CLT 特征之间存在不依赖于具体提示的全局连接,即前面提到的虚拟权重 w_{s->t}
中不经过 Attention 的那部分(仅通过残差连接)。这可以看作是特征之间的“默认连接强度”。
挑战:干扰权重(Interference Weights)
直接分析这些全局虚拟权重也遇到了问题。由于叠加现象,即使两个特征在真实数据中从不一起激活,它们之间也可能存在很大的虚拟权重(可能是正也可能是负)。这些权重在实际计算中几乎不起作用,但它们会“干扰”我们的分析,使得很难仅凭权重大小找出真正有意义的连接。
尝试的解决方案:结合共激活统计
论文尝试通过结合特征在大量数据上的共激活统计来过滤掉这些干扰权重。例如,计算一种叫做 TWERA 的加权权重,它会降低那些很少共同激活的特征对之间的权重值。结果显示,使用 TWERA 排序后,一个特征连接到的其他特征在语义上相关性更高了。
对加法任务的全局权重分析也比较成功,可以清晰地看到代表不同数字、不同运算步骤的特征是如何普遍连接在一起的。
局限性(Limitations)¶
这项工作非常出色,但也存在一些重要的局限性,论文也坦诚地指出了这些:
- 缺失 Attention 电路:该方法冻结了 Attention 模式,无法解释模型为什么会关注(Attend to)某些 Token。对于那些关键计算依赖于动态 Attention 的任务(比如根据上下文判断代词指代),归因图可能会“错过重点”。
- 重构误差/暗物质:CLT 无法完美重构 MLP 输出,留下了无法解释的“误差”或“暗物质”。如果关键计算步骤恰好落在这个误差里,归因图就无法揭示它。
- 未激活特征与抑制:有时,某个特征没有被激活(可能被其他特征抑制了)反而是关键。目前的归因图默认只显示激活的特征,难以分析抑制性电路。
- 图的复杂性:即使经过剪枝,归因图仍然可能非常复杂,理解起来需要花费大量精力。
- 特征抽象层次问题:学习到的特征可能过于具体(特征分裂,Feature Splitting),比如模型可能有很多个略微不同的“说‘the’”的特征,而不是一个统一的特征。这使得跨样本泛化分析变得困难。超节点是手动解决此问题的一种尝试。
- 全局电路理解困难:如前所述,干扰权重和 Attention 的复杂性使得理解全局、通用的计算模式仍然非常困难。
- 机制忠实度(Mechanistic Faithfulness):替代模型(即使是局部替代模型)模拟原始模型行为的方式(机制)可能与原始模型不同。干预实验有助于验证,但不能完全保证两者机制一致。
总结与展望¶
《Circuit Tracing》这篇论文为我们打开了一扇观察 LLM 内部运作的新窗口。通过引入跨层转码器(CLT)来学习更可解释的特征,并构建(局部)替代模型,该方法能够生成针对特定提示的、基于线性归因的计算图。
这种方法的核心优势在于:
- 绕过神经元的多义性/叠加问题:使用更可能单义的、稀疏的特征作为分析单元。
- 线性化交互:在局部替代模型中,特征间的直接影响是线性的,使得归因清晰明确。
- 可视化计算流:归因图直观地展示了信息如何从输入流经中间特征到达输出。
- 可验证性:可以通过干预实验在原始模型中验证归因图揭示的机制。
尽管存在局限性(尤其是对 Attention 机制的简化处理和全局模式的理解困难),这项工作代表了机制可解释性(Mechanistic Interpretability)领域的一个重要进展。它不仅提供了一套强大的分析工具,其产生的洞见(如加法算法、事实回忆通路等)也加深了我们对 LLM 如何执行复杂任务的理解。
未来,我们期待看到这些方法被进一步改进(比如更好地整合 Attention 分析、提高特征质量、提升机制忠实度),并被应用于更广泛、更复杂的模型行为上,最终帮助我们彻底揭开这些强大 AI 的神秘面纱。
希望这篇解读能帮助你理解这篇论文的核心思想!如果你对细节感兴趣,强烈推荐阅读原论文及其附录,还有他们发布的配套论文《On the Biology of a Large Language Model》,里面有更多精彩的案例分析。