【论文笔记】Transcoders Find Interpretable LLM Feature Circuits

Abstract

机制可解释性(mechanistic interpretability)的核心目标是路径分析(circuit analysis)：在模型中找出与特定行为或能力对应的稀疏子图。
然而，MLP 子层使得在基于 Transformer 的语言模型中进行细粒度的路径分析变得困难。具体而言，可解释特征（例如通过稀疏自动编码器 SAEs 得到的特征）通常是非常多神经元的线性组合，而每个神经元又带有各自的非线性激活函数。在这种情况下进行路径分析，要么会导致路径规模过于庞大而难以处理，要么无法有效区分局部行为和整体行为。
为了解决这一问题，本文探索了transcoders，旨在使用更宽、激活稀疏的 MLP 层来忠实地近似原本激活密集的 MLP 层。本文提出利用transcoder在MLP子层中进行基于权重的路径分析。由此得到的路径结构可以清晰地分解为依赖输入的部分和与输入无关的部分。
github repo

1 Introduction

LLMs的内部机制本身不透明，机制可解释性(mechanistic interpretability)致力于理解这些机制，其核心在于将模型分解为路径(circuit)：即对模型特定行为负责的可解释的子计算过程。

细粒度路径分析(circuit analysis)中的一个核心问题是如何纳入MLP子层。
直接分析 MLP 神经元会面临“多义性”问题：即神经元倾向于对许多不相关的概念产生激活。
以前工作中：

1. SAE提取的特征往往具有可解释性，但这些向量通常是大量神经元的稠密线性组合。要在机制上理解一个SAE特征（位于一个或多个 MLP 层之前）如何影响后续的 SAE 特征，可能需要考虑数量难以承受的神经元及其非线性变换。
2. 采用因果干预与基于梯度的近似方法来分析 MLP 层，这些方法缺乏输入不变性：特征之间的连接只能在某一特定输入下描述，无法推广到整个模型层面。尝试通过在多个输入上取平均来解决这一问题，反而会丧失提供输入依赖信息的能力，无法说明某一连接在单个输入上的重要性。

本文提出transcoder解决上述问题：它是一种对模型原始 MLP 子层进行宽且稀疏激活的近似方法。
具体而言，MLP transcoder 是一种具有一个隐藏层的宽 ReLU MLP，逼近原始较窄MLP子层的输出，同时在神经元激活上施加 L1 正则化，以鼓励稀疏激活。
主要动机是通过 MLP 子层实现输入不变的特征级路径分析，从而理解并解释涉及 MLP 子层的路径的一般行为。

2 Transformers preliminaries

首先，模型将输入的 token（及其位置）映射为嵌入向量$x_{\text{mid}}^{(0,t)}\in {\mathbb{R}}^{d_{\text{model}}}$，其中$t$为token编号，$d_{\text{model}}$是模型维度。模型应用一系列layers，将上一个block的隐藏层信息映射到新的隐藏层，可以表示为：
$$x_{\text{mid}}^{(l,t)}=x_{\text{pre}}^{(l,t)}+\sum _{\text{head}h}{\text{attn}}^{(l,h)}(x_{\text{pre}}^{(l,t)};x_{\text{pre}}^{(l,1:t)})\text{\hspace{1em}(1)}$$
$$x_{\text{pre}}^{(l+1,t)}=x_{\text{mid}}^{(l,t)}+{\text{MLP}}^{(l)}(x_{\text{mid}}^{(l,t)})\text{\hspace{1em}(2)}$$
其中$l$是layer的编号，$t$是token的编号，${\text{attn}}^{(l,h)}(x_{\text{pre}}^{(l,t)}$是注意力头$h$在第$l$个layer下，给定起始token $x_{\text{pre}}^{(l,1:t)}$和目标token $x_{\text{pre}}^{(l,t)}$后得到的输出。${\text{MLP}}^{(l)}(x_{\text{mid}}^{(l,t)})$表示第$l$个layer的MLP的输出。
Eq.1 展示了注意力子层如何更新位置为$t$的 token 的隐藏状态，Eq.2展示了 MLP 子层如何更新隐藏状态。
重要的是，每个子层的输出都会加到当前的隐藏状态上。因此，隐藏状态始终可以被加性地分解为所有先前子层输出的总和。因此将每个 token 的隐藏状态称为其残差流(residual stream)，每个子层都会对该残差流进行“读取”和“写入”操作。

3 Transcoders

3.1 Architecture and training

Transcoder旨在学习一个 MLP 子层的“稀疏化”近似：它们将 MLP 子层的输出近似为特征向量的稀疏线性组合。形式上，transcoder 的架构可以表示为：
$$z_{\text{TC}}(x)=\text{ReLU}(W_{\text{enc}}x+b_{\text{enc}})\text{\hspace{1em}(3)}$$
$$\text{TC}(x)=W_{\text{dec}}{z}_{\text{TC}}(x)+b_{\text{dec}}\text{\hspace{1em}(4)}$$
其中$x$是MLP子层的输入，$W_{\text{enc}}\in {\mathbb{R}}^{d_{\text{features}}\times d_{\text{model}}}$，$W_{\text{dec}}\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{features}}}$，$b_{\text{enc}}\in {\mathbb{R}}^{d_{\text{features}}}$，$b_{\text{dec}}\in {\mathbb{R}}^{d_{\text{model}}}$，$d_{\text{features}}$是transcoder中的特征向量数量，$d_{\text{model}}$是MLP输入激活的维度。通常$d_{\text{features}}$远大于$d_{\text{model}}$。

