从神经网络中定向恢复权重空间机制

# 从神经网络中定向恢复权重空间机制 来源：https://arxiv.org/html/2607.13047

###### 摘要

参数分解（PD）将神经网络分解为可解释的计算组件，这些组件忠实反映原始网络的操作。然而，将PD扩展到大型模型需要大量计算，使其成为一项成本高昂且风险较大的工作。本文提出靶向PD（tPD），该方法通过引入一个高秩的"兜底"组件来处理所有非目标数据，从而仅识别处理特定感兴趣输入的组件——从单个提示词到大型子任务。我们在玩具模型和基于The Pile训练的Transformer语言模型上验证了tPD，它能够恢复可重复、机制上忠实的电路。我们使用约7%的已发表分解所需的FLOPs，提取了一个4块Transformer的纯CSS子模型，并在12块Transformer中精准消融和重连记忆化序列，对其他输入的副作用可忽略不计。

机理可解释性、参数分解、电路、语言模型

## 1 引言

参数分解（PD）旨在通过逆向工程恢复神经网络实现的算法，而非注释其内部表示 (Braun et al., 2025 (https://arxiv.org/html/2607.13047#bib.bib2); Bushnaq et al., 2025 (https://arxiv.org/html/2607.13047#bib.bib4), 2026 (https://arxiv.org/html/2607.13047#bib.bib3))。每个权重矩阵被分解为由秩1子组件U_c V_c^T组成的机制，每个机制实现单一功能角色。分解必须满足四个要求：(i) 子组件的和等于原始权重；(ii) 非活跃子组件可以任意组合消融而不改变模型输出；(iii) 每个子组件的计算尽可能简单；(iv) 每个输入激活最少数量的组件 (Bushnaq et al., 2025 (https://arxiv.org/html/2607.13047#bib.bib4), 2026 (https://arxiv.org/html/2607.13047#bib.bib3))。一个联合训练的辅助网络学习哪些子组件在哪些输入上是活跃的（即因果重要的）。在玩具模型上的初始工作之后，这个复杂的优化问题最近被扩展到包含2800万非嵌入参数的4块Transformer (Bushnaq et al., 2026 (https://arxiv.org/html/2607.13047#bib.bib3))。要达到前沿LLM需要数个数量级的规模扩展，且没有真实基准来指导工程。本文证明PD框架无需完整模型分解即可发挥作用。我们引入**靶向参数分解**（tPD），该方法选择性地识别处理选定数据子集（即**目标数据**）的机制，大幅降低计算需求。

参见图注
参见图注

图1：A：TMCC玩具模型第一个矩阵的因果重要性模式：全数据和针对3个特征的靶向分解。B：当提供可变数量的子组件槽位时，TMCC玩具模型靶向分解的重建均方误差（MSE），作为分解成功的代理指标。C：TMCC全数据分解和靶向分解在训练过程中的MSE。D：针对 numpy/pandas 联合任务，两个随机种子间匹配组件的U和V向量的余弦相似度。水平条表示均值。灰色线和阴影区域表示通过将第一个分解与第二个分解未经训练、随机初始化的版本进行比较得到的零分布的均值±标准差。作为对比，同一模型不同种子的全数据分解 (Bushnaq et al., 2026 (https://arxiv.org/html/2607.13047#bib.bib3)) 在U和V向量上分别产生0.48和0.51的平均最大余弦相似度。E：嵌套目标集分解得到的子组件之间的成对匹配：`import numpy as`（{np}）、`import pandas as`（{pd}）或同时两者（{np, pd}）。

## 2 方法

一个自然的起点是仅使用目标数据运行PD。这能恢复在目标数据上正确重建模型输出的组件 (Christensen & Riggs, 2025 (https://arxiv.org/html/2607.13047#bib.bib5))，但这会使解欠定：在狭窄的目标数据集上，激活只覆盖模型在通用数据上所见空间的一个子空间。组件在该子空间内受到约束，但在正交方向上则是任意的。因此，它们的消融会对非目标输入产生不可预测的影响——阻碍模型编辑，并阻止我们检查它们被触发的更广泛上下文。为获得可解释的机制，tPD在两条并行数据流上训练PD——一个较小的**目标**数据集（我们想要解释的输入）以及一个从模型原始训练分布中采样的**非目标**数据流。分解模型被训练为在两条数据流上重建模型的输出：
- • 对于目标数据，重建必须仅使用一组稀疏的秩1子组件，如同标准PD。然而，这些子组件不一定需要求和为原始权重——它们只需捕捉在目标数据上触发的机制。
- • 对于非目标数据，每个权重矩阵被分配一个满秩的**兜底**组件Δ，定义为原始权重与子组件之和的残差。这些Δ组件捕捉所有在目标数据上从不活跃的机制，因此它们不被显式分解。处理目标数据的子组件也可用于非目标重建，因为某些机制可能在目标和非目标数据上都活跃。为实现这一点，Δ组件在目标批次上被对抗性消融（每次在0到100%之间），以最大化重建损失，迫使目标输出仅由子组件重建。在非目标批次上，Δ组件始终完全启用。重要的是，目标和非目标数据流都对重要性最小化损失有贡献。这有两个后果：(i) 任何在非目标数据上触发但不在目标数据上触发的机制被移入Δ组件（这些不计入重要性最小化损失）；(ii) 处理目标数据的子组件被塑造为尽可能不干扰非目标数据。为确保分解后的电路机制上忠实于原始计算，非活跃组件和Δ组件都被对抗性消融以最大化重建损失，迫使模型输出对于原始矩阵和消融矩阵的任何混合保持不变。更多实现细节见附录A (https://arxiv.org/html/2607.13047#A1)，关于“TMS-5-2-id”玩具模型 (Bushnaq et al., 2025 (https://arxiv.org/html/2607.13047#bib.bib4)) 的完整工作示例见附录B (https://arxiv.org/html/2607.13047#A2)，其中说明了为什么需要非目标批次和Δ组件来捕捉正确的机制。

## 3 结果

### 3.1 tPD 在玩具模型上快速恢复真实机制

我们首先在玩具模型上验证tPD。为展示计算收益，我们使用Braun等人 (2025 (https://arxiv.org/html/2607.13047#bib.bib2)) 提出的压缩计算玩具模型（TMCC）。在该模型中，参数分解识别出MLP输入矩阵的100种机制：每个输入特征对应一种（图1A，左）。我们仅使用输入5、10和15作为目标数据运行靶向版本。正如预期，我们精确恢复了处理这些特征的三种机制（图1A，右）。关键的是，CI网络学习了正确的激活模式，因此子组件不会在非目标输入上虚假激活。计算成本降低有两个原因：
- • 首先，分解模型只需足够大以容纳目标机制，内存占用大大减少。虽然完整的TMCC每矩阵至少需要100个子组件槽位，但针对3个特征子集的tPD可使用小得多的模型（最少5个子组件槽位，图1B）。
- • 其次，我们发现tPD收敛更快，大概是因为这是一个简单得多的优化问题。在其他条件相同的情况下，单层TMCC的全数据PD需要约2500个优化步骤才能收敛；针对3输入目标的tPD在约500步内收敛（图1C）。

### 3.2 Transformer上的一致性检验

接下来，我们在Transformer语言模型上验证该方法，其训练数据远比玩具模型复杂。虽然没有真实基准可供比较，但我们运行了两个一致性检验：第一，如果tPD发现了可辨识的机制，从不同随机种子初始化的分解应收敛到相似的电路。第二，嵌套集合的分解应具有连贯性：处理A∪B的电路应该是处理A的电路和处理B的电路的并集（由于高秩机制可能存在例外，见附录B (https://arxiv.org/html/2607.13047#A2)-C (https://arxiv.org/html/2607.13047#A3)）。我们在目前存在全数据分解的最大模型上执行此验证：一个在The Pile上训练的4块Transformer (Bushnaq et al., 2026 (https://arxiv.org/html/2607.13047#bib.bib3))。为在极端相反的情况下验证tPD，我们仅针对两个输入进行分解：“import numpy as”和“import pandas as”。这些序列在Python代码中无处不在，分解模型正确地将它们补全为“np”和“pd”。使用这两个提示作为目标，通用Pile作为非目标，我们执行tPD两次，从两个不同随机种子初始化。然后我们基于权重的最大余弦相似度匹配子组件，并计算它们各自U和V向量之间的余弦相似度（图1D）。总体而言，一致性很强。特别是V向量（输入方向）高度相似，表明两个分解收敛到基本相同的实现。注意力层Q和K矩阵中的U向量一致性稍低，表明用于形成注意力模式的表示约束稍弱。注意，子组件在不同初始化下相同与否可能有多种原因。进一步讨论见附录C (https://arxiv.org/html/2607.13047#A3) 以及Bushnaq等人 (2026 (https://arxiv.org/html/2607.13047#bib.bib3)) 的附录B.3。

为检查嵌套目标间的一致性，我们分别针对两个提示中的每一个运行靶向分解，并将获得的子组件与同时针对两个提示的分解得到的子组件进行比较。在联合分解 {np, pd} 的77个子组件中，有69个在其中一个或两个单提示分解中有强匹配：18个与“numpy”分解 {np} 共享，19个与“pandas”分解 {pd} 共享，32个似乎被两个提示共同使用（图1E）。

参见图注

图2：A：仅由CSS目标分解识别出的子组件构成的模型，与原始模型在各种语言上的KL散度。值越小表示输出越接近原始模型。B：原始序列、原始模型预测的下一个词元以及纯CSS模型预测的下一个词元的比较，样本来自CSS代码、Python代码和英语文本。提示行左移一个位置，使真实下一个词元与预测的下一个词元对齐。高亮部分为原始模型和纯CSS模型之间每个词元的KL散度，截断至[0,1]。

参见图注

图3：来自纯CSS分解的三个CSS特异性子组件示例。左：在CSS代码与其他语言上，当子组件被单独从模型中减去时，与原始模型输出的KL散度。右：这些组件的激活模式（显示为消融时每个词元的KL散度，截断至[0,1]）。

### 3.3 提取纯CSS子模型

在确定tPD产生一致的分解之后，我们转向一个更广泛的案例：提取处理CSS代码的每一个机制。使用相同的4块Transformer，我们将目标设置为The Pile中限制为CSS代码的子集，并去除注释以避免拾取普通语言机制。使用已发表全数据分解约7%的FLOPs（附录D (https://arxiv.org/html/2607.13047#A4)），我们匹配并超过了其重建精度（平均对抗性KL散度：0.45，因果重要性L0范数：89.6）。tPD返回1,638个子组件，而全数据分解为9,972个。这1,638个子组件推测捕捉了模型用于预测CSS代码的算法。为验证这一点，我们使用仅有这些组件的“纯CSS”模型，在来自The Stack (Kocetkov et al., 2022 (https://arxiv.org/html/2607.13047#bib.bib15)) 和WikiText (Merity et al., 2016 (https://arxiv.org/html/2607.13047#bib.bib21)) 的多种语言面板上进行测试。如图2A所示，纯CSS模型在CSS数据上的行为与原始模型相似（KL散度约0.6）。相比之下，在每种其他语言上的行为都崩溃了（KL散度≥2），并且模型倾向于回退到单个主导词元（英语为“the”，Python为“count”）。

使用非目标批次和Δ组件进行靶向PD的一个声称的好处是，识别出的机制对目标数据之外的样例应无副作用，除非它们被积极用于处理这些样例。如果确实如此，我们应该能够选择性地从模型中摧毁CSS特定能力，同时保留其他能力。这并不简单，因为许多在CSS上活跃的子组件也在非CSS上下文中活跃，因此简单地消融全部1,638个子组件并不可行。我们首先筛选出在CSS数据上频繁触发但在通用Pile数据上很少触发的子组件。对该短名单的手动检查识别出具有特定句法角色的子组件，例如数字前加单位或用于缩进的空白（图A4 (https://arxiv.org/html/2607.13047#A6.F4)、A5 (https://arxiv.org/html/2607.13047#A6.F5)、A6 (https://arxiv.org/html/2607.13047#A6.F6)）。引人注目的是，它们高度特定于CSS上下文：例如，第3块MLP-up中的子组件98在CSS代码中对缩进激活，但在Python代码中对缩进不激活（图A6）。我们推测这些子组件除了识别表面模式外，还编码了某种形式的CSS上下文意识，尽管我们将对此工作的更机械理解留给未来研究。因此，消融这些子组件会降低模型在CSS数据上的输出，但保持其他语言不变（图3左）。

参见图注

图4：A：消融单个子组件的效果，这些子组件分别通过本文引入的tPD方法（“靶向”）或仅使用目标数据（无非目标批次）的标准PD（“朴素”）获得。x轴和y轴显示在两个目标序列上KL(原始||消融)的最大值。点的颜色代表...