【论文笔记】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 mid ( 0 , t ) ∈ R d model x_\text{mid}^{(0,t)}\in\mathbb{R}^{d_\text{model}} xmid(0,t)​∈Rdmodel​，其中 t t t为token编号， d model d_\text{model} dmodel​是模型维度。模型应用一系列layers，将上一个block的隐藏层信息映射到新的隐藏层，可以表示为：

x mid ( l , t ) = x pre ( l , t ) + ∑ head  h attn ( l , h ) ( x pre ( l , t ) ; x pre ( l , 1 : t ) ) (1) 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)})\tag{1} xmid(l,t)​=xpre(l,t)​+head h∑​attn(l,h)(xpre(l,t)​;xpre(l,1:t)​)(1)

x pre ( l + 1 , t ) = x mid ( l , t ) + MLP ( l ) ( x mid ( l , t ) ) (2) x_\text{pre}^{(l+1,t)}=x_\text{mid}^{(l,t)}+\text{MLP}^{(l)}(x_\text{mid}^{(l,t)})\tag{2} xpre(l+1,t)​=xmid(l,t)​+MLP(l)(xmid(l,t)​)(2)

其中 l l l是layer的编号， t t t是token的编号， attn ( l , h ) ( x pre ( l , t ) \text{attn}^{(l,h)}(x_\text{pre}^{(l,t)} attn(l,h)(xpre(l,t)​是注意力头 h h h在第 l l l个layer下，给定起始token x pre ( l , 1 : t ) x_\text{pre}^{(l,1:t)} xpre(l,1:t)​和目标token x pre ( l , t ) x_\text{pre}^{(l,t)} xpre(l,t)​后得到的输出。 MLP ( l ) ( x mid ( l , t ) ) \text{MLP}^{(l)}(x_\text{mid}^{(l,t)}) MLP(l)(xmid(l,t)​)表示第 l l l个layer的MLP的输出。

Eq.1 展示了注意力子层如何更新位置为 t t t的 token 的隐藏状态，Eq.2展示了 MLP 子层如何更新隐藏状态。

重要的是，每个子层的输出都会加到当前的隐藏状态上。因此，隐藏状态始终可以被加性地分解为所有先前子层输出的总和。因此将每个 token 的隐藏状态称为其残差流(residual stream)，每个子层都会对该残差流进行“读取”和“写入”操作。

3 Transcoders

3.1 Architecture and training

Transcoder旨在学习一个 MLP 子层的“稀疏化”近似：它们将 MLP 子层的输出近似为特征向量的稀疏线性组合。形式上，transcoder 的架构可以表示为：

z TC ( x ) = ReLU ( W enc x + b enc ) (3) z_\text{TC}(x)=\text{ReLU}(W_\text{enc}x+b_\text{enc})\tag{3} zTC​(x)=ReLU(Wenc​x+benc​)(3)

TC ( x ) = W dec z TC ( x ) + b dec (4) \text{TC}(x)=W_\text{dec}z_\text{TC}(x)+b_\text{dec}\tag{4} TC(x)=Wdec​zTC​(x)+bdec​(4)

其中 x x x是MLP子层的输入， W enc ∈ R d features × d model W_\text{enc}\in\mathbb{R}^{d_\text{features}\times d_\text{model}} Wenc​∈Rdfeatures​×dmodel​， W dec ∈ R d model × d features W_\text{dec}\in\mathbb{R}^{d_\text{model}\times d_\text{features}} Wdec​∈Rdmodel​×dfeatures​， b enc ∈ R d features b_\text{enc}\in\mathbb{R}^{d_\text{features}} benc​∈Rdfeatures​， b dec ∈ R d model b_\text{dec}\in\mathbb{R}^{d_\text{model}} bdec​∈Rdmodel​， d features d_\text{features} dfeatures​是transcoder中的特征向量数量， d model d_\text{model} dmodel​是MLP输入激活的维度。通常 d features d_\text{features} dfeatures​远大于 d model d_\text{model} dmodel​。

Transcoder中的每个特征由两个向量组成： W enc W_\text{enc} Wenc​的第 i i i行是特征 i i i的编码特征向量， W dec W_\text{dec} Wdec​的第 i i i列是特征 i i i的解码特征向量。 z TC ( x ) i z_\text{TC}(x)_i zTC​(x)i​被称作特征 i i i的激活值。

直观地说，对于每个特征，编码器向量用于确定该特征应激活的程度；解码器向量随后按该激活程度进行缩放，所有解码器向量的加权和就是 transcoder 的输出。

本文中， f enc ( l , i ) f_\text{enc}^{(l,i)} fenc(l,i)​和 f dec ( l , i ) f_\text{dec}^{(l,i)} fdec(l,i)​被用于表示transcoder中第 l l l个layer的第 i i i个编码特征向量和解码特征向量。

由于希望 transcoder 学会用特征向量的稀疏线性组合来近似 MLP 子层的计算，因此对 transcoder 的训练采用如下损失函数，其中 λ 1 \lambda_1 λ1​是一个超参数，用于在稀疏性与逼真性之间进行权衡：

L TC ( x ) = ∥ MLP ( x ) − TC ( x ) ∥ 2 2 ⏟ faithfulness loss + λ 1 ∥ z TC ( x ) ∥ 1 ⏟ sparsity penalty (5) \mathcal{L}_\text{TC}(x)=\underbrace{\|\text{MLP}(x)-\text{TC}(x)\|_2^2}_{\text{faithfulness loss}}+\underbrace{\lambda_1\|z_\text{TC}(x)\|_1}_{\text{sparsity penalty}}\tag{5} LTC​(x)=faithfulness loss ∥MLP(x)−TC(x)∥22​​​+sparsity penalty λ1​∥zTC​(x)∥1​​​(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 TC ( l , i ) ( x mid ( l , t ) ) z_\text{TC}^{(l,i)}(x_\text{mid}^{(l,t)}) zTC(l,i)​(xmid(l,t)​)代表第 l l l层transcoder中第 i i i个特征在第 t t t个 token 上的标量激活值，它是该层第 t t t个 token 的 MLP 输入 x mid ( l , t ) x^{(l,t)}_{\text{mid}} xmid(l,t)​的函数。对于layer l