Transcoder中的每个特征由两个向量组成：$W_{\text{enc}}$的第$i$行是特征$i$的编码特征向量，$W_{\text{dec}}$的第$i$列是特征$i$的解码特征向量。$z_{\text{TC}}(x{)}_i$被称作特征$i$的激活值。
直观地说，对于每个特征，编码器向量用于确定该特征应激活的程度；解码器向量随后按该激活程度进行缩放，所有解码器向量的加权和就是 transcoder 的输出。
本文中，$f_{\text{enc}}^{(l,i)}$和$f_{\text{dec}}^{(l,i)}$被用于表示transcoder中第$l$个layer的第$i$个编码特征向量和解码特征向量。

由于希望 transcoder 学会用特征向量的稀疏线性组合来近似 MLP 子层的计算，因此对 transcoder 的训练采用如下损失函数，其中${\lambda }_1$是一个超参数，用于在稀疏性与逼真性之间进行权衡：
$${\mathcal{L}}_{\text{TC}}(x)=\underset{\text{faithfulness loss}}{\underbrace{\Vert \text{MLP}(x)-\text{TC}(x){\Vert }_2^2}}+\underset{\text{sparsity penalty}}{\underbrace{{\lambda }_1\Vert z_{\text{TC}}(x){\Vert }_1}}\text{\hspace{1em}(5)}$$

3.2 Circuit analysis with transcoders

本文在此提出一种利用 transcoder 进行特征级路径分析的新方法，该方法能够以可扩展且可解释的方式识别不同层中的 transcoder 特征是如何连接起来以完成特定任务的。重要的是，这种方法能够揭示 MLP 子层的一般输入输出行为，这是基于 SAE 的方法所无法实现的。

路径分析的主要目标是识别模型计算图中的一个子图，该子图对模型在特定任务中的（大部分）行为负责，这需要一种方法来评估该计算子图对任务的重要性。为了确定哪些边应包含在该子图中，我们必须计算每条边的归因：即前一个节点对后一个节点自身贡献的影响程度。因此，使用 SAE 进行电路分析就意味着要通过 MLP 计算 pre-MLP SAE 特征对 post-MLP SAE 特征的归因。

3.2.1 Attribution between transcoder feature pairs

首先展示如何计算一对transcoder特征之间的归因。这种归因由两个部分的乘积构成：前一特征的激活值（依赖模型输入）以及前一特征的解码器向量与后一特征的编码器向量的点积（无关模型输入）。

