Matryoshka Attribution：把语言模型输出归因到表示与权重

如何把语言模型输出归因到真正有因果作用的内部计算，同时兼顾因果性、可并行和简单。

• 场景：作者认为模型内部有大量计算不会以人能读懂的方式说出来。审计对齐、可靠编辑模型，以及理解哪些内部计算导致了输出、各组件起什么功能，都需要这种归因。
• 现有方法的不足：因果干预要对每个变量单独前向，完整因果说明要对 n 个组件的幂集做 2^n 次前向。基于梯度的方法一次反向就能覆盖全部变量，但实践中可能找不到有因果效果的计算，也不一定优化下游可解释性指标。掩码学习可并行且能用梯度下降，但额外的稀疏损失和梯度估计带来大量手工调参和训练不稳定；要扫不同稀疏度还得另训掩码。
• 核心做法：把归因定义成找出嵌套的内部组件子集，使下游损失最小。Matryoshka Attribution（MAttr）用可微 sigmoid top-k 参数化掩码，训练时随机采样预算 k，从而一次学出按归因分数排序的组件顺序，不再使用显式稀疏损失，也不用直通估计器或 Gumbel 技巧去近似硬掩码的梯度。
• 应用：在 Mechanistic Interpretability Benchmark（MIB）的电路定位上评测表示归因；再用强化学习把微调前后的行为变化归因到权重差。作者称恢复 Llama-3.1-8B-Instruct 1% 的权重到基座状态，就足以去掉拒答并维持能力。

论文在 §2 把归因方法分成三类，并在 §4 把其中若干方法当作基线重调。

因果干预

• 因果抽象与交换干预（Geiger et al., 2020; 2021; 2025；同期还有 Vig et al., 2020; Meng et al., 2022; Chan et al., 2022; Goldowsky-Dill et al., 2023）——把反事实表示换进基座表示，用输出变化衡量直接因果效应（Pearl, 2001）。后续用于语言与数学行为、导向和训练。
• 代价：每个组件一次前向；并行需要不精确近似（Nanda, 2023）。组件不必线性组合（McGrath et al., 2023），完整因果说明要对幂集做 2^n 次前向（Sundararajan & Najmi, 2020）。

梯度归因

• 早期显著图与 I×G、Integrated Gradients（Simonyan et al., 2014; Baehrens et al., 2010; Shrikumar et al., 2017; Sundararajan et al., 2017）以及 LRP（Bach et al., 2015）——用梯度或泰勒近似给内部计算分配功劳，并被改写到语言模型上。
• 论文的批评：梯度归因可经反向传播并行，但可能是因果干预的不可靠近似，也不一定优化关心的下游可解释性指标（Bilodeau et al., 2024）。

掩码学习

• Hard concrete 掩码（Louizos et al., 2018）及后续用于解释（De Cao et al., 2020; Cao et al., 2021 等）和剪枝（Voita et al., 2019; Michel et al., 2019 等）——在任务损失上加稀疏损失，用梯度下降掩掉不重要部分。
• 论文的批评：额外稀疏项和梯度估计造成调参与不稳定（Gale et al., 2019; Yin et al., 2019; Bhaskar et al., 2024）；扫稀疏度要另训掩码，而能给出归因排序的方法在评测时可以免费扫稀疏度。

基线与基准

• 基线方法：Interchange Interventions、I×G、Integrated Gradients（IG）、Expected Gradients（EG）、GIM、AttnLRP、RelP、DBM、Node Pruning；参数实验另有 GRPO、GRPO+KL、Abliteration、GRP-Oblit、EG、Base（URIAL）。
• 基准/数据集：MIB 的 IOI、arithmetic、MCQA、ARC-Easy、ARC-Challenge；MIB+ 增加主谓一致与更多算术任务；拒答实验用 StrongREJECT、SORRY-Bench、GSM8K、IFEval、MMLU。

作者把 MAttr 定位成同时追求因果性、可并行和简单的掩码学习：用随机 k 和 sigmoid top-k 去掉显式稀疏损失与硬掩码梯度近似，并在 §3 说明限制训练步数时它与中心化 I×G、中心化 Integrated Gradients 的关系。

MAttr把归因分数 S 学成一个顺序：按这个顺序把变量加回未干预集合，就能在各个稀疏度上压低给定损失。训练时每步随机抽预算 k，用 sigmoid top-k 做成嵌套软掩码，再对非 top-k 变量做软交换干预。

问题设定与软交换干预

