Dual Path Attribution: Efficient Attribution for SwiGLU-Transformers through Layer-Wise Target Propagation

1. 研究背景与问题 (Problem)

核心挑战：
随着基于 Transformer 的大语言模型（LLMs）的广泛应用，理解其内部机制对于可靠部署至关重要。然而，现有的**归因（Attribution）**方法（即识别哪些输入或内部组件对模型预测负责）面临以下主要瓶颈：

• 计算成本高昂： 现有的忠实度（Faithfulness）较高的方法（如激活修补 Activation Patching、积分梯度 Integrated Gradients）通常需要大量的前向传播或反事实推理。对于密集组件归因（Dense Component Attribution），即需要分析模型中成千上万个注意力头和神经元时，计算复杂度随组件数量线性增长（$O(M)$），导致在长序列和大规模模型上不可行。
• 近似误差与噪声： 基于梯度的方法容易受到梯度饱和或破碎的影响，而基于注意力的方法往往无法准确反映因果机制。
• 忽略关键组件： 许多方法忽略了模型内部复杂的非线性交互，特别是 SwiGLU（Swish-Gated Linear Unit）架构中的门控机制。

目标：
开发一种既能保持高忠实度（准确识别因果机制），又具有极高计算效率（不依赖反事实样本，计算复杂度低）的归因框架，专门针对当前主流的 SwiGLU 架构 Transformer。

---

2. 方法论 (Methodology)

作者提出了 双重路径归因（Dual Path Attribution, DPA） 框架。其核心思想是将归因视角从“将每个组件前向传播”逆转为“将目标向后传播”。

2.1 核心原理：目标传播 (Target Propagation)

DPA 不针对每个组件单独运行前向传播，而是：

1. 前向传播（一次）： 模型处理输入，缓存所有中间激活值（Activations）。
2. 反向传播（一次）： 将目标 Token 的未嵌入向量（Unembedding Vector）作为目标，递归地通过 Transformer 模块向后传播，计算每个残差位置（Residual Position）的有效目标（Effective Target）。
3. 归因计算： 通过计算原始输入/组件激活与传播回来的有效目标的点积，直接得到归因分数。

2.2 双重路径分解 (Dual Path Decomposition)

针对 SwiGLU 架构，DPA 利用其双线性结构，将计算图解析为两条 distinct 的路径，并分别进行线性化逆推：

1. 内容路径 (Content Pathway)：

   • 负责传播表示信息。
   • 在 GLU 中：对应 Up 投影和 Down 投影的组合（$W_U \cdot W_D$）。
   • 在 MHSA 中：对应 Value 投影。
   • 该路径假设激活值固定，是线性的。

2. 控制路径 (Control Pathway)：

   • 负责路由和门控决策。
   • 在 GLU 中：对应 Gate 投影（通过 SiLU 激活函数）。DPA 通过局部线性化（利用激活值与预激活值的比率）来近似逆推。
   • 在 MHSA 中：对应 Query 和 Key 投影（通过 Softmax 注意力机制）。DPA 采用基于泰勒展开的分解来稳定 Softmax 的逆推过程，追踪注意力分数的敏感性。

2.3 算法流程

• 输入： 输入序列 $X$，目标 Token 的未嵌入向量 $t^{(L)}$。
• 步骤 1： 执行一次前向传播，缓存注意力权重 $\alpha$、归一化输入 $\tilde{x}$、GLU 预激活 $s$ 和门控激活 $\alpha$。
• 步骤 2： 从输出层 $L$ 到输入层 $0$ 进行递归反向传播：
  • 在每一层，分别计算 GLU 和 MHSA 的逆线性变换。
  • 引入权重参数 $\mu$ 来平衡内容路径和控制路径的贡献（例如，$\mu_{up} + \mu_{gate} = 1$）。
  • 更新当前层的有效目标 $t^{(l)}$。
• 输出： 得到输入 Token 的归因分数 $S(x^{(0)}) = (x^{(0)})^\top t^{(0)}$，以及各组件（注意力头、神经元）的归因分数。

---

3. 主要贡献 (Key Contributions)

