🤖 machine learning

WriteSAE: Sparse Autoencoders for Recurrent State

本文介绍了 WriteSAE，这是首个专为分解和编辑状态空间及混合循环语言模型的矩阵缓存写入而设计的稀疏自编码器，它通过将解码器原子分解为其原生秩 1 更新形状，从而实现精确的行为干预和高精度的效果预测。

原作者： Jack Young

发布于 2026-05-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Jack Young

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一个大型语言模型（就像那些写故事或回答问题的模型）是一个庞大、高速的工厂。在这个工厂内部，有一块特殊的“记忆板”，机器在进入下一步之前会将当前的想法写在这块板上。

长期以来，科学家们一直试图使用称为*稀疏自编码器（SAEs）*的工具来理解这块板上写了什么。你可以把这些工具想象成一副眼镜，让研究人员能够看到机器正在思考的单个“想法”或“特征”。然而，存在一个问题：这些眼镜只能用于工厂的输出*，而无法用于某些先进模型中机器写入*其记忆板的具体方式。

本文介绍了一种名为WriteSAE的新工具。以下是其工作原理的简单解释：

1. 问题：眼镜的形状不对

在旧模型中，机器将想法写成一个简单的数字列表。旧眼镜（标准 SAE）被设计用来读取列表。
但在更新、更快的模型（如 Gated DeltaNet、Mamba-2 和 RWKV-7）中，机器不再写入列表。相反，它使用一种称为“秩 1 更新”的特定数学技巧，将矩阵（数字网格）写入。

类比：想象机器正在绘制一幅画。旧工具试图通过观察完成的画布来描述这幅画。但新机器是通过在复杂的网格上一次添加一个特定的笔触（单条色线）来绘画的。旧工具无法看到那个单一的笔触，因为它们寻找的是一整列颜色，而不是单一线条。

2. 解决方案：WriteSAE

作者构建了WriteSAE，这是一副专门设计用来观察那个单一笔触的新眼镜。

形状匹配：WriteSAE 的“原子”（它寻找的组成部分）不是试图读取列表，而是其形状与机器使用的笔触完全一致。它们是微小的、单一线条的图案。
结果：由于形状完美匹配，WriteSAE 能够精确隔离机器在任何给定时刻写入记忆的内容。

3. 实验：交换笔触

为了证明这行之有效，研究人员对机器的记忆进行了“手术”。

交换：他们找到了机器写入特定笔触的时刻。他们擦除该笔触，并用他们新字典（WriteSAE）中的“学习过”的笔触替换它。
测试：他们将此与另外两种情况进行了比较：
1. 完全擦除（留下空白）。
2. 放入随机的涂鸦。
结果：当他们用 WriteSAE 笔触进行交换时，机器几乎完全像以前一样继续工作（92.4% 的情况下）。当他们擦除它或使用随机涂鸦时，机器会感到困惑并犯错。
隐喻：这就像用完美契合的定制齿轮替换时钟中的特定齿轮。时钟继续滴答作响。如果你移除齿轮或放入一块随机石头，时钟就会停止。

4. 他们的发现

两种类型的笔触：他们发现机器主要使用两种笔触：
- 寄存器（Registers）：这些是精确、专注的笔触，执行特定任务（如标记句子的开头或专有名词）。
- 束（Bundles）：这些是更分散的笔触，似乎执行混合的任务。
预测未来：他们发现了一个数学公式，可以精确预测改变特定笔触将如何改变机器的下一个词。这就像知道如果微调这一个特定齿轮，时钟就会提前一秒报时。
编辑机器：他们成功利用该工具“安装”了新行为。例如，他们可以通过将正确的笔触插入记忆板，强制机器持续谈论一个特定主题（如一个“中位”词），而这是它通常不会选择的。

5. 局限性

本文非常谨慎地指出，该工具最适用于以这种特定“单笔触”（秩 1）方式写入的模型。

如果模型以更复杂的方式写入（例如同时使用两个笔触或对角线模式），该工具的效果就不那么完美，尽管它仍能发现一些模式。
该工具在 Qwen3.5 模型（一种特定类型的 AI）上表现极佳，并且他们展示了它也适用于其他类似模型，但预测未来的“魔法公式”会根据模型架构的不同而变化。

总结

WriteSAE是一种新工具，让我们能够观察和编辑先进 AI 模型写入其内部记忆的具体方式。通过将工具的形状与机器写入的形状相匹配，研究人员可以交换特定的想法，预测机器的反应，甚至引导机器说出它通常不会说的话，而不会破坏机器。这是科学家首次成功地对这类特定模型的内存写入点直接进行此类“手术”。

技术摘要：用于循环状态分解的 WriteSAE

问题陈述

状态空间模型和混合循环语言模型（例如 Gated DeltaNet、Mamba-2、RWKV-7、Qwen3.5）利用矩阵缓存进行状态传播，其中每个 token 将一个秩为 1 的外积（ $k_t v_t^\top$ ）写入 $d_k \times d_v$ 矩阵中。现有的稀疏自编码器（SAE）作用于残差流，在输出后读取状态。因此，标准的 SAE 无法直接访问或分解发生在循环上游的原生写入机制。此外，由于状态是通过秩为 1 的外积更新的，标准的基于向量的 SAE 解码器原子无法在不破坏因果干预（缓存修补）所需结构对应关系的情况下，替换单个缓存写入。

方法论

本文介绍了 WriteSAE，这是首个旨在分解和编辑循环模型矩阵缓存写入的稀疏自编码器。

架构与训练

