Matryoshka attribution: Learning to attribute language model outputs to representations and weights
Aryaman Arora, Kirill Acharya, Nathan Hu, Yanzhe Zhang, Noah Goodman, Dan Jurafsky, Christopher Potts
cs.CL, cs.LG
2026-09-22
斯坦福把归因做成可学的套娃式掩码排序。MAttr在MIB官方秘密测试集拿下5.6分，亚军1.95；把Llama 3.1 8B Instruct的1%权重还原到base，拒答几乎消失，能力基本还在。
要把语言模型的输出归因到内部计算，现成工具分成三摊。因果干预给每个组件单独换一次激活，前向次数随组件数线性涨，要覆盖非线性组合更是 2^n 次前向。梯度归因一次反向就能扫完全部变量，但经常找不到真正有因果作用的计算。可学习掩码能用梯度下降，通常要加稀疏损失、straight-through estimator 或 Gumbel，调参和训练都不稳。
没有一种方法同时满足因果、可并行、实现简单。斯坦福把问题改写成：给内部变量学一个排序，按这个排序一层层把变量加回未干预集合，让下游损失尽量小。套娃（Matryoshka）指不同稀疏度下保留的集合是嵌套的。
MAttr 学一组标量分数 S。每一步从均匀分布抽一个预算 k，用可微的 sigmoid top-k 把分数变成和为 k 的软掩码：分数高的变量接近 1，保留原计算；其余按 α 在 base 输入和 source 输入的激活之间插值。这就是 soft interchange intervention，一种按权重混合两次前向激活的干预。
关键设计有两条。
因为 k 变大时每个变量的掩码值只升不降，学到的是一个全局排序。只训一步时，对 S 的梯度等价于 centered Input×Gradient；SGD 一步之后，期望上等价于按 p(k) 加权路径的 Integrated Gradients。继续训就能离开 IG 的公理约束：Bilodeau et al. 证明，complete 且 linear 的归因方法可以在局部反事实上不优于随机。MAttr 学的是对常数平移不敏感的排序，并且直接优化任务损失。
表示归因用配对输入（base / counterfactual）和 logit difference。参数归因把插值从激活空间换到权重空间：θ = θb + α ⊙ (θs − θb)，损失改成不 clip 的 GRPO，用奖励模型给生成打分，把优势回传到 S。基是每张权重矩阵里形状为 dmodel 的行，外加 layernorm 向量。
Mechanistic Interpretability Benchmark（MIB）的电路定位赛道上，edge-level MAttr 在官方秘密测试集平均 5.6，第二名 1.95，各模型–任务组合相对亚军提升 45.9% 到 343.0%。公开测试集上对 I×G、IG、Expected Gradients、AttnLRP、Interchange Interventions、DBM、Node Pruning 等全部基线做配对 Wilcoxon，Holm–Bonferroni 校正后 p < 0.05。DBM 和 Node Pruning 若要扫稀疏度，代价大约是 MAttr 的 50 倍，扫完仍追不上。
他们另报 compactness：按对数网格对 interchange intervention accuracy 求面积，奖励「小电路就保住行为」。各粒度上 MAttr 的 CPR 都最高；compactness 多数设置胜过或接近 IG。把 k 改成 log-uniform 采样，compactness 还能再压过 IG，代价是 CPR 略降。
跨任务迁移在 Llama 上很清楚：减法任务学到的排序，在加法和月份算术上仍高度有效，注意力头 a15.h13、a16.h21 被反复排到前面，和已有电路文献对得上。梯度方法里，和 MAttr 分歧最大的是 MLP 块的排序，注意力头相对一致。
参数级实验更刺眼。在 Llama 3.1 8B Instruct 上，把 1% 的权重行还原到 base，StrongREJECT 从 Instruct 的 2.6 升到 84.0，SORRY-Bench 从 27.1 升到 99.3；GSM8K 81.0→79.5，IFEval 74.7→71.7，MMLU 69.1→69.5。Abliteration 要改 27.6% 的参数才做到类似去拒答。同预算的 Expected Gradients 把 IFEval 砸到 47.5。1B 模型上 MAttr 改 2% 参数：StrongREJECT 75.8，MMLU 仍是 43.6。
| 方法 | 改动比例 | StrongREJECT | GSM8K | MMLU |
| Instruct（8B） | 0 | 2.6 | 81.0 | 69.1 |
| Abliteration | 27.6% | 68.6 | 81.5 | 68.6 |
| MAttr | 1.0% | 84.0 | 79.5 | 69.5 |
| Expected Gradients | 1.0% | 59.5 | 76.5 | 67.8 |
附录里，emergent misalignment 的权重增量用 0.1% 几乎收回 held-out 损失；ViT-B/16 上 Faith log-AUC 为 5.2，2% 的 patch 就够让 basenji 变成 argmax。
可解释性写成一个可反传的目标函数之后，找电路不再绑死在手工公理上。从业者如果要做 MIB 式电路定位，MAttr 现在是官方榜第一，一张排序能扫全部稀疏度。如果要审计「base 和 instruct 成对开源」的风险，这篇给出一个定量结论：拒答可以挤在约 1% 的权重差分里，GSM8K、IFEval、MMLU 大部分不在这 1% 上。这既是对齐审计工具，也是去安全微调的说明书。
方法本身是渐进的。价值在把三摊工具收成一个可学目标，以及把 mask learning 接到参数差分和强化学习奖励上。
作者自己说，目前只学数据集上的标量分数 S，没有做成按样本、按 checkpoint 变化的归因元模型。表示归因仍要配对反事实输入；参数归因要奖励模型。CPR 本身有病：logit difference 无界，评价网格在对数稀疏度上却按线性面积积分，随机排序单靠 p=0.5 到 1.0 就能白捡约 0.25 分，所以「5.6 对 1.95」的倍数不能直接当电路质量倍数。
去拒答实验的训练提示来自 StrongREJECT 的 held-out 划分，评测也在同一分布再加 SORRY-Bench。没有证明换到完全不同的有害请求分布仍然 1% 就够。GRPO 无 clip、无 KL，稳定性依赖实现。行级权重在 1B 和 8B 上够稀疏，70B 级有没有同样的 1% 口袋，论文没测。