• 变量与基：模型 M 是变量 V 上的有向无环图，变量 H 由其父节点经机制 F_H 计算。归因所覆盖的变量集合称为 basis H，例如 MLP 神经元或注意力头。
• 软交换干预：给定基座输入 b、源输入 s 和掩码权重 α_H ∈ [0, 1]，把机制换成 α_H F_H(u) + (1−α_H) h(s)。叶节点（输入）则用 h(b) 替换 F_H(u)。输出记为 M_{H←H*}(b, s, α)。α 接近 1 表示保留基座计算，接近 0 表示换成源运行上的值。
• 优化目标：在期望上最小化各稀疏度的损失。目标写成 min_S E_{k∼p(k)}[L(M_{H←H*}(b, s, σ(S+c_k)), y_b, y_s)]。因为掩码随 k 单调地从“被干预”变到“不干预”，子集是嵌套的，所以称为 Matryoshka。

掩码与训练

• 分数与预算：S 初始化为 0。每步从 Uniform(1, |S|−1) 抽 k，再从数据集抽一对 ⟨b, s, y_b, y_s⟩。
• Sigmoid top-k：α^{(k)} = σ(S+c_k)，其中 c_k 用二分法求解，使各变量的掩码之和等于 k（Wijk et al., 2025）。论文称因此不再需要 L0 的近似稀疏项，也不需要直通估计器或 Gumbel。
• 损失与更新：对非 top-k 变量做软干预，损失 L 相对基座标签 y_b 计算，再用优化器按 ∂ℓ/∂S 更新 S，重复 T 步。MIB 上梯度基线与训练方法都以 logit difference L_MIB 为目标，训练方法最大化它。节点级用 Adam 时主要调学习率；更细的 basis 另设 ε=10^{-2}。§A.3 的理论称小 ε 的 Adam 会忽略梯度幅度，可能损害数据效率。MIB+ 上各方法统一为 5000 个训练步。
• 与梯度法的关系：作者称若限制训练，第一步对 S 的梯度等价于在掩码前向上计算的中心化 I×G；在 SGD 下训练一步的期望，分数等价于按预算分布 p(k) 加权路径的中心化 Integrated Gradients。证明放在 §A。作者还称 Bilodeau et al. (2024) 关于 Integrated Gradients 的不可能性定理不适用于 MAttr：排序对分数的常数平移不敏感，因而不受那些归因公理约束；并且 MAttr 直接优化从业者指定的损失。

评测时的电路与指标

• 硬掩码：评测用硬掩码 β^{(k)} 干预非 top-k 变量。Faith(k) 是归一化的 logit difference，分母是完全不干预与全部干预之差。
• CPR：在比例网格 P={.001, .002, .005, .01, .02, .05, .1, .2, .5, 1} 上对 Faith 用梯形法则求面积，k_i=⌊p_i |H|⌋。
• Compactness：先定义 IIA(k)，即干预后基座标签 logit 高于源标签 logit 的比例；再在 log k 上等间距对 IIA 求面积。作者引入它，是因为 CPR 的 logit difference 无上界，且线性面积会被大电路区间主导。
• Basis：MIB 原有 node（每个注意力头、MLP 块和输入）和 edge（节点间残差项）。本文另加 mlp（每个 MLP 神经元和 token）、mlp+attn（每个 token 上的 MLP 神经元和注意力头）、sae（每个 MLP 输出 SAE 特征 and token）。SAE 误差项当作叶节点，单独有一个归因分数。

参数归因与强化学习

• 权重空间插值：给定两个检查点 θ_b、θ_s，软干预为 θ* = θ_b + α ⊙ (θ_s − θ_b)。S 与 k 仍按算法 1 生成 α^{(k)}。评测时用硬掩码把权重差的 top-k 部分 graft 回去。
• GRPO：奖励函数 r_ϕ 给完成打分，不裁剪更新。每步按当前 k 准备干预后的权重，对批内每个提示采样一组完成，用组内归一化优势和长度归一化对数概率做策略梯度，梯度经 α^{(k)} 回到 S。只保留组内奖励标准差不低于 ε 的组。
• 拒答设置：basis 是每个权重矩阵中形状为 d_model 的行（例如对应单个 MLP 神经元的 gate、up、down 投影向量）以及 layernorm 向量。作者称选择这些行是因为它们通常是特权基。训练用 StrongREJECT 提示和模型默认对话模板，奖励来自基于 gemma-2b 微调的 StrongReject 打分器，输出是 [1, 5] 上的有用性概率分布。在 StrongREJECT 的小子集上评测，其余提示用于训练。

实验设置

