Query Lens：利用间接效应解释稀疏键值特征

# Query Lens: 通过间接效应解读稀疏键-值特征 来源：https://arxiv.org/html/2606.07617 ###### 摘要 尽管稀疏自编码器提供的特征比单个神经元更易解释，但可靠地表征这些特征仍然具有挑战性。我们提出 Query Lens，它扩展了 Logit Lens，以实现对稀疏特征更全面、更忠实的解释。通过联合考虑编码器侧的键特征和解码器侧的值特征，我们既识别了激活一个特征的输入，也识别了该特征所促进的输出。我们还考虑了当特征被下游模块处理时产生的间接的、模块介导的效应，超越了 Logit Lens 所捕获的直接影响。在实验中，我们发现 Query Lens 为那些在 Logit Lens 下仍然无法解释的特征产生了一致的令牌签名。最后，我们提出了子空间通道假设，该假设认为下游模块通过特定层的子空间来读取特征。 机器学习，ICML ## 1 引言 解释大型语言模型（LLMs）的内部运作机制仍是机械可解释性领域的一个核心挑战。该研究方向的一个核心目标（Bau 等人，2019 (https://arxiv.org/html/2606.07617#bib.bib31)；Mu 和 Andreas，2021 (https://arxiv.org/html/2606.07617#bib.bib32)；Dai 等人，2022 (https://arxiv.org/html/2606.07617#bib.bib33)；Park 等人，2025 (https://arxiv.org/html/2606.07617#bib.bib34)）是给 LLM 的内部表示（即特征）赋予人类可解释的描述。最近在稀疏字典学习方面的进展，特别是通过稀疏自编码器（SAEs），通过提供一个更易分析的目标——表示为字典元素稀疏组合的内部激活（Huben 等人，2024 (https://arxiv.org/html/2606.07617#bib.bib29)），加速了这一研究方向。 参见图注 (a) 残差流动态示意图。 参见图注 (b) Logit Lens 与 Query Lens 的比较。 图 1: Query Lens 概述。(a) 写入残差流的特征被下游模块作为*查询*读取，产生*间接效应*。(b) Logit Lens 将特征直接投影到词汇空间，错过了这些间接效应。Query Lens 考虑了它们，并提供了更忠实的解释。 表征 SAE 特征的一种常见方法是数据驱动的：在大型语料库上运行模型，识别强烈激活目标特征的输入，并使用这些高激活上下文来推断特征的语义。尽管先前大量的工作（Bills 等人，2023 (https://arxiv.org/html/2606.07617#bib.bib20)；Bricken 等人，2023 (https://arxiv.org/html/2606.07617#bib.bib19)；Choi 等人，2024 (https://arxiv.org/html/2606.07617#bib.bib22)；Paulo 等人，2025 (https://arxiv.org/html/2606.07617#bib.bib21)）依赖这种做法，但它存在两个显著的限制。首先，获取特征敏感的样本通常需要在大型语料库上进行详尽的模型运行；在某些情况下，由于隐私限制，甚至无法访问底层数据（Dar 等人，2023 (https://arxiv.org/html/2606.07617#bib.bib10)）。其次，仅基于输入条件的描述并非输出导向的：它们通常无法充分捕捉特征对模型生成的因果效应，限制了其在诸如引导（Gur-Arieh 等人，2025 (https://arxiv.org/html/2606.07617#bib.bib17)）等下游应用中的可靠性。 解释 SAE 特征的一个突出替代方案是直接将它们投影到词汇空间，主要是通过应用 Logit Lens（Nostalgebraist, 2020 (https://arxiv.org/html/2606.07617#bib.bib37); Bloom 和 Lin, 2024 (https://arxiv.org/html/2606.07617#bib.bib39)）。虽然这种方法避免了收集激活上下文的复杂性，并提供了输出导向的总结，但它本身也有两个基本缺点。(1) **完整性**：虽然 Logit Lens 捕获了特征方向如何影响输出 logits，但它没有解释哪些输入最初激活了该特征，即输入侧的因果性。(2) **忠实性**：很大一部分 SAE 特征在 Logit Lens 下仍然难以捉摸。特别是，许多特征，尤其是在早期层中的特征，表现出分散的令牌模式或被无法解释的令牌主导，而不是收敛到一个连贯的语义概念。 在这项工作中，我们旨在解决上述两个问题。为了提高完整性，我们将探针扩展到 SAE 解码器之外，包括之前相对较少受到关注的 SAE 编码器。采用键-值记忆视角（Geva 等人，2021 (https://arxiv.org/html/2606.07617#bib.bib9)），我们将编码器特征称为**键特征**，解码器特征称为**值特征**。具体来说，键特征响应于输入计算激活，而值特征则根据这些激活加权后添加到残差流中。因此，这种结构直接映射到特征因果角色的两个方面：键特征表征哪些输入激活了特征，而值特征指定了它促进哪些输出。 在忠实性方面，我们认为，根据设计，Logit Lens 仅揭示了输入-特征和特征-输出交互的一个子集。具体来说，当特征方向被添加到残差流时，它对输出分布的影响分解为两个部分：直接影响和间接影响。¹¹ 直接效应通过残差流传导到输出 logits，而**间接效应** 则在下游模块消耗残差流中的特征时产生，引起输出 logits 的额外变化。Logit Lens 关注直接效应而在很大程度上忽略了间接效应，导致出现大量无法解释的令牌分布。 为此，我们提出 Query Lens，这是一个框架，旨在以更全面和更忠实的方式扩展 Logit Lens，在嵌入空间中解释 SAE 特征。首先，为了表征特征在输入和输出两侧的因果角色，Query Lens 自适应地在键特征和值特征之间切换：前者捕获什么激活了特征，而后者揭示了它促进了什么。此外，Query Lens 同时考虑了特征的直接效应和间接效应，在更广泛的情况下提供对其功能的更可靠解释。总之，这些改进使得令牌级别的解释能够更准确地反映 SAE 特征的因果足迹。通过明确考虑间接效应，Query Lens 进一步支持了对 LLM 中特征如何被下游模块转换和读取的机制分析。我们观察到，尽管特征作为静态向量被添加到残差流中，但它在不同 Transformer 组件上产生了显著不同的影响。受此发现启发，我们引入了**子空间通道假设**。我们假设特征向量的异质性处理源于选择性读出：每个模块并非统一地消耗完整特征，而是从特征的低维子空间（称为通道）中提取信息。我们通过学习从特征到模块响应的低秩线性映射来研究这种现象，揭示特征的读出是由特定层的通道介导的。 我们的代码可在 https://github.com/HYU-NLP/query-lens 获取。 ## 2 背景 ### 2.1 Transformer 的残差流视角 我们采用 Transformer 语言模型的残差流视角，其中隐藏状态形成一个单一流，仅通过残差相加进行更新（Elhage 等人，2021 (https://arxiv.org/html/2606.07617#bib.bib15)）。每一层包含两个残差块，一个多头自注意力块 \(R_{\text{A}}(\cdot)\) 和一个 MLP 块 \(R_{\text{M}}(\cdot)\)，两者都向流中添加残差更新。形式上， 
\[
h_{\text{mid}}^{l}=h_{\text{pre}}^{l}+R_{\text{A}}^{l}(h_{\text{pre}}^{l}),\quad h_{\text{post}}^{l}=h_{\text{mid}}^{l}+R_{\text{M}}^{l}(h_{\text{mid}}^{l}).
\]
然后残差流通过设置 \(h_{\text{pre}}^{l+1}=h_{\text{post}}^{l}\) 传播到下一层。为简洁起见，我们省略了 LayerNorm。 ### 2.2 稀疏自编码器 解释神经网络的一个瓶颈是许多神经元是*多语义*的，由于*叠加*（Elhage 等人，2022 (https://arxiv.org/html/2606.07617#bib.bib25)；Bricken 等人，2023 (https://arxiv.org/html/2606.07617#bib.bib19)）而响应多个不相关的解释。为了将这些特征解缠到更局部的基中，最近的工作采用了**稀疏自编码器（SAEs）**。SAEs 在鼓励稀疏特征激活的同时重建原始残差流向量。给定目标向量 \(h_{\text{post}} \in \mathbb{R}^{d_m}\)，一个 SAE 计算 
\[
\hat{h}_{\text{post}} = W_{\text{dec}} f\!\left(W_{\text{enc}}^{\top} h_{\text{post}}\right), \qquad (1)
\]
其中 \(W_{\text{enc}}, W_{\text{dec}} \in \mathbb{R}^{d_m \times d_{\text{dict}}}\) 是编码器和解码器权重矩阵，\(f(\cdot)\) 是逐点非线性函数，\(\hat{h}_{\text{post}}\) 是 \(h_{\text{post}}\) 的 SAE 重建。通过使用具有稀疏激活的过完备字典（\(d_{\text{dict}} \gg d_m\)），SAE 学习到的特征比单个神经元更单语义，因此更易表征，这促使了对它们的分析。 ### 2.3 键-值记忆 我们采用 MLP 块的**键-值记忆**视角（Geva 等人，2021 (https://arxiv.org/html/2606.07617#bib.bib9)），并将其扩展到稀疏字典。SAE 计算（式 (1)）可以写成*子更新*的和。令 \(k_i\) 和 \(v_i\) 分别表示 \(W_{\text{enc}}\) 和 \(W_{\text{dec}}\) 的第 \(i\) 列。然后，SAE 重建可以分解为 
\[
\hat{h}_{\text{post}} = \sum_{i=1}^{d_{\text{dict}}} a_i\!\left(h_{\text{post}}\right) v_i, \quad a_i\!\left(h_{\text{post}}\right) = f\!\left(\langle h_{\text{post}}, k_i \rangle\right). \qquad (2)
\]
每个子更新通过计算输入向量 \(h_{\text{post}}\) 与编码器列向量 \(k_i\) 的内积来获得标量激活 \(a_i\)，并使用该激活对解码器的相应向量 \(v_i\) 进行加权。这产生了一个注意力式的类比：编码器列向量 \(\{k_i\}\) 充当**键特征**，从输入产生稀疏激活，而解码器列向量 \(\{v_i\}\) 充当**值特征**，使用这些激活进行组合。 ## 3 Query Lens ### 3.1 Logit Lens 的关键概念和局限性 Logit Lens（Nostalgebraist, 2020 (https://arxiv.org/html/2606.07617#bib.bib37)）最初被提出作为一种检查 Transformer 中*中间隐藏状态*的方法，通过询问如果从中间层进行解码，模型会预测什么。它将残差流向量投影到词汇空间，表示为 
\[
y^{l} = U^{\top} h_{\text{post}}^{l} \in \mathbb{R}^{|\mathcal{V}|},
\]
其中 \(U \in \mathbb{R}^{d_m \times |\mathcal{V}|}\) 是解嵌入矩阵，\(|\mathcal{V}|\) 是词汇表 \(\mathcal{V}\) 的大小。先前的工作进一步将 Logit Lens 应用于模型*参数*（Geva 等人，2021, 2022）。由于 \(h_{\text{post}}^{l}\) 可以像式 (2) 那样用子更新之和来估计，投影单个子更新可以隔离其在词汇空间中的贡献：
\[
U^{\top}\!\left(a_i^{l}(h_{\text{post}}^{l}) \, v_i^{l}\right) = a_i^{l}(h_{\text{post}}^{l}) \, U^{\top} v_i^{l}.
\]
由于标量激活 \(a_i^{l}(h_{\text{post}}^{l})\) 仅重新缩放 logits 而不改变其排名，静态值向量通过 \(U^{\top} v_i^{l}\) 决定了被促进的令牌签名。 尽管应用广泛，但使用 Logit Lens 解释 SAE 特征存在两个关键局限性。首先，许多 SAE 特征在 Logit Lens 下产生无法解释的令牌。我们将此部分归因于 Logit Lens 测量的内容：它主要反映了将特征方向添加到残差流的**直接影响**，而忽略了当下游模块消耗受扰动的流并通过其计算传播其影响时产生的**间接效应**。其次，当前实践很少强调 SAE 编码器中的键特征。尽管少数开创性研究（Geva 等人，2021；Dar 等人，2023）分析了 MLP 键向量的投影，但这一研究方向与基于激活示例的 SAE 特征解释的主流数据驱动方法联系薄弱。这些局限性促使我们解决以下问题：如何通过 (1) 表征什么激活了特征以及特征促进了什么，以及 (2) 同时考虑直接效应和间接的、模块介导的效应，来解释特征的因果角色？ ### 3.2 残差流动态 我们首先使用残差流来框架化这个问题。令 \(x\) 表示输入令牌的独热指示符，使得输入嵌入通过查找得到 \(e = Ex\)，其中 \(E \in \mathbb{R}^{d_m \times |\mathcal{V}|}\) 是嵌入矩阵。令 \(a\) 是前向传递中的特征激活，\(y\) 是输出 logits，其中 \(y = U^{\top} h\) 且 \(U \in \mathbb{R}^{d_m \times |\mathcal{V}|}\) 是解嵌入矩阵。激活依赖于输入，输出依赖于激活，即 \(x \mapsto a\) 和 \(a \mapsto y\)。我们的目标是通过将输入 \(x\) 的变化与特征激活 \(a\) 的变化联系起来，以及将 \(a\) 的变化与输出 logits \(y\) 的变化联系起来，来表征这些依赖关系是如何沿残差流表达的。 #### 3.2.1 前向动态 我们首先形式化特征激活的局部扰动如何向前传播到输出 logits。考虑第 \(l \in \{1, ..., L\}\) 层的特征 \(i\)，其后激活为 \(a_i^{\, l}\)，即应用非线性后获得的标量。由将该激活从 \(a^{*}\) 扰动 \(\Delta a\) 所引致的输出 logits \(y\) 的变化可以近似为 
\[
y(a^{*} + \Delta a) \approx y(a^{*}) + \left. \frac{\partial y}{\partial a_i^{\, l}} \right|_{a_i^{\, l}=a^{*}} \Delta a.
\]
这是一阶线性化，其中每单位标量激活变化所引致的 logit 变化由 \(\partial y / \partial a_i^{\, l}\) 固定。根据 \(y = U^{\top} h_{\text{post}}^{L}\) 和 \(h_{\text{post}}^{l} \approx \sum_i a_i^{\, l} v_i^{\, l}\)，并应用链式法则，它可以展开为 
\[
\frac{\partial y}{\partial a_i^{\, l}} = \frac{\partial y}{\partial h_{\text{post}}^{L}} \frac{\partial h_{\text{post}}^{L}}{\partial h_{\text{post}}^{l}} \frac{\partial h_{\text{post}}^{l}}{\partial a_i^{\, l}} = U^{\top} \frac{\partial h_{\text{post}}^{L}}{\partial h_{\text{post}}^{l}} v_i^{\, l}.
\]