原生形状匹配：WriteSAE 的解码器原子被约束为秩为 1 的外积（ $v_i w_i^\top$ ），以匹配宿主架构的原生写入基元（ $k_t v_t^\top$ ）。这使得单个原子能够替换特定的缓存槽位更新。
双线性编码器：为了高效编码矩阵状态 $S_t$ ，该模型使用双线性编码器（ $a_i = v_i^\top S_t w_i$ ），而非展平的全连接编码器。这将每个原子的参数量从 $d_k d_v$ （对于 $128 \times 128$ 为 16,384）减少到 $d_k + d_v$ （256），实现了 64 倍的缩减。
训练目标：模型在去均值状态（ $x = \text{vec}(S_t - M)$ ）上进行训练，使用 TopK 稀疏约束和辅助损失以复活死特征。训练语料库由通过特定层和头处理的 Qwen3.5-0.8B（及其他模型）的 OpenWebText 序列组成。

特征分类

根据原子写入方向与原生写入方向之间的中值余弦相似度，将训练好的字典划分为两类：

寄存器（Registers）：与原生写入具有高余弦对齐度的原子（可从缓存中恢复）。
束（Bundles）：在缓存中具有分散写入方向的原子。

因果干预测试

本文提出了三种主要的验证机制：

缓存槽位替换：用匹配弗罗贝尼乌斯范数的 SAE 原子替换原生写入（ $\hat{S}_t = S_t - \Delta_{nat} + \Delta_{atom}$ ），并测量下游的 KL 散度。
闭式 Logit 偏移：推导一个三因子表达式以预测每个 token 的 logit 偏移： $\Delta \ell \approx G_{t_0 \to t} \langle w_i, q_t \rangle \langle v_i, W_U[\text{tok}] \rangle$ ，其中 $G$ 是门控乘积， $q_t$ 是读取查询， $W_U$ 是未嵌入行。
直接缓存干预：在特定位置擦除或安装原子，以观察 token 概率和生成行为的变化。

主要贡献

架构匹配字典：一个原子为秩为 1 外积的字典，能够直接替换原生缓存写入，这是残差流 SAE 此前不具备的能力。
替换验证：证明了学习到的原子可以高保真地替换原生写入。在 Qwen3.5-0.8B（第 9 层，第 4 头）上，原子在 4,851 次触发中击败了匹配范数消融法，占比 92.4%。
闭式预测：一个无参数、三因子的公式，可预测每次触发的 logit 偏移，在 200 个“原子 - epsilon"单元中，中值 $R^2$ 达到 0.98。
跨架构迁移：验证了替换成功率和特征划分（寄存器与束）在不同的矩阵循环架构（Gated DeltaNet、RWKV-7、Mamba-2）中均成立，尽管成功率随底层的“写入秩”而变化。

实验结果

替换性能

Qwen3.5-0.8B (Gated DeltaNet)：在 L9 H4 处，原子在 92.4% 的触发中击败了消融法。在 87 个原子的群体测试中，该比例保持在 89.8%。严格的排序 $KL_{atom} < KL_{ablate} < KL_{random}$ 在 89.5% 的触发中成立。
Mamba-2-370M：尽管 Mamba-2 使用对角状态更新而非完整的秩为 1 外积，但在 2,500 次触发中，替换成功率仍达到 88.1%。
RWKV-7：表现出中等性能，与其秩为 2 的写入规则一致。

Logit 偏移预测

闭式表达式准确跟踪了测量到的效应。在 L9 H4 处，预测与测量的 logit 偏移之间的中值 $R^2$ 为 0.98。然而，该闭式形式在 Mamba-2 和 Qwen3.5-4B 上失效（ $R^2$ 为负），表明门控因子 $G$ 是特定于底层的。

行为干预

擦除：在其自然触发位置擦除特定的寄存器原子（F412），使目标 token（"space"）的对数概率中值降低了 -0.116 nats（ $p=1.07 \times 10^{-6}$ ）。
生成引导：在三个连续位置持续安装闭式方向，在贪婪解码下，将中位 rank token 的“目标 - 延续”率从 33.3% 提升至 100%。
段落级放大：在 4B 模型中放大与边界相关的特征，在 5 倍剂量下将换行生成减少了 33%（从每 400 个 token 的 16.8 次降至 11.2 次）， demonstrating 了对文档结构的控制。

架构范围

本文建立了“写入秩”有效性的层级：

Gated DeltaNet (秩 -1)：最高的寄存器余弦分离度（0.262）和替换保真度。
RWKV-7 (秩 -2)：中等分离度（0.180）。
Mamba-2 (对角/标量)：最低分离度（0.0575），但仍实现了约 88% 的替换成功率。

意义与主张

本文声称提供了首个针对矩阵循环写入位点的因果字典，超越了残差流分析。它证明了：

写入秩是决定性的：将状态分解为可解释的“寄存器”的能力与原生写入规则的秩相关（秩 -1 > 秩 -2 > 对角）。
因果替换是可能的：学习到的原子可以替换原生写入，且下游 KL 散度最小，验证了循环模型中“寄存器”的概念。
闭式可解释性：对于 Gated DeltaNet，特征对 logit 的影响可以在不拟合参数的情况下进行解析预测，将字典直接与模型的门控和读取机制联系起来。

作者谦逊地指出了局限性：闭式系数无法在不同底层之间迁移（例如在 Mamba-2 上失效），且每个原子的身份是特定于种子的（尽管类别级别的属性是稳定的）。这项工作为状态空间模型的机制可解释性建立了新的基准，表明“写入”基元是字典学习和因果编辑的可行目标。