使用分块SVD寻找可用的权重机制

# 1 引言

Source: https://arxiv.org/html/2608.06969

Finding Usable Weight Mechanisms with Tiled SVD

Ash Manvi

Aquin Labs

ash@aquin\.app

Samreena Tajreen

Aquin Labs

samreena@aquin\.app

摘要

> 机制可解释性的主流方法是训练代理字典（如稀疏自编码器），并从最大激活文本中标注特征。最好的这类图谱能够识别概念，但该身份存在于学习到的字典中，而非网络权重本身。我们提出直接从线性位点通过按列分块 SVD（column-tiled SVD）提取“机制挂载点”（mechanism mounts）：每个挂载点是一个三元组 \((v,u,\sigma)\)，分别读作触发器、写入方向和强度。身份就是权重规则本身。我们使用一套预先注册的评测套件来评估挂载点，判定标准是全写能量提升（full-write energy lift）而非分块局部提升（tile-local lift）。在 Gemma-2-2B 与 WikiText-2（16,384 token 子采样）上，七个线性映射均被评分：残差写入（mlp.down, attn.o）获得完整的 A/B/C 评分，可在子层后 RMSNorm 之后进行引导（steer），并通过 52/52 个位点层；其他映射仅获得 A/B 评分（mlp.gate/attn.q/attn.k/effective mlp.up/attn.v 各 26/26）。总计：182/182 通过。我们发布了库代码、语料构建器、实验入口和单元测试。