1. 提出 DPA 框架： 一种新颖的基于分解的归因框架，仅需一次前向传播和一次反向传播即可计算所有组件的归因分数，无需反事实样本。
2. 理论创新： 首次展示了 SwiGLU Transformer 的计算图可以被分解为独立的控制路径和内容路径，并推导了相应的线性化逆传播公式。
3. 效率突破： 实现了相对于模型组件数量 $M$ 的 $O(1)$ 时间复杂度。这意味着无论模型有多少个注意力头或神经元，归因计算的时间成本是恒定的，极大地扩展了密集归因的可行性。
4. 实证验证： 在多个基准测试中证明了 DPA 在忠实度和效率上均优于现有最先进（SOTA）的方法。

---

4. 实验结果 (Results)

实验在 Llama-3.1-8B, Llama-2-7B, Qwen3-4B, Qwen2.5-32B, Mistral-7B 等多个模型上，针对 Known 1000（事实知识）、SQuAD v2.0（阅读理解）、IMDb（情感分析）和 IOI（间接对象识别）等任务进行评估。

4.1 归因忠实度 (Attribution Faithfulness)

• 指标： 使用“全面性（Comprehensiveness/Disruption）”和“充分性（Sufficiency/Recovery）”指标。
• 表现：
  • Token 级归因： DPA 在所有数据集上均取得了最高的总忠实度分数（Total Score）。例如在 IMDb 情感分析任务中，DPA 的总分为 30.17，远超次优的 Integrated Gradients (11.12)。
  • 组件级归因： 在识别因果相关的注意力头和 MLP 神经元方面，DPA 表现出极高的精度。在 IOI 任务中，DPA 实现了零干扰（0.00 disruption）和极高的恢复率（86.87+），能够精准定位到负责特定电路的组件。

4.2 计算效率 (Efficiency)

• 时间复杂度： 理论分析表明 DPA 为 $O(1)$，而 Activation Patching (AP) 和 DePass 为 $O(M)$。
• 实测数据：
  • 在 Llama-3.1-8B 模型上进行长序列（SQuAD）归因时，DPA 仅需 129 秒。
  • 相比之下，Activation Patching 和 DePass 分别需要 5370 秒 和 5452 秒。
  • 加速比： DPA 比现有方法快约 40 倍。
• 内存开销： DPA 的峰值内存占用与基线方法相当，主要开销在于前向传播时的激活缓存，并未因反向传播而显著增加。

4.3 敏感性分析

• 研究发现，对于 Token 级归因，平衡内容路径和控制路径（$p \approx 0.5$）效果最好。
• 对于细粒度的组件级归因，增加控制路径（门控和查询/键）的权重能显著提高识别精度，表明控制路径在决定组件功能中起关键作用。

---

5. 意义与局限性 (Significance & Limitations)

意义

• 填补空白： DPA 成功解决了“高忠实度”与“高效率”之间的权衡难题，使得在大规模模型上进行密集组件归因（Dense Component Attribution）成为可能。
• 可解释性工具： 为研究人员提供了一种强大的工具，能够以前所未有的速度和精度揭示 LLM 内部的因果推理路径（Causal Mechanisms），例如识别出特定的注意力头或神经元在事实回忆或逻辑推理中的作用。
• 架构通用性潜力： 虽然目前针对 SwiGLU 优化，但其“自顶向下目标传播”的核心思想具有架构无关性，未来可推广至其他架构。

局限性

1. 架构特异性： 当前的数学分解专门针对标准的稠密 SwiGLU Transformer。对于混合专家（MoE）或状态空间模型（如 Mamba），需要手动推导新的逆函数。
2. 近似误差： 为了线性化，DPA 丢弃了二阶交叉导数。在高度非线性或组件间强耦合的区域（高曲率区域），归因的忠实度可能会下降。
3. Logit 级范围： DPA 直接追踪目标 Logit，而非 Softmax 后的概率分布。因此，它可能难以捕捉那些通过抑制竞争 Token 的 Logit 来间接提升目标概率的机制。
4. 内存需求： 尽管时间效率高，但为了进行反向传播，前向传播时需要缓存所有中间激活值，这对显存（VRAM）仍有较高要求，特别是在超长上下文窗口下。

总结

Dual Path Attribution (DPA) 是一项重要的进展，它通过创新的“双重路径”分解和“目标反向传播”机制，将 Transformer 归因的计算复杂度从线性降低到常数级。这不仅大幅提升了计算效率，还保持了甚至超越了现有方法的忠实度，为大规模语言模型的机械可解释性研究提供了新的范式。