• 电路模型：MIB 有 4 个基座模型：gpt2-small、Qwen2.5-0.5B、gemma-2-2b、Llama-3.1-8B，共 12 个任务–模型对。任务为 IOI（4 个模型）、arithmetic（仅 Llama）、MCQA（除 GPT-2 外）、ARC-Easy（Llama 与 Gemma）、ARC-Challenge（仅 Llama）。
• MIB+：三个新 basis 按 token 归因，只用 Llama；另有 Marks et al. (2025a) 的 4 个主谓一致任务和 Feucht et al. (2026) 的 4 个算术任务。MLP 输出 SAE 用 Llama Scope。论文未在正文给出各任务样本量。
• 电路基线：重调并重评 I×G、IG，并加入 Interchange Interventions、EG、GIM、AttnLRP、DBM、Node Pruning。梯度法沿用 MIB 原预算；MIB+ 统一 5000 步。DBM 与 Node Pruning 在多个稀疏度上分别训练，扫稀疏度的成本作者称约为 MAttr 的 50 倍。超参数见 §B.1，正文未列出具体学习率表。
• 电路指标：Faithfulness、CPR、Compactness、IIA。显著性用双侧配对 Wilcoxon signed-rank，并对每个 MIB 子任务和每个基线做 Holm–Bonferroni（§E.3）。评测设置沿用 MIB 代码库。官方榜使用秘密测试集，正文只报告公开测试集分数；秘密集只给了平均分。
• 拒答模型与基线：Llama-3.2-1B-Instruct 与 Llama-3.1-8B-Instruct，对照同一模型的基座检查点。基线为 GRPO（无裁剪、无 KL）、GRPO+KL（β=0.01）、Abliteration、GRP-Oblit、EG，以及带 URIAL 模板的基座。1B 的 GRP-Oblit 全参数训练；GRPO、GRPO+KL 和 8B 的 GRP-Oblit 用 LoRA，r=32，α=64，学习率 10^{-5}。
• 拒答指标：StrongREJECT、留出的 SORRY-Bench，以及 GSM8K、IFEval、MMLU；同时报告被改参数比例 ∥∆θ∥0。Figure 5a 给出均值 ± 误差，蓝色表示与 Instruct 无显著差异。论文未说明重复次数。

主结果

Figure 2 标注的是跨模型/任务平均 CPR 与成本，不是逐模型数字。正文另给出官方秘密测试集的总分。

表格较宽，可左右滑动

设置	指标	MAttr	对照	出处
官方秘密测试集，edge-level，跨模型与基准平均	榜分数	5.6	次名 1.95	§4 正文
相对次名，各模型–基准组合	提升	45.9%–343.0%	次名	§4 正文
公开测试集，node 与 edge	CPR	高于所测基线	p<0.05	Figure 2，§4
Llama-3.1-8B-Instruct，恢复 1% 权重	StrongREJECT / GSM8K	84.0±2.0 / 79.5±2.9	Instruct 2.6±1.5 / 81.0±2.8	Figure 5a
Llama-3.2-1B-Instruct，恢复 2% 权重	StrongREJECT / GSM8K	75.8±3.2 / 28.5±3.2	Instruct 4.3±2.1 / 29.5±3.2	Figure 5a

• 作者对 MIB 的解读：作者称 MAttr 在公开测试集的节点级和边级电路定位上统计上优于所有被测基线，并在各评测指标和电路粒度上表现稳定。官方秘密测试集平均分 5.6，次名 1.95；各模型–基准组合相对次名的百分比提升为 45.9% 到 343.0%。Figure 2 顶部还标出若干方法的平均 CPR，例如 edge 侧可见 6.44、2.49、1.56、1.60、1.61，node 侧可见 2.30、2.06、1.85 等；图注未把每个柱子逐一对应到方法名以外的图例位置，正文未逐项抄出全部柱值。
• 成本：Figure 2 底部是每个任务所需反向次数（Interchange Interventions 为前向次数）。图中可见多组梯度法为 0.1–1k，部分掩码法为 3k 或 0.5k，另有 63–423k 的区间。作者称 MAttr 成本低于基线。
• 拒答：作者称两种模型上，常规去拒答方法改动的参数都超过 27.6%，而 MAttr 把 1B 的 2%、8B 的 1% 恢复到基座状态后，有害性评测与这些方法相当，并维持能力。EG 使用相同的参数比例，但作者称它未能维持非拒答能力。