稀疏自编码器及相关的代理字典已成为机制可解释性中为神经网络中的方向命名的标准工具 [1 (https://arxiv.org/html/2608.06969#bib.bib1), 2 (https://arxiv.org/html/2608.06969#bib.bib2)]。它们在激活上训练，并从最大激活文本中标注，产生的概念图谱对于描述模型表示的内容很有用。此后，大量工作不断改进字典学习、扩展和自动标注，同时也提高了对多语义神经元的覆盖。

这些图谱通常通过学到的字典原子及其语言标签来识别一个特征。将解释对齐到一个单独训练的码本，而不是网络内部的某个特定权重矩阵，意味着“这个方向是什么意思”这个问题是在代理空间中回答的。实际写入残差流的权重规则仍然是隐式的。最近的工作表明，通过 unembedding 读出的 MLP 和注意力矩阵的奇异向量通常能形成可解释的 token 簇，并且可以被编辑 [3 (https://arxiv.org/html/2608.06969#bib.bib3)]，同时也有工作将奇异模式框架化为检测器-效应器单元 [4 (https://arxiv.org/html/2608.06969#bib.bib4)]。然而，在大多数此类设置中，SVD 被用作一种透镜或电路原语 [5 (https://arxiv.org/html/2608.06969#bib.bib5)]，而不是作为对 \(W\) 的哪种分块方式能产生可用的在分布（on-distribution）机制的一种公平测试。

在这项工作中，我们提出直接从线性位点通过按列分块 SVD 提取“机制挂载点”。每个挂载点是一个三元组 \((v,u,\sigma)\)，分别读作触发器、写入方向和强度；身份就是权重规则本身。我们用一个测量栈来评估挂载点，其判据是全写能量提升（full-write energy lift）而非分块局部提升——后者会同义反复地偏好单列分块——再加上覆盖饱和，以及针对残差写入的、以深度为条件的、与最终 unembedding 对齐的引导检查。在 Gemma-2-2B 上，使用 WikiText-2 子采样，该套件覆盖每一层的全部七个线性映射，并在 \(mlp.up\) 和 \(attn.v\) 使用有效路径挂载点后通过了全部 182 个位点层。

## 2 方法

### 2.1 模型与位点

我们评估 `google/gemma-2-2b`（26 层，索引 \(0,\ldots,25\)）。每一层暴露七个线性映射。其中两个是*残差写入*：`mlp.down`（\(W\in\mathbb{R}^{2304\times 9216}\)）和 `attn.o`（\(W\in\mathbb{R}^{2304\times 2048}\)）。它们接受完整的分块、覆盖和因果检查套件（A/B/C），其中实验 C 在 Gemma-2 子层后 RMSNorm 之后注入，使被引导的方向进入残差流。其余五个映射（`mlp.gate`, `mlp.up`, `attn.q`, `attn.k`, `attn.v`）仅接受 A/B：它们的输出不是残差写入，因此 unembed 对齐不是正确的因果度量。相同的全写能量提升和覆盖标准适用于所有位点；只有 C 依赖于位点类型。

语料文本来自 WikiText-2（原始训练集），由 `scripts/build_corpus.py` 构建。前向过程收集所有 token，然后用种子 0 子采样到 16,384 个 token 以节省内存。

参见图 1 标题：图 1：一个 Gemma 风格解码器块中的线性位点（VisualTorch 的缩放替身）。橙色 ABOnlyLinear：attn.q/k/v, mlp.gate/up（仅 A/B）。绿色 ResidualWriteLinear：attn.o 和 mlp.down（A/B/C 残差写入）。

### 2.2 分块 SVD 挂载点

设 \(T\) 为分块宽度，\(k\) 为每个分块的模态数。默认值与位点相关：残差和 MLP 映射 \(T=512\)，`attn.k`/`attn.v` \(T=256\)，`attn.q` \(T=128\)（在困难层有 64 的回退）。实验 A 和 C 使用 \(k=2\)；实验 B 扫描 \(k\in\{1,2,4,8,16\}\)。

将 \(W\) 的输入列划分为分块 \([s_t,e_t)\)，并对每个分块进行分解：

\[
B_t = W_{:, s_t:e_t} = U_t \Sigma_t V_t^\top.
\]

对于模式 \(i<k\)：

\[
u = \frac{U_t[:,i]}{\|U_t[:,i]\|_2},\quad \sigma = \Sigma_t[i],\quad v = (V_t^\top)[i,:].
\]

一个挂载点是一个三元组 \((v,u,\sigma)\)，并带有位点和层元数据。身份就是这条权重规则，而不是语言标签。挂载点 id：`tile:{t}:sv{i}`。

在匹配的挂载点预算下，我们比较四种构造：

参见图 2 标题：图 2：线性权重的按列分块 SVD。每个分块产生模态 \((v,u,\sigma)\)，分别读作触发器、写入方向和强度。

### 2.3 触发与全写能量提升

在真实前向中，给定位点输入 \(x\)，触发系数为：

\[
a_{t,j} = x_{t, s_j:e_j} \cdot v_j.
\]

分块写入使用对应的列块：

\[
\Delta h^{\mathrm{tile}} = x_{:, s:e} B^\top.
\]

SVD 恒等式（健全性检查，而非证明）要求 \(a\) 与 \(\Delta h^{\mathrm{tile}} u\) 的相关性高于 0.99，相对斜率误差低于 0.05。这个恒等式几乎总是成立的，因此无法区分可用挂载点和未使用挂载点。

分块局部能量提升

\[
L_{\mathrm{tile}} = \mathbb{E}_t\!\left[\frac{(a_t \sigma)^2}{\|\Delta h^{\mathrm{tile}}_t\|_2^2 + \varepsilon}\right] - \mathbb{E}_t\!\left[\frac{(\Delta h^{\mathrm{tile}}_t \cdot \tilde{u})^2}{\|\Delta h^{\mathrm{tile}}_t\|_2^2 + \varepsilon}\right]
\]

同义反复地偏好单列分块（列采样时 \(L_{\mathrm{tile}}\approx 1\)）。因此，通过标准使用对位点写入张量 \(\Delta h\) 的*全写*能量提升：

\[
L_{\mathrm{full}}(u) = \mathbb{E}_t\!\left[\frac{(\Delta h_t \cdot u)^2}{\|\Delta h_t\|_2^2 + \varepsilon}\right] - \mathbb{E}_t\!\left[\frac{(\Delta h_t \cdot \tilde{u})^2}{\|\Delta h_t\|_2^2 + \varepsilon}\right],
\]

其中随机方向 \(\tilde{u}\) 使用种子 \(10007 + j \cdot 997\) 生成。

### 2.4 覆盖饱和

权重覆盖是指每个分块的 rank-\(k\) 重建所保留的 \(\|W\|_F^2\) 比例。稀疏写入覆盖构建一个单位挂载方向字典，每个 token 选择 top-\(k_{\mathrm{active}}\)（默认 8）个挂载点，用最小二乘法重建 \(\Delta h\)，并报告解释能量。覆盖提升是与匹配大小的随机字典之间的差距。饱和（B1）要求提升峰值达到位点相关的下限（残差 0.25，其他映射 0.15，有效 up/v 路径 0.08），早期扫描窗口位于峰值的 0.08 以内（残差 \(m\in\{1,2\}\)；其他 \(m\in\{1,2,4\}\)），并且最终点不低于峰值 0.10。

### 2.5 针对 unembed 的因果引导

仅对残差写入，我们通过 unembed 透镜读取写入方向 \(u\) 的最终对数几率几何：先近似最终 RMSNorm，然后 \(t = W_{\mathrm{lm}} u\)。引导是在 post-attention 或 post-FF 的 RMSNorm 模块上添加 \(\alpha u\)，其中 \(\alpha=2\)（不是仅在 `o_proj`/`down_proj` 上），在最多八段文本上对最后 token 的 \(\Delta\) 对数几率取平均，并报告 Spearman \(\rho(\overline{\Delta}, t)\) 和 top-20 Jaccard。C1 通过条件为 \(\rho \ge 0.05\) 或 \(J_{20} \ge 0.05\)，并且仅当位点是残差写入且层 \(\ell \ge 6\) 时*要求*通过。早期残差层仍会报告 C；当对齐较弱时它们不会失败。非残差位点跳过 C。

### 2.6 有效路径挂载点

原始的 `mlp.up` 和 `attn.v` 模块权重无法通过残差形态的 A/B：实际在分布中使用的映射并非原始矩阵。因此默认从有效映射进行挂载。对于 `mlp.up`，\(W^{\star} = \mathrm{diag}(\bar{g})\, W_{\mathrm{up}}\)，其中 \(\bar{g}\) 是语料平均的门控激活（回退方案：门控混合分块 SVD；或经 down 合成 \(W_{\downarrow} \mathrm{diag}(\bar{g}) W_{\mathrm{up}}\)）。对于 `attn.v`，使用岭最小二乘 \(x \mapsto\) mixed-\(v\)（回退方案：\(W^{\dagger} = W_o \, \mathrm{expand}_{\mathrm{GQA}}(W_v)\)）。评分使用对应的写入张量（门控乘积或 mixed-\(v\)/\(o\) 写入）。

参见图 3 标题：图 3：有效路径挂载点。左：`mlp.up` 的均值门控 \(W^{\star}\)。右：`attn.v` 的岭 lstsq \(x \mapsto\) mixed-\(v\)（合成-\(o\) 回退）。

### 2.7 通过标准

当一个位点层的所有适用检查均通过时，该位点层被接受：

判定器位于 `src/atlas/mount/paper_eval.py`。运行器搜索位点相关的分块大小和有效路径池，保留最佳通过试验，并聚合全部七个位点 \(\times\) 26 层的结果。

## 3 实验

### 3.1 设置

我们构建 WikiText-2 语料，并在模型加载后运行完整套件：

```
python scripts/build_corpus.py --out data/corpus/train.jsonl
python scripts/run_paper_experiments.py --layers all --sites all \
  --device cuda --texts data/corpus/train.jsonl \
  --out-dir data/eval/paper_experiments_all
```

默认设置：位点相关的分块大小，A/C 每个分块 \(k=2\) 个模态，B 的模态扫描为 \(\{1,2,4,8,16\}\)，\(n_{\mathrm{steer}}=8\)，\(\alpha=2\)，从收集到的 86,109 个 token 中子采样 16,384 个 token。输出位于 `{site_slug}/L{n}/` 下，并包含聚合文件 `sites.csv`。

报告的运行结果：全部七个位点 \(\times\) 26 层均通过：182/182（残差 A/B/C 52/52；其他 A/B 130/130）。

### 3.2 实验 A：分块

在两个残差写入位点上，分块全写提升在每个深度都超过全矩阵 SVD、列采样和随机，除了 `mlp.down` 第 25 层，在那里分块 \(\approx\) 全矩阵，而 A4 仍按比例通过。`attn.o` 使用的挂载点更少（\(d_{\mathrm{in}}=2048\)）但仍在所有位置赢得 A。列采样的 `tile_lift≈0.999` 被忽略；A1–A4 只使用全写提升。

参见图 4 标题：图 4：实验 A。残差写入的均值全写能量提升随深度变化。分块 SVD 超过全矩阵 SVD、列采样和随机。`W` 中的局部列结构并不能被单个全局 SVD 很好地概括出在分布写入能量。在固定挂载点数量下，分块在两个残差写入上都恢复了更高能量的写入方向。

### 3.3 实验 B：覆盖与模态

在残差写入上，覆盖提升在 \(m=1\) 或 \(m=2\) 时已经很高，之后趋于平坦：B1 在早期饱和规则下通过全部 52 个残差位点层。额外的模态主要增加冗余。

参见图 5 标题：图 5：实验 B。覆盖提升与每分块模态数的关系。提升在 \(m=1\) 或 \(m=2\) 时已经很高，之后持平（饱和）。

### 3.4 实验 C：因果深度

使用 post-norm 残差注入后，与 unembed\((u)\) 的平均 Spearman 相关性在两个残差位点上都随深度上升。中深度的 `attn.o` 不再失败 C1。最终层达到 \(\rho \approx 0.91\)（`mlp.down`）和 \(\rho \approx 0.75\)（`attn.o`）。对 \(\ell<6\) 的 C1 豁免仍然由这条曲线支持。

参见图 6 标题：图 6：实验 C。post-norm 注入后，被引导的 \(\Delta\) 对数几率与 unembed\((u)\) 的平均 Spearman。灰色带：\(L<6\)（C1 豁免）。

### 3.5 总体结论

使用全写能量提升、早期覆盖饱和和深度感知的 post-norm 引导与 unembed 对比，分块 SVD 挂载点在全部七个线性映射的每个位点层上都通过，前提是 `mlp.up` 和 `attn.v` 使用有效写入映射而非原始模块权重。

## 4 讨论

测量栈是这项工作中最持久的部分。全写能量提升在匹配预算下区分了分块 SVD 与全矩阵 SVD、列采样和随机，同时不奖励抬高局部提升的单列同义反复。残差写入上的覆盖早期饱和。使用 post-norm 注入后，引导与 unembed 的对齐在 `mlp.down` 和 `attn.o` 上都呈现出清晰的深度曲线，并且同一个判定器在 `mlp.up` 和 `attn.v` 从它们各自的有效写入映射挂载后接受了全部七个线性映射。这是一个诚实的多位点协议，而不是仅限 MLP 的演示。

新颖性较薄。Transformer 权重的奇异向量、检测器-效应器单元以及 unembed 读出已经存在 [3 (https://arxiv.org/html/2608.06969#bib.bib3), 4 (https://arxiv.org/html/2608.06969#bib.bib4), 5 (https://arxiv.org/html/2608.06969#bib.bib5)]。我们并不声称提供人类可读的概念名称，也不声称要取代用于概念发现的稀疏自编码器 [1 (https://arxiv.org/html/2608.06969#bib.bib1), 2 (https://arxiv.org/html/2608.06969#bib.bib2)]。我们的切入点在于公平的分块、关于分块局部度量的负结果、覆盖饱和，以及包装为可复现套件的深度条件因果检查。

范围仍然清晰。实验 C 仅适用于 \(u\) 是残差方向的情况。原始 `mlp.up`/`attn.v` 模块权重按设计无法通过残差形态的 A/B；受支持的对象是有效路径。所有数字均针对 Gemma-2-2B 在 WikiText-2 子采样上的结果。

## 5 局限性

本研究使用单一模型家族和规模（Gemma-2-2B）。实验 C 仅适用于残差写入位点；gate、up、q、k、v 仅以 A/B 判定。WikiText-2 子采样（86,109 个 token 中的 16,384 个）可能会影响哪些挂载点看起来更强。能量提升不是人类含义：挂载点不携带语义标签。引导使用短文本、固定 \(\alpha=2\) 以及子层后 RMSNorm 之后最后 token 的对数几率。对 \(\ell<6\) 的 C1 豁免从深度曲线看是有原则的，但仍然是一个设计选择；早期 \(\rho\) 应始终报告。原始 `mlp.up` 和 `attn.v` 模块权重无法通过残差形态的 A/B，需要有效路径挂载点。`mlp.down` 第 25 层在 A4 上处于边缘（分块 \(\approx\) 全矩阵），但仍通过；`attn.o` 运行的挂载点比 `mlp.down` 更少。

## 6 结论

我们从 Gemma-2-2B 的线性位点提取分块 SVD 机制挂载点 \((v,u,\sigma)\)，并用全写能量提升、覆盖饱和以及针对残差写入的、以深度为条件的 unembed 引导检查来评分。在匹配的挂载点预算下，分块优于全矩阵 SVD、列采样和随机；写入覆盖在每个分块一到两个模态时即饱和；post-norm 残差引导与 unembed\((u)\) 的对齐随深度增强。在全部七个线性映射和 26 层上，一旦 `mlp.up` 和 `attn.v` 使用有效路径挂载点，套件通过 182/182 个位点层。我们发布了库、语料构建器、实验入口和单元测试，其中挂载点身份被定义为权重规则本身。

## 致谢