用于特征发现与长上下文归因的回合平均SAEs

# 用于特征发现与长上下文归因的轮次平均SAE  
来源：https://arxiv.org/html/2606.28548  

###### 摘要  
稀疏自编码器（SAEs）已成为从语言模型中提取可解释特征的有用工具。然而，标准的SAE架构作用于单个token的激活，这意味着活跃特征的数量随上下文长度线性增长，使得研究长模型转录变得困难。我们引入了轮次平均SAE，它通过学习重构一个轮次（Human或Assistant回合）内的平均模型激活，用固定数量的特征表示单个轮次。我们发现，由LLM评判时，轮次平均特征比逐token特征更完整地描述了单个轮次的高层特性。我们还展示了轮次平均SAE大大简化了SAE的常见下游用途，如归因图。总体而言，轮次平均SAE使可解释性技术在长上下文场景下变得实用。  

## 1 引言  
稀疏自编码器将模型激活分解为稀疏、可解释的特征，并在电路分析、行为引导和安全监控等领域得到广泛应用。现有的SAE方法通常作用于单个token的激活，产生的特征数据随序列长度线性增长。由于许多可解释性技术归根结底是分析工具，逐token的缩放带来了实际限制。例如，考虑理解SAE特征如何影响模型输出的目标：在给定文本片段上激发的特征似乎与输出有很强的相关性，但仅凭激活值并不能建立因果关系。归因图通过估计跨层特征之间的一阶效应来应对这一问题，但在逐token分辨率下，长上下文和多轮对话产生的图变得难以检查和理解。  

除了缩放挑战，可解释性研究者通常希望理解模型为何产生特定响应。因此，拥有能够识别表征整个响应的高层属性的工具是有价值的。这些属性的例子包括主题、风格、人设以及与安全相关的特征（如对抗性提示检测）。至少部分此类信号（包括与人设相关的信号）分布在多个token上，而不是局限于少数位置（Chen et al., 2025）。相比之下，许多逐token SAE特征充当特定单词或句法模式的检测器，仅在少量token位置上激发。由于逐token SAE独立地重构每个位置的激活，训练目标中没有任何东西强制这些高层特性与所学特征中的token级细节分离。  

我们提出一个简单的修改：训练SAE作用于token跨度上的平均隐藏状态，而不是单个token激活。这通过将分析单位从单个token缩小到跨度来解决缩放挑战。该技术实现起来很简单，因为平均保持输入的维度不变，因此现有的SAE架构和训练代码无需修改。尽管我们的实验集中在人类-助手对话回合上，但该方法可以推广到任意长度的跨度，正如我们的长上下文结果所示（第4节（https://arxiv.org/html/2606.28548#S4））。我们的贡献如下：  

1. 我们展示了在轮次平均激活上训练SAE能够产生捕捉轮次高层特性的特征，并通过定性和定量评估与逐token特征进行比较。  
2. 我们引入了一种嵌套SAE架构，使用类似Matryoshka的损失函数在单个模型中联合训练轮次平均特征和逐token特征。  
3. 我们将归因图适配到轮次级别，实现了对多轮对话的紧凑分析。通过干预实验以及完备性和替换充分性指标来验证边权重的有效性，这些指标衡量特征对模型计算的解释程度。  

## 2 背景与相关工作  

### 2.1 面向可解释性的稀疏自编码器  
Bricken等人（2023）和Cunningham等人（2023）证明了稀疏自编码器能从语言模型激活中提取可解释特征。这些特征可用于引导模型行为，包括安全相关维度，并可扩展到生产级模型（Templeton et al., 2024）。此后发展出了多种稀疏性机制，包括我们在此工作中使用的BatchTopK（Bussmann et al., 2024）。特征描述通常通过用特征的最高激活示例提示LLM来生成（Bills et al., 2023; Gao et al., 2024）。  

### 2.2 时间与多尺度表示  
**时间SAE**：Bhalla等人（2025）引入了时间SAE，对指定子集的特征应用对比损失，鼓励这些特征在同一序列的相邻token上具有相似的激活。对比损失有利于激活在token维度上平滑变化的特征，作者将其与语义内容而非句法内容相关联。我们针对的是不同的属性：我们不鼓励相邻token上的平滑激活，而是训练作用于整个跨度的平均隐藏状态，产生表征整个跨度的特征，无论其中单个token的激活如何变化。我们的方法基于以下观察：逐token SAE的概念特征不一定会强烈激发在相邻token上——这类特征通常分散在上下文中的不同位置，而不是平滑变化。Lubana等人（2025）观察到标准SAE假设token位置是独立同分布的，忽略了模型激活中的时间结构，并提议将激活分解为由过去上下文导出的可预测部分和捕捉当前token新增内容的新颖部分。  