Figure 5a 中与能力、有害性直接相关的其他数字：8B 上 MAttr 的 SORRY-Bench 为 99.3±0.4（Instruct 27.1±2.1），IFEval 71.7±1.9（Instruct 74.7±1.9），MMLU 69.5±2.0（Instruct 69.1±2.0）；EG 在 1.0% 参数下 StrongREJECT 59.5±3.6、GSM8K 76.5±3.0、IFEval 47.5±2.1。1B 上 MAttr 的 SORRY-Bench 为 99.6±0.3，IFEval 46.2±2.1，MMLU 43.6±2.2；Abliteration 改动 38.9% 参数，StrongREJECT 70.2±3.0。Base（URIAL）在 8B 上 GSM8K 为 4.0±1.4、StrongREJECT 52.7±4.2。

CPR 与 Compactness

• 测了什么：Figure 3 在 node、MLP、MLP+Attn、SAE（MLP out）上画 MIB+ 任务平均的 CPR 与 Compactness，比较 MAttr、MAttr+log k、IG、I×G、NP、DBM、Random。
• 结果：作者称在所有 basis 上 MAttr 的 CPR 最高；Compactness 上它优于 IG 或紧随其后，取决于 basis。若改用对 k 的 log-uniform 采样，作者称可以用 CPR 换取高于 IG 的 Compactness。图中纵轴 CPR 标到 2.0 与 6，横轴 Compactness 标到约 0.6 与 0.5；正文未给出每个点的精确坐标。
• 作者的解读：固定稀疏度的 DBM 与 Node Pruning 会过度优化 CPR，Compactness 甚至低于 I×G。增大稀疏系数能提高 Compactness 但损害 CPR（§E.1）。多稀疏度汇总对两个指标都有帮助，但仍不足以追上 MAttr。作者称这两种方法在不同稀疏度上找到的变量集合几乎是嵌套的（§E.2）。

排序一致性与跨任务迁移

• MLP 分歧：Figure 4a 是节点级 MIB 上，各方法组件排序与 MAttr 的 Spearman ρ，并在任务上平均。作者称 MAttr 与表现差的梯度法在 MLP 块上分歧明显，在注意力头上则不明显；IntInv、EG、AttnLRP 和掩码方法与 MAttr 聚在一起。图中可读到的 ρ 包括注意力头上 IntInv 与 I×G 为 .86、MAttr 与 IntInv 为 .46，MLP 上 MAttr 与 IntInv 为 .48、与 I×G 为 −.19。
• 迁移：用源任务的 MAttr 排序在目标任务上评测，覆盖节点级 Llama 的全部任务对。Figure 4b 的单元格是用源任务分数在目标任务上恢复的 IIA AUC 比例。作者称减法、加法与 months 高度共享头 a15.h13 和 a16.h21。Figure 4c 中跨任务归因分数的 Pearson r 与 Compactness 恢复比例相关，图上标出 r=0.90，并标了 ARC-E–ARC-C、Arith.–Add.、Arith.–Simple 等配对。
• 参数扫描：Figure 5b 扫 ∥∆θ∥0，画 GSM8K 对 StrongREJECT，圈出表格中的 MAttr 点，并与 EG 比较，覆盖 1B 与 8B。作者称很大一段参数比例都能维持 GSM8K，直到出现性能下跌。权重位置和生成示例放在 §I；正文未给出下跌处的精确比例。

其他消融与分析

• log-uniform 的 k：相对均匀采样，作者称 Compactness 可超过 IG，CPR 下降（Figure 3）。论文未在正文给出该变体的精确 CPR。
• Adam 的 ε：更细 basis 设 ε=10^{-2}；节点级主要调学习率，作者称较大范围内都能成功。未报告具体学习率网格的分数。
• DBM / Node Pruning 的稀疏系数：增大系数提高 Compactness、降低 CPR（§E.1）；多稀疏度汇总仍低于 MAttr。
• 一步梯度：作者称 MAttr 第一步梯度等于中心化 I×G，SGD 一步期望等于路径加权的中心化 IG（§A）。正文未给该等价的数值误差。
• EG 对照：8B 上同样改 1.0% 参数，StrongREJECT 59.5±3.6，低于 MAttr 的 84.0±2.0，IFEval 从 Instruct 的 74.7±1.9 降到 47.5±2.1（Figure 5a）。
• Base（URIAL）：不加微调权重时 8B 的 GSM8K 为 4.0±1.4、IFEval 为 27.5±1.9，低于 MAttr 恢复 1% 后的 79.5±2.9 与 71.7±1.9。

作者指出的局限与后续方向