令$z_{\text{TC}}^{(l,i)}(x_{\text{mid}}^{(l,t)})$代表第$l$层transcoder中第$i$个特征在第$t$个 token 上的标量激活值，它是该层第$t$个 token 的 MLP 输入$x_{\text{mid}}^{(l,t)}$的函数。对于layer $l<l^{\prime }$，token $t$ 上，第$l$个transcoder中第$i$个特征对第$l^{\prime }$个transcoder中第$l^{\prime }$个特征的贡献可以表示为：
$$\underset{\text{input-dependent}}{\underbrace{z_{\text{TC}}^{(l,i)}(x_{\text{mid}}^{(l,t)})}}\underset{\text{input-invariant}}{\underbrace{(f_{\text{dec}}^{(l,i)}\cdot f_{\text{enc}}^{(l^{\prime },i^{\prime })})}}\text{\hspace{1em}(7)}$$
Derivation：
需要搞清楚在token $t$上是什么导致layer $l^{\prime }$的transcoder的特征$i^{\prime }$被激活。这个特征的激活可以写作：
$$\text{ReLU}(f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot x_{\text{mid}}^{(l^{\prime },t)}+b_{\text{enc}}^{(l^{\prime },i^{\prime })})$$
其中$f_{\text{enc}}^{(l^{\prime },i^{\prime })}$是layer $l^{\prime }$ transcoder中$W_{\text{enc}}$的第$i$行，$b_{\text{enc}}^{(l^{\prime },i^{\prime })}$是layer $l^{\prime }$ transcoder上特征$i^{\prime }$的已学习的编码偏差(encoder bias)。无视常值偏差项$b_{\text{enc}}^{(l^{\prime },i^{\prime })}$，假设这个特征被激活（可以无视ReLU），于是特征$i^{\prime }$能否被激活完全取决于$f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot x_{\text{mid}}^{(l^{\prime },t)}$。
由于Transformer中存在残差连接，$x_{\text{mid}}^{(l^{\prime },t^{\prime })}$可以被分解为模型中所有前面组成部分输出的总和。假设在2层的模型中，有：
$$x_{\text{mid}}^{(2,t)}=\sum _h{\text{attn}}^{(2,h)}(x_{\text{pre}}^{(2,t)};x_{\text{pre}}^{(2,1:t)})+{\text{MLP}}^{(1)}(x_{\text{mid}}^{(1,t)})+\sum _h{\text{attn}}^{(1,h)}(x_{\text{pre}}^{(1,t)};x_{\text{pre}}^{(1,1:t)})$$
由于线性，这意味着${\text{MLP}}^{(1)}(x_{\text{mid}}^{(1,t)})$对$f_{\text{enc}}^{(2,i^{\prime })}\cdot x_{\text{mid}}^{(2,t)}$的贡献为：
$$f_{\text{enc}}^{(2,i^{\prime })}\cdot {\text{MLP}}^{(1)}(x_{\text{mid}}^{(1,t)})$$
只要$l<l^{\prime }$，这一结论通常适用于理解第$l$层 MLP 对第$l^{\prime }$层 transcoder 中第$i^{\prime }$个特征激活的贡献。现在，如果第$l$层的transcoder能够足够好地近似第$l$层的 MLP，我们就可以用前者替代后者：
$$f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot {\text{MLP}}^{(l)}(x_{\text{mid}}^{(l,t)})\approx f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot {\text{TC}}^{(l)}(x_{\text{mid}}^{(l,t)}).$$
然后进一步分解为独立的transcoder特征：
$${\text{TC}}^{(l)}(x_{\text{mid}}^{(l,t)})=\sum _{\text{feature}j}{z}_{\text{TC}}^{(l,j)}(x_{\text{mid}}^{(l,t)})f_{\text{dec}}^{(l,j)}$$
利用线性的优势，有
$$\begin{array}{rl}{f}_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot {\text{MLP}}^{(l)}(x_{\text{mid}}^{(l,t)})& \approx f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot \sum _{\text{feature}j}{z}_{\text{TC}}^{(l,j)}(x_{\text{mid}}^{(l,t)})f_{\text{dec}}^{(l,j)}\\ & =\sum _{\text{feature}j}{z}_{\text{TC}}^{(l,j)}(x_{\text{mid}}^{(l,t)})(f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot f_{\text{dec}}^{(l,j)})\end{array}$$
于是，transcoder $l$在token $t$上的特征$i$的贡献为
$$z_{\text{TC}}^{(l,j)}(x_{\text{mid}}^{(l,t)})(f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot f_{\text{dec}}^{(l,j)})$$

3.2.2 Attribution through attention heads

Transcoder特征可以通过注意力头(attention head)进行中介传递。因此，我们将把上述分析扩展到通过注意力头的OV路径(OV circuit)来计算transcoder特征的归因。

