VASAE：使用词汇对齐锚定命名SAE字典方向

# VASAE: 使用词汇对齐锚定命名SAE字典方向
来源：https://arxiv.org/html/2606.27941
###### 摘要

稀疏自编码器（SAE）能有效分解Transformer残差流，但学到的特征通常事后命名，而非直接与Transformer的词元词汇表建立联系。我们提出词汇对齐稀疏自编码器（VASAE），该方法在词汇对齐锚定下训练SAE特征，并为每个特征赋予一个内在词元名称：与该特征最近的词元嵌入的字符串。在不降低与标准SAE相比的重建质量的情况下，VASAE生成具有词汇对齐特征的字典。使用0.8的最近词元对齐分数截断值，在GPT-2-small后残差流上训练的字典在层0到10中约90%的特征对齐。在Llama-3.1-8B中，代表性的浅层和中层字典包含强对齐特征，其中浅层达92.8%，而代表性的最终层字典对齐有限。在减去句子级平均稀疏编码后，案例研究表明许多剩余的内在词元名称与附近的输入词元相关。这些结果表明，词汇对齐锚定可以在训练期间将学到的特征与内在词元名称连接起来，补充已学字典的事后解释。

稀疏自编码器，机制可解释性，词汇对齐

## 1 引言

Transformer架构是生成式AI的核心模型类别（Vaswani等人，2017 (https://arxiv.org/html/2606.27941#bib.bib28)；Brown等人，2020 (https://arxiv.org/html/2606.27941#bib.bib6)），并且语言模型性能随模型和数据规模提升（Hoffmann等人，2022 (https://arxiv.org/html/2606.27941#bib.bib17)）。然而，它们的行为由高维内部状态产生，难以直接检查（Belinkov和Glass，2019 (https://arxiv.org/html/2606.27941#bib.bib2)；Ghandeharioun等人，2024 (https://arxiv.org/html/2606.27941#bib.bib13)）。在仅解码器Transformer中，一个重要状态是*残差流*：每个Transformer块读取和写入的向量通道（Elhage等人，2021 (https://arxiv.org/html/2606.27941#bib.bib7)）。理解这个流很重要，因为它携带后续层计算下一个词元预测的中间信息（Geva等人，2021 (https://arxiv.org/html/2606.27941#bib.bib12)；Ghandeharioun等人，2024 (https://arxiv.org/html/2606.27941#bib.bib13)）。

稀疏自编码器（SAE）越来越多地用于将残差流分解为关于已学字典的稀疏编码（Huben等人，2024 (https://arxiv.org/html/2606.27941#bib.bib18)；Gao等人，2025 (https://arxiv.org/html/2606.27941#bib.bib11)）。在本文中，字典是解码器权重矩阵，一个*特征*是该矩阵的一列或一个方向。*稀疏编码*是特征系数的向量。每个条目对应一个特征。对于给定的残差流向量，大多数条目为零，更大的非零值表示相应特征对重建贡献更强。最近的工作使用SAE特征作为可解释性工具进行电路发现（Marks等人，2025 (https://arxiv.org/html/2606.27941#bib.bib21)；Karvonen等人，2025 (https://arxiv.org/html/2606.27941#bib.bib29)）。然而，标准SAE是为重建和稀疏性而学习特征，而非直接命名（Huben等人，2024 (https://arxiv.org/html/2606.27941#bib.bib18)；Gao等人，2025 (https://arxiv.org/html/2606.27941#bib.bib11)）。特征标签通常在训练后通过检查稀疏编码对应元素大的上下文来分配，或通过运行单独的自动解释程序（Bills等人，2023 (https://arxiv.org/html/2606.27941#bib.bib4)；Paulo等人，2025 (https://arxiv.org/html/2606.27941#bib.bib24)），使得解释成为一个独立的事后步骤。词汇对齐稀疏自编码器（VASAE）为每个学习的解码器方向分配一个训练时的最近词元标签。

技术瓶颈在于字典几何和词汇几何通常是脱节的。先前工作表明，语言模型表示空间具有非平凡的几何结构，包括各向异性和主导方向（Ethayarajh，2019 (https://arxiv.org/html/2606.27941#bib.bib9)；Mu和Viswanath，2018 (https://arxiv.org/html/2606.27941#bib.bib23)）。另外，权重绑定工作表明，输入和输出嵌入是语言模型中核心的词汇接口（Press和Wolf，2017 (https://arxiv.org/html/2606.27941#bib.bib25)；Inan等人，2017 (https://arxiv.org/html/2606.27941#bib.bib19)）。一个学到的特征可以重建激活，同时远离任何词元嵌入，因此没有内在的词汇级名称。相反的设计会强制解码器等于词元嵌入矩阵；在本文中我们测试这种硬绑定解码器是否足够。因此，缺失的目标是一个可学习的字典，既能保持重建的灵活性，又能足够接近词汇方向以提供基于词元的候选名称。词元名称是方向的一个几何派生标签。

我们的关键观察是，当词元嵌入与SAE解码器空间共享维度时，语言模型内部已经包含一组模型内部、词汇索引的参考方向：固定的输入词元嵌入（Vaswani等人，2017 (https://arxiv.org/html/2606.27941#bib.bib28)）。我们使用这些嵌入作为锚点，而非冻结的特征。基于这一观察，我们提出VASAE，这是一种SAE，其学到的特征通过通常的稀疏重建目标加上词汇对齐锚定目标来训练。每个特征被鼓励保持靠近其最近的词元嵌入，并且在训练后获得一个*内在词元名称*：其嵌入最接近该特征的词元字符串。这里“内在”意味着该名称由训练后的特征相对于模型自身词汇嵌入的位置诱导。

经验上，VASAE在测试设置下实现了与标准SAE相当的重建指标，同时添加了几何词汇对齐信号。在不同层上，它在GPT-2-small（0.9650.965）和Llama-3.1-8B（0.9310.931）上具有可比的方差解释（VE），而将解码器硬绑定到词元嵌入则重建指标较低。在GPT-2-small上，L0L0到L10L10层中8989–94%94\\%的特征超过诊断对齐截断值。使用更大的测试锚系数（λanchor=5×10−3\\lambda_{\\text{anchor}}=5\\times 10^{-3}），Llama-3.1-8B在L0L0层达到92.8%92.8\\%，但在最终代表层中不稳定。案例研究显示在上下文中分配给活跃特征的内在词元名称。

我们的贡献是：

- •我们引入VASAE：一个软词汇锚定目标，为学到的SAE解码器方向分配最近词元名称。
- •我们评估词汇锚定下的重建，并测试硬绑定解码器到词元嵌入是否足够。
- •我们分析嵌入空间特征对齐和词汇覆盖，并使用案例研究展示内在词元名称作为几何派生标识符。

## 2 背景

残差流SAE、词元嵌入和词汇读出都依赖于Transformer残差状态的几何。VASAE利用这种几何为学到的SAE解码器方向定义词汇索引的锚点。

### 2.1 Transformer残差流

Transformer模型维持一个持久的残差流，它在各层之间传递信息（Vaswani等人，2017 (https://arxiv.org/html/2606.27941#bib.bib28)；Elhage等人，2021 (https://arxiv.org/html/2606.27941#bib.bib7)）。图1 (https://arxiv.org/html/2606.27941#S3.F1)面板A显示了本设置中使用的标准Transformer结构。令V\\mathcal{V}为词汇空间，WE∈R|V|×dmodel\\mathbf{W}_{E}\\in\\mathbb{R}^{|\\mathcal{V}|\\times d_{\\mathrm{model}}}为词元嵌入矩阵，其中dmodeld_{\\mathrm{model}}为模型维度。在整个论文中，我们使用行向量约定以与实现保持一致。对于词元v∈Vv\\in\\mathcal{V}，我们用wv∈Rdmodel\\mathbf{w}_{v}\\in\\mathbb{R}^{d_{\\mathrm{model}}}表示其嵌入，并将初始后残差流写为h\(0\)=wv\\mathbf{h}^{(0)}=\\mathbf{w}_{v}。

在每一层l\\ell，残差流通过添加注意力和多层感知器（MLP）模块的输出来更新：

h\(l\)=h\(l−1\)+Attnl⁡(h\(l−1\))+MLPl⁡(h\(l−1\))。\\mathbf{h}^{(\\ell)}=\\mathbf{h}^{(\\ell-1)}+\\operatorname{Attn}_{\\ell}(\\mathbf{h}^{(\\ell-1)})+\\operatorname{MLP}_{\\ell}(\\mathbf{h}^{(\\ell-1)})。\(1\)此示意性更新省略了底层Pre-LN Transformer块的顺序子层结构和LayerNorm操作。

我们称h\(l\)∈Rdmodel\\mathbf{h}^{(\\ell)}\\in\\mathbb{R}^{d_{\\mathrm{model}}}为第l\\ell层后残差流：第l\\ell层添加其注意力和MLP输出后的残差流。在这项工作中，所有SAE变体都在后残差流上训练和评估，而不是在它们加回残差流之前的中间注意力或MLP输出上。

对于自回归语言模型，最终预测logits通过将最后一层后残差流投影经过解嵌入矩阵WU∈R|V|×dmodel\\mathbf{W}_{U}\\in\\mathbb{R}^{|\\mathcal{V}|\\times d_{\\mathrm{model}}}获得：

u=h\(L\)WU⊤。\\mathbf{u}=\\mathbf{h}^{(L)}\\mathbf{W}_{U}^{\\top}。\(2\)
嵌入和解嵌入权重绑定是语言模型的标准设计选择（Press和Wolf，2017 (https://arxiv.org/html/2606.27941#bib.bib25)；Inan等人，2017 (https://arxiv.org/html/2606.27941#bib.bib19)）。在此处使用的生成预训练Transformer 2（GPT-2）-small检查点中（Radford等人，2019 (https://arxiv.org/html/2606.27941#bib.bib26)），嵌入和解嵌入矩阵是绑定的（WU=WE\\mathbf{W}_{U}=\\mathbf{W}_{E}）。在此处使用的Llama-3.1-8B检查点中（Grattafiori等人，2024 (https://arxiv.org/html/2606.27941#bib.bib14)），它们不是绑定的。对于输入嵌入和输出解嵌入未绑定的模型，锚矩阵的选择不是唯一的。输入嵌入提供了靠近模型输入的词汇索引残差方向，而解嵌入方向可能更适合后层。在这项工作中，我们专注于输入嵌入锚点，并将基于解嵌入的锚点视为可能的扩展。

### 2.2 词汇对齐探测

Transformer残差流的一个常见诊断是将中间后残差流投影到词汇logit空间。在我们的辅助LogitLens风格读出中，与用于报告读出指标的实现匹配，我们直接通过解嵌入矩阵投影来自中间层的后残差流：

u\(l\)=h\(l\)WU⊤。\\mathbf{u}^{(\\ell)}=\\mathbf{h}^{(\\ell)}\\mathbf{W}_{U}^{\\top}。\(3\)
这衡量了在此直接读出下，哪些词元logits可以从中间后残差流线性解码。它与交叉熵（CE）恢复指标分开，后者中补丁残差流在最终logits读出前经过剩余模型计算。这个背景为VASAE使用的参考框架提供动机：词汇索引方向可用于将残差流几何与词元空间相关联。与logit镜头读出不同，VASAE不将激活解码为词元预测；它软锚定学到的SAE解码器方向到词汇索引参考方向。

### 2.3 用于残差流分析的SAE

SAE将后残差流映射到稀疏编码，然后将稀疏编码解码回残差流空间。图1 (https://arxiv.org/html/2606.27941#S3.F1)面板B显示了这种编码器-解码器结构。SAE变体在编码器参数化和强制稀疏性方式上有所不同。本文使用top-kk SAE（Makhzani和Frey，2014 (https://arxiv.org/html/2606.27941#bib.bib20)；Gao等人，2025 (https://arxiv.org/html/2606.27941#bib.bib11)），其中稀疏性通过保留ReLU后最大的kk编码器坐标来实现。

编码器使用WE∈Rdmodel×dsparse\\mathbf{W}_{\\mathcal{E}}\\in\\mathbb{R}^{d_{\\mathrm{model}}\\times d_{\\mathrm{sparse}}}和bE∈Rdsparse\\mathbf{b}_{\\mathcal{E}}\\in\\mathbb{R}^{d_{\\mathrm{sparse}}}，应用学到的仿射映射后跟ReLU和top-kk选择：

z=E\(h\)=TopKk⁡(ReLU⁡(hWE+bE))。\\mathbf{z}=\\mathcal{E}(\\mathbf{h})=\\operatorname{TopK}_{k}\\!\\left(\\operatorname{ReLU}(\\mathbf{h}\\mathbf{W}_{\\mathcal{E}}+\\mathbf{b}_{\\mathcal{E}})\\right)。\(4\)
这里E\\mathcal{E}表示完整的稀疏编码器，包括ReLU非线性和top-kk选择。向量z\\mathbf{z}是稀疏编码。每个标量zi\\mathbf{z}_{i}是稀疏编码向量中的一个元素，是特征ii的系数。算子TopKk\\operatorname{TopK}_{k}保留最大的kk坐标，将所有其他坐标设为零，因此每个稀疏编码满足‖z‖0≤k\\left\\lVert\\mathbf{z}\\right\\rVert_{0}\\leq k，其中∥⋅∥0\\left\\lVert\\cdot\\right\\rVert_{0}表示非零坐标的数量。

解码器使用WD∈Rdmodel×dsparse\\mathbf{W}_{\\mathcal{D}}\\in\\mathbb{R}^{d_{\\mathrm{model}}\\times d_{\\mathrm{sparse}}}和bD∈Rdmodel\\mathbf{b}_{\\mathcal{D}}\\in\\mathbb{R}^{d_{\\mathrm{model}}}将稀疏编码映射回残差流空间：

h~=zWD⊤+bD。\\tilde{\\mathbf{h}}=\\mathbf{z}\\mathbf{W}_{\\mathcal{D}}^{\\top}+\\mathbf{b}_{\\mathcal{D}}。\(5\)在线实验直接使用提取的后残差流；我们不在编码器之前减去单独的存储激活均值。常数偏移可以由学到的仿射编码器和解码器偏置项表示。

对于包含NN个后残差流{hi}i=1N\\{\\mathbf{h}_{i}\\}_{i=1}^{N}的训练集，其中NN计数评估的词元位置，SAE训练损失是在此稀疏编码器下的经验重建误差：

Lrecon=1N∑i=1N‖hi−h~i‖22。\\mathcal{L}_{\\text{recon}}=\\frac{1}{N}\\sum_{i=1}^{N}\\left\\lVert\\mathbf{h}_{i}-\\tilde{\\mathbf{h}}_{i}\\right\\rVert_{2}^{2}。\(6\)
稀疏性意味着对于给定的后残差流，稀疏编码中只有少数元素非零，从而产生具有小活跃特征集的分解。在3节 (https://arxiv.org/html/2606.27941#S3)中，我们保留相同的重建问题与稀疏编码，并向字典添加词汇对齐锚定。

## 3 方法

我们使用一个学到的字典分解后残差流，该字典的特征通过词汇对齐目标锚定到固定的词元嵌入方向。稀疏性通过稀疏编码引入：对于给定的后残差流，只有少数字典特征在稀疏编码中有对应的非零元素。我们提出的VASAE方法保持SAE字典可学习，并仅使用词元嵌入作为命名锚点。该方法总结在图1 (https://arxiv.org/html/2606.27941#S3.F1)中。

参见图标题图1：使用VASAE的Transformer残差流分解。模型在重建