• 标量分数：本文只在数据集上学习标量归因分数 S。作者认为更值得做的是训练特征器 f，把表示或提示特征变成示例级、检查点级或任务级分数，即归因元模型；该想法类似 Bushnaq et al. (2025) 的因果重要性函数（§6 Future work）。
• 联合训练：可以把模型或特征器与标量归因分数一起训练，类似 optimal ablation，用 MAttr 代替 L0 惩罚（§6）。
• SAE 与 DAS：作者写到也可以把 MAttr 用于 SAE 或 DAS 训练（§6）。
• 讨论中的定位：作者认为可解释性对深度学习的采用偏慢，主张把可解释性写成可用梯度下降优化的目标，而不是只在单一设定上用专门方法做主张（§6）。这是立场，不是列出的实验局限。
• 结论中的范围：作者称通过多种基准和案例说明 MAttr 是最先进的归因方法，并称附录还有 emergent misalignment、玩具模型中的 interference weights 和 Vision Transformer 上的实验（§7）。这些是已做或放入附录的工作，不是局限清单。论文正文的 Limitations 小节未单独列出失败条件。

实验覆盖范围

• 电路定位：公开测试集覆盖 gpt2-small、Qwen2.5-0.5B、gemma-2-2b、Llama-3.1-8B 共 12 个 MIB 任务–模型对，basis 包括 node、edge，以及仅在 Llama 上的 mlp、mlp+attn、sae（§4）。
• MIB+：主谓一致 4 个任务、Feucht et al. (2026) 的算术 4 个任务，以及 Llama Scope 的 MLP 输出 SAE；跨任务迁移图用节点级 Llama 的任务对（§4，Figure 4）。
• 指标与检验：主指标为 CPR，另报告 Compactness 与 IIA；节点级和边级相对每个被测基线做了配对 Wilcoxon 与 Holm–Bonferroni（§4，§E.3）。官方秘密测试集只报告平均分 5.6 对 1.95，以及相对次名 45.9%–343.0% 的区间。
• 参数归因：Llama-3.2-1B-Instruct 与 Llama-3.1-8B-Instruct 相对各自基座，basis 为 d_model 行和 layernorm；训练奖励来自 StrongREJECT 的 gemma-2b 打分器，评测含 StrongREJECT 小子集、SORRY-Bench、GSM8K、IFEval、MMLU（§5.1，Figure 5）。
• 训练预算：MIB 梯度法沿用原预算，MIB+ 为 5000 步；拒答侧 GRPO 类方法的 LoRA 设置为 r=32、α=64、学习率 10^{-5}，GRPO+KL 的 β=0.01。论文未在正文报告 MAttr 的查询次数或与人工评判的一致性。

Matryoshka Attribution（MAttr）把“哪些内部计算或权重差导致了输出”写成可以用梯度下降学习的嵌套子集问题。

• 问题：因果干预要按组件甚至按幂集做前向，梯度归因可能对不上因果重要的计算，掩码学习则要调稀疏损失并在每个稀疏度上另训。作者希望一个方法同时具有因果性、可并行性和简单性。
• 方法：为每个变量学习分数 S。每步从均匀分布抽预算 k，用二分法确定偏移，使 sigmoid 掩码之和等于 k，再对非 top-k 变量做软交换干预，使输出靠近基座行为。掩码随 k 单调嵌套，因此一次训练得到可在评测时任意取稀疏度的排序。限制训练步数时，作者称它与中心化 I×G 和中心化 Integrated Gradients 有精确对应。参数版本在两个检查点之间插值，并用 GRPO 优化不可微奖励。
• 电路结果：作者称 edge-level MAttr 在 MIB 官方秘密测试集上平均分为 5.6，次名为 1.95，各模型–基准组合相对次名提升 45.9%–343.0%。公开测试集上，节点级与边级 CPR 高于所测基线，双侧配对 Wilcoxon 给出 p<0.05。作者称各 basis 上 CPR 最高，Compactness 则优于或接近 IG；log-uniform 的 k 可以换得更高 Compactness。跨任务时，归因分数的 Pearson r 与 Compactness 恢复比例相关，Figure 4c 标出 r=0.90。
• 拒答结果：把 Llama-3.1-8B-Instruct 相对基座的 1% 权重差按 MAttr 分数恢复后，StrongREJECT 为 84.0±2.0、SORRY-Bench 为 99.3±0.4，GSM8K 为 79.5±2.9（Instruct 为 81.0±2.8），MMLU 为 69.5±2.0。1B 上对应比例是 2%，StrongREJECT 为 75.8±3.2，GSM8K 为 28.5±3.2（Instruct 为 29.5±3.2）。同样只改 1% 的 EG 在 8B 上 IFEval 降到 47.5±2.1。作者称这是拒答特有、且因果上局部的权重差。
• 作者的结论：作者认为 MAttr 把可解释性收成一个可学习目标，并鼓励沿梯度下降做可解释性。后续方向是示例级归因元模型、与模型联合训练，以及用于 SAE 或 DAS。