**Matryoshka SAE**：Matryoshka表示学习（Kusupati et al., 2022）训练维度递增的嵌套子空间，每个子空间独立地充当表示。Bussmann等人（2025）将此原理应用于SAE，使用嵌套字典，较小的分区捕捉一般概念，较大的分区捕捉具体细节。时间SAE使用Matryoshka风格的分区将受对比损失约束的特征与其余逐token特征分开。类似的公式可以适用于强制单个SAE内的特征表示高层特性或逐token细节，正如我们的工作所探索的。  

### 2.3 信息在token位置间的集中  
轮次平均在计算平均隐藏状态时平等对待每个token位置；然而，一些研究表明语言模型并不总是均匀地将信息分布在token位置上。Gurnee和Tegmark（2023）发现，句号token处的隐藏状态编码了前一个句子的摘要信息，超出其他位置存在的信号。类似地，Geva等人（2023）发现，对于多token实体名称，模型将关于该实体的信息集中到名称的最后一个token。关于信息如何集中在特定token以及这与轮次平均如何相互作用，是一个相关方向，我们留给未来工作。  

### 2.4 归因与电路分析  
Ameisen等人（2025）引入了归因图，对于给定的输入，量化每一层特征如何影响后续层特征和模型输出。每个节点代表一个活跃特征或一个输出logit，有向边使用一阶近似估计一个节点激活对另一个节点的影响。边权重通过模型的线性化反向传播计算，捕捉每个源节点激活对目标节点的贡献。完备性和替换充分性指标衡量特征节点相对于未解释的错误节点对模型计算的解释程度。这类图通常使用SAE特征构建，因此其大小取决于活跃特征的数量和它们描述的激活粒度。当SAE以逐token分辨率运行时，生成的图节点数量随序列长度增长。我们将归因图适配到轮次平均特征，产生每层每轮次的节点，而不是每层每token的节点，从而得到紧凑的图，其节点以高层特性表示模型行为。作为替代方法，Arora等人（2026）证明了电路追踪也可以直接在神经元基中进行，无需学习特征。  

## 3 方法  

### 3.1 轮次平均SAE  
我们在一个token跨度上的平均隐藏状态激活上训练稀疏自编码器，而不是在单个逐token激活上训练。我们将这类SAE称为轮次平均SAE，其训练目标是标准的SAE重构损失——与逐token SAE的唯一区别是输入。给定一个token跨度 \(T\)，令 \(x_t \in \mathbb{R}^{d_{\text{model}}}\) 为位置 \(t\) 的残差流。计算轮次平均：  
\[
\bar{x} = (1/|T|) \sum_{t \in T} x_t
\]  
SAE在 \(\bar{x}\) 上使用标准重构损失进行训练：  
\[
\mathcal{L} = \|\bar{x} - \text{dec}(\text{enc}(\bar{x}))\|^2
\]  
其中 \(\text{enc}(\bar{x}) = \text{BatchTopK}(W_{\text{enc}}(\bar{x} - b_{\text{dec}}) + b_{\text{enc}})\)，\(\text{dec}(z) = W_{\text{dec}} z + b_{\text{dec}}\)。由于 \(\mathbb{R}^{d_{\text{model}}}\) 中一组向量的均值与任何单个向量的形状相同，SAE架构和训练代码无需更改即可接受轮次平均样本。在我们的实验中，每个跨度对应一个多轮聊天记录中的单个轮次。该技术适用于任何token序列的分段，包括任意长度的跨度。  