一些transformer的前置总结
注意力分数：${\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})$表示位置$t$与$s$之间的相关性权重，由$Q{K}^{\top }$内积经过softmax归一化得到：
$${\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})=\text{softmax}(\frac{(x_{\text{pre}}^{(l,t)}{W}_Q^{(l,h)})((x_{\text{pre}}^{(l,s)}{W}_K^{(l,h)}{)}^{\top }}{\sqrt{d_K}})$$
Value变换：$W_{OV}^{(l,h)}{x}_{\text{pre}}^{(l,s)}$

1. 先通过Value矩阵$W_V^{(l,h)}$投影：$V_s^{(h)}=x_{\text{pre}}^{(l,s)}{W}_V^{(l,h)}$
2. 再通过输出矩阵$W_O^{(l,h)}$投影：$W_O^{(l,h)}{V}_s^{(h)}=W_O^{(l,h)}(x_{\text{pre}}^{(l,s)}{W}_V^{(l,h)})$
3. 合并为组合矩阵：$W_{OV}^{(l,h)}=W_V^{(l,h)}{W}_O^{(l,h)}$

首先要清楚是什么导致$l^{\prime }$ transcoder的特征$i^{\prime }$在token $t$ 上被激活。给定layer $l$的特征头$h$ ($l<l^{\prime }$)。结合上一节的讨论可知，给到特征$i^{\prime }$的特征头贡献为
$$\begin{array}{rl}& f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot {\text{attn}}^{(l,h)}(x_{\text{pre}}^{(l,t)};x_{\text{pre}}^{(l,1:t)})\\ =& f_{\text{enc}}^{(l^{\prime },i^{\prime })}\cdot (\sum _{\text{source token}s}{\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})W_{OV}^{(l,h)}{x}_{\text{pre}}^{(l,s)})\\ =& \sum _{\text{source token}s}{\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})((f_{\text{enc}}^{(l^{\prime },i^{\prime })}{)}^{\top }{W}_{OV}^{(l,h)}{x}_{\text{pre}}^{(l,s)})\\ =& \sum _{\text{source token}s}{\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})(((W_{OV}^{(l,h)}{)}^{\top }{f}_{\text{enc}}^{(l^{\prime },i^{\prime })})\cdot x_{\text{pre}}^{(l,s)})\end{array}$$
从这里可以看出token $s$在layer $l$上经过head $h$给到的贡献为
$${\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})(((W_{OV}^{(l,h)}{)}^{\top }{f}_{\text{enc}}^{(l^{\prime },i^{\prime })})\cdot x_{\text{pre}}^{(l,s)})\text{\hspace{1em}(22)}$$
由于$x_{\text{pre}}^{(l,s)}$可以被分解成其他MLP子层（或者说transcoder的特征）的输出+注意力头的输出+原始的token embedding。这些之前层的组成部分对原始特征的贡献既经过QK路径，也经过OV路径，意味着之前层的组成部分可以带来非线性的贡献。
为了解决这个问题，本文follow了A Mathematical Framework for Transformer Circuits (Elhage et al.) 这篇工作，将QK路径分数${\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})$视为定值，尽看OV路径的贡献。尽管这样我们无法具体观察QK路径带来的贡献，但是OV路径的信息已经足够丰富。如果QK路径决定从哪个token提取信息，那么OV路径决定的就是从这个token中获取什么样的信息。这在路径分析中非常有用。

假设QK score是固定的。在Eq.22中，如果$y$是一些之前层组件的输出，其中包含残差流$x_{\text{pre}}^{(l,s)}$，那么$y$通过layer $l$的OV路径的注意力头$h$对原始transcoder特征$i^{\prime }$的贡献为$y\cdot p^{\prime }$，其中
$$p^{\prime }={\text{score}}^{(l,h)}(x_{\text{pre}}^{(l,t)},x_{\text{pre}}^{(l,s)})p\text{\hspace{1em}(23)}$$
$$p=(W_{OV}^{(l,h)}{)}^{\top }{f}_{\text{enc}}^{(l^{\prime },i^{\prime })}\text{\hspace{1em}(24)}$$
$p^{\prime }$可以被看做是一个特征向量。就像transcoder的特征，特征向量$p^{\prime }$被给定向量$y$的激活程度由$y$和$p^{\prime }$的内积给出。将$p^{\prime }$看作一个特征向量这个角度可以推广到前面的所有分析。