**设置**。我们在Qwen-2.5-7B-Instruct[20（https://arxiv.org/html/2606.28548#bib.bib20）]（\(d_{\text{model}} = 3584\)）上进行训练，使用LMSYS-Chat-1M[23（https://arxiv.org/html/2606.28548#bib.bib23）]作为训练语料。所有SAE均使用BatchTopK，\(d_{\text{sae}} = 32,768\)，\(k = 128\)。训练在4×H100上大约需要1.5小时完成一个epoch（约158万个轮次平均样本），而相同数据上的逐token SAE需要处理2.66亿个token。更多训练细节见附录A（https://arxiv.org/html/2606.28548#A1）。  

**直观理解**。为什么平均会迫使SAE学习高层特征？许多逐token SAE特征是低层的——例如，单token检测器以及词汇或句法模式——因为token表示的大部分必须编码它在上下文中的位置以及接下来的内容。当对多个位置的激活取平均时，这些逐token特性因跨token变异而被平滑掉。如果我们将一个跨度上的隐藏状态视为token维度上的信号，那么像主题和行为属性这样的宽泛概念通常会在多个位置上反映出来，并对应于该信号的低频分量。取均值充当低通滤波器，捕捉这些分量，同时衰减逐token噪声。被丢弃的高频分量对应于单个token特有的细节。  

### 3.2 混合架构  
轮次平均SAE捕捉高层特征，但无法重构逐token细节——平均丢弃了信号的高频分量。一个自然的问题是，单个模型能否分别学习捕捉这两个分量的特征：低频的轮次级结构和细粒度的逐token残差。我们探索了两种架构：一种解耦模型，在正交信号分量上训练独立的SAE；一种嵌套模型，使用类似Matryoshka的损失训练单个SAE。  

#### 3.2.1 解耦SAE  
逐token激活可以写成轮次平均隐藏状态与残差之和：  
\[
x_t = \bar{x} + (x_t - \bar{x})
\]  
解耦模型训练两个独立的SAE：一个轮次平均SAE作用于 \(\bar{x}\)，一个逐token SAE作用于残差 \((x_t - \bar{x})\)。每个SAE在其各自的输入上使用标准重构损失：  
\[
\mathcal{L}_{\text{TA}} = \|\text{dec}_c(\text{enc}_c(\bar{x})) - \bar{x}\|^2
\]  
\[
\mathcal{L}_{\text{PT}} = \|\text{dec}_f(\text{enc}_f(x_t - \bar{x})) - (x_t - \bar{x})\|^2
\]  
组合的逐token重构将两个解码相加：  
\[
\hat{x}_t = \text{dec}_c(\text{enc}_c(\bar{x})) + \text{dec}_f(\text{enc}_f(x_t - \bar{x}))
\]  
注意，轮次平均特征对轮次中的每个token贡献相同的重构，而逐token特征随位置变化。两个SAE完全解耦——它们没有共享梯度或参数。轮次平均特征完全在信号的均值分量上训练，逐token特征则专门处理残差。每个SAE使用BatchTopK，\(k = 64\)，因此每个token总共有128个特征，与其他架构的特征预算相匹配。然而，完全解耦有显著缺点。两个底层模型的分离使得典型的分析技术（如归因）无法使用——因为逐token SAE的输入 \(h_t - \bar{x}\) 对由于轮次平均特征消融引起的均匀偏移是不变的。此外，为了在固定的top-k预算下与其他架构进行比较，必须将容量分配给两个模型，这引入了一个额外的超参数，并且在相同的总k值下产生比单个模型更差的重构。  

#### 3.2.2 嵌套SAE  
注意到解耦模型的缺点，我们受到以下高层目标的启发：在单个SAE中，能否有一组特征最优地重构激活的均值，而全部特征则重构逐token激活？也就是说，某些特征被迫捕捉信号的低频分量，而所有特征则捕捉整个信号。为此，我们将SAE特征空间划分为内部分区 \([0:h]\) 和外部分区 \([h:d_{\text{sae}}]\)，并采用以下损失：  
\[
\mathcal{L} = \sum_t \|\text{dec}(\text{enc}(x_t)) - x_t\|^2 + \alpha \cdot \|\text{dec}_h(\text{enc}(\bar{x})[0:h]) - \bar{x}\|^2
\]  
其中 \(\text{dec}_h\) 仅使用前 \(h\) 个解码器方向，\(\text{enc}(\bar{x})[0:h]\) 在TopK之前将特征 \([h:]\) 掩码为零。与解耦模型不同，两个分区共享一个编码器和解码器，损失作为软约束来区分轮次平均特征和逐token特征。我们的损失公式类似于Matryoshka SAE方法[9（https://arxiv.org/html/2606.28548#bib.bib9）]，不同之处在于