3.2.3 Finding computational subgraphs

前面介绍了给定输入和transcoder的特征$i^{\prime }$，前面layer的transcoder特征$i$对$i^{\prime }$激活的重要程度是怎么计算的。一旦确认哪些特征$i$与$i^{\prime }$有关，通过重复确认的过程就可以这些重要的特征所形成的路径。（贪心BFS在形成的图上面寻找$N$条贡献最大的路径）

图2：寻找路径的算法的可视化。

3.2.4 De-embeddings: a special case of input-invariant information

解嵌入(de-embeddings)：一个转译器特征的解嵌入向量是一个包含模型词表中每个token的词嵌入对该转译器特征的直接影响的向量。
第$l$层transcoder中特征$i$的解嵌入向量由 $W_E^{\top }{f}_{\text{enc}}^{(l,i)}$ 给出，其中$W_E$是模型的token嵌入矩阵。
重要的是，这个向量提供了关于每个可能的输入token对该特征激活的直接贡献程度的输入无关信息。给定一个解嵌入向量，查看模型词库中哪些token具有最高的解嵌入分数，可以体现该特征的一般行为。

4 Comparison with SAEs

Transcoder最初被构想为SAE的一种变体，因此二者之间存在许多相似之处，区别仅在于训练目标的不同。
由于这些相似性，SAE可以用与transcoder完全相同的方法进行定量评估（如稀疏性和保真度）和定性评估（如特征可解释性），这些方法都是标准的SAE评估手段。在这些评估指标上对比SAE和transcoder，发现transcoder的表现与SAE相当，甚至更优。

4.1 Blind interpretability comparison of transcoders to SAEs

为了评估transcoder的可解释性，作者手动尝试解释了Pythia-410M第15层transcoder中的 50 个随机特征，以及在 MLP 输入上训练的同一层 SAE 中的 50 个随机特征。
对于每个特征，预先计算了 OpenWebText 语料库的一个子集中最能激活该特征的样本。
然后，将 SAE 和 transcoder 中的所有特征随机打乱。对于每个特征，展示其最强激活的样本，但不告知该特征来自 SAE 还是 transcoder 。
记录每个特征是否表现出可解释的模式，并在查看完所有特征之后，才揭示每个特征的来源。表 1 中的结果表明，transcoder的特征大致与 SAE 的特征一样具有可解释性。
这进一步说明，相较于 SAE，transcoder在可解释性上没有损失。

Table 1：可解释特征、可能可解释特征和不可解释特征在transcoder与基于 MLP 输入的 SAE 中的数量统计。在可解释特征中，还将其中 6 个transcoder特征和 16 个 SAE 特征认定为“上下文无关”，即这些特征似乎只在单个 token 上激活，且没有明显的依赖上下文的模式。

4.2 Quantitative comparison of transcoders to SAEs

4.2.1 Evaluation metrics

作者从定性和定量两个方面对transcoder进行评估。定性方面评估其特征的可解释性，由人工评估者判断；定量方面则评估其激活的稀疏性以及对原始 MLP 计算的忠实度。

作为特征可解释性的定性代理指标，假设可解释的特征在激活它们的示例中应表现出可解释的模式。为此，可以在一个大型文本数据集上运行transcoder，观察哪些示例会激活某个特征，并判断这些 token 是否存在可解释的模式。
尽管这种方法并不完美，但它仍是衡量这个本质上定性概念的合理代理手段。

为了衡量transcoder的稀疏性(sparsity)，可以在一个输入数据集上运行transcoder，并计算每个 token 上平均被激活的特征数量（即激活值的平均 L0 范数）。

为了衡量transcoder的忠实度(fidelity)，可以按以下步骤操作：首先，在一个大型输入数据集上运行原始模型，并计算其在该数据集上的下一 token 预测的交叉熵损失。然后，用transcoder替换模型中对应的 MLP 子层，并计算修改后的模型在数据集上的平均损失。此时，transcoder的忠实度可以通过修改后模型的损失与原始模型损失之间的差值来量化。

4.2.2 Results

图3：Transcoder与 SAE 在 GPT2-small、Pythia-410M 和 Pythia-1.4B 上的稀疏性-准确率的trade-off情况。图中的每个点对应一个训练好的 SAE 或transcoder，并标注了训练时使用的 L1 正则化项系数${\lambda }_1$。