机制化数据归因：追溯可解释大语言模型单元的训练起源

用影响函数追踪 LLM 可解释单元到训练样本的因果起源

前置知识

机制可解释性

机制可解释性旨在将神经网络逆向工程为可理解的功能电路，揭示模型内部实现特定算法行为的机制。传统研究主要采用事后分析范式，描述推理时电路如何运作，但忽视了这些机制在训练语料中的因果起源。近年来，研究者已成功识别出多种可解释单元，如负责上下文复制的归纳头、存储事实关联的知识神经元，以及通过稀疏自编码器解纠缠的单语义特征。

本文的核心贡献是将机制可解释性从静态事后分析扩展到动态因果追溯，因此理解该领域的基本概念是读懂本文的前提。

归纳头

归纳头是 Transformer 中的一类特定注意力头，负责执行上下文内的复制-粘贴匹配功能。当序列中出现 AB...A 模式时，归纳头会提高对 token B 的预测概率。具体而言，它们通过追踪前一个 token 的转换来工作，被认为是上下文学习能力的关键基础组件。Olsson 等人在 2022 年的研究中首次系统描述了这一机制。

本文的主要实验对象就是归纳头，通过 MDA 框架追溯其训练起源并验证其与 ICL 的因果关系。

影响函数

影响函数是一种经典统计工具，用于估计上加权某个训练样本对测试样本损失的影响。标准公式为 $I(z_{train}, z_{test}) = -\nabla_{\theta} L(z_{train})^\top H_{\theta}^{-1} \nabla_{\theta} L(z_{test})$，其中 $H_{\theta}$ 是损失函数的 Hessian 矩阵。对于大型语言模型，精确计算 Hessian 逆矩阵计算成本过高，因此需要使用近似方法如 EK-FAC。

MDA 框架的核心技术基础就是影响函数，本文对其进行了特定单元的改造，使其能够计算训练样本对可解释单元的影响而非全局损失。

EK-FAC 近似

EK-FAC 是一种可扩展的 Hessian 矩阵近似方法，通过逐层使用协方差矩阵的 Kronecker 积来近似 Hessian，从而实现高效的逆 Hessian-向量积估计。该方法由 George 等人在 2018 年提出，Grosse 等人在 2023 年将其应用于大型语言模型的训练数据归因。

EK-FAC 是 MDA 框架能够扩展到大规模语言模型的关键技术，使得在 Pythia 14M-160M 规模上的归因计算成为可能。

前一个 token 头

前一个 token 头是 Transformer 中的另一类重要注意力头，它们将注意力集中在序列中紧邻的前一个 token 上。这类头与归纳头密切相关，因为归纳头的计算路径通常依赖前一个 token 头作为上游组件。在本文的实验中，前一个 token 头与归纳头一起作为主要的可解释单元进行分析。

理解前一个 token 头与归纳头的关系有助于理解本文关于归因范围的讨论，以及训练数据如何通过级联效应影响不同类型的注意力头。

研究动机

当前机制可解释性研究主要停留在静态分析阶段，能够逆向工程出模型内部电路的计算逻辑，但无法追溯这些计算的因果起源。例如，我们已经知道归纳头负责上下文复制、知识神经元存储事实关联，但不知道是哪些训练样本导致了这些特定机制的形成。这种局限性在两个方面造成问题：从科学理解角度，无法用因果视角观察专门化电路如何被训练语料的统计特性塑造；从实践治理角度，缺乏精确的数据级干预手段来控制模型内部表示的形成，传统数据归因方法通常聚焦于全局模型行为，而非特定可解释单元。

本文的目标是本文提出机制化数据归因框架，旨在将训练数据归因的分析粒度从全局模型行为下移到单个可解释单元层面，实现对特定机制形成过程的精确追溯。具体目标包括：建立可扩展的梯度归因方法来识别对可解释单元功能行为影响最大的训练样本；通过因果干预实验验证归因结果的有效性；揭示归纳头形成的数据分布模式和动态特性；建立归纳头与上下文学习能力的因果联系；基于归因发现开发实用的数据增强管道。

与已有工作不同的是，MDA 的独特切入角度在于将机制可解释性与训练数据归因这两个原本独立的研究领域桥接起来。传统训练数据归因方法使用标准验证损失来衡量影响，而 MDA 创新性地将其替换为针对特定可解释单元设计的探测函数 $f_{probe}(\theta, D_{probe})$，同时将计算范围限制在目标单元的参数子空间 $\theta_{sub}$ 内。这种从电路存在到数据导致电路形成的视角转变，使得分析从描述性分析转向发展性追溯，为理解 LLM 的内部机制提供了一个全新的因果分析范式。

核心方法

MDA 框架的整体思路是将标准影响函数的应用场景从全局模型行为细化到单个可解释单元，通过三阶段流程实现训练数据到内部机制的因果追溯。直觉上，如果我们想知道某个特定注意力头是如何形成的，就需要计算训练语料中每个样本对该头功能行为的影响程度。技术路线包括：首先通过监控指标定位目标可解释单元并确定其参数子空间；然后在该子空间内计算每个训练样本的影响分数；最后利用这些分数指导因果干预实验来验证归因的有效性。

MDA 的核心创新在于将影响函数的计算从全局参数空间限制到目标单元的参数子空间 $\theta_{sub} \subseteq \theta$，并将标准损失函数替换为针对特定单元设计的探测函数 $f_{probe}$。形式化地，MDA 影响分数定义为：$$I_{MDA}(z_{train}, D_{probe}) = -\nabla_{\theta_{sub}} L(z_{train})^\top \hat{H}_{\theta_{sub}}^{-1} \nabla_{\theta_{sub}} f_{probe}(\theta, D_{probe})$$ 其中 $\hat{H}_{\theta_{sub}}^{-1}$ 是在参数子空间内使用 EK-FAC 近似的逆 Hessian 矩阵。这与已有方法的本质区别在于：传统 TDA 方法量化训练样本对整个参数空间内全局模型损失的影响，而 MDA 能够归因到细粒度的组件级行为，实现更局部化、更具机制性的模型发展理解。

方法步骤详情

MDA 框架包含三个阶段的完整流程。Stage 1 定位可解释单元，由三元组 $(\mu, \pi, f_{probe})$ 形式化表征，首先定义监控指标用于识别特定可解释单元，在该指标指导下定位目标单元并隔离其关联的参数子空间，然后基于对已识别单元的机制理解设计探测函数和相应的评估数据集来评估目标单元的功能效能。Stage 2 进行单元特定影响计算，用探测函数替换标准影响函数中的验证损失，结合指定的参数子空间计算训练样本对可解释单元的影响。Stage 3 进行理解和干预，利用影响分数进行因果干预，包括数据删除和数据增强实验来验证归因的有效性。

技术新颖性

MDA 的技术新颖性体现在多个层面。首先，这是首个将训练数据归因与机制可解释性系统性结合的框架，填补了知道电路存在但不知道其训练起源的研究空白。其次，MDA 影响函数的推导实现了从全局行为到组件级行为的归因粒度跃迁，通过对参数子空间的限制和探测函数的引入，使得归因结果直接反映特定机制的形成原因而非泛化的模型性能。第三，MDA 展示了亚线性的计算扩展性，GPU 时间与模型参数量的对数关系的 R-squared 达到 0.9640，这使得框架能够应用到当代十亿参数级语言模型。第四，基于 MDA 发现的数据增强管道实现了跨模型规模的泛化，从 14M 模型识别的模式可以有效加速 160M 模型的电路形成。

**Figure 1**: The mechanistic data attribution (MDA) framework

实验结果

本文通过在 Pythia 模型家族上的大量实验验证了 MDA 框架的有效性和发现了多项重要机制洞察。因果验证实验表明，删除高影响样本显著抑制或延迟归纳头和前一个 token 头的出现，而随机排除几乎无影响；增强高影响样本则加速这些头的形成，随机插入基线无此效果。具体数据显示，在 Pythia-14M 上，归纳头分数从基线 0.432 提升到增强后的 0.485，提升 12.3%；在 Pythia-70M 上从 0.304 提升到 0.352，提升 15.8%。影响分数分布呈现幂律特性，前 10% 的样本贡献了高达 50% 的累积影响。高影响样本展现出高度重复的结构模式，主要来自 LaTeX 和 XML 等数据。跨头一致性分析发现，不同归纳头的高影响样本存在显著重叠，且与非归纳头的样本明显分离。动态分析揭示归纳头形成遵循稳定累积模型而非突然相变，高影响样本主要调节形成速率。最重要的是，实验证明了归纳头与上下文学习的双向因果耦合：增强归纳头能力同时提升 ICL 性能，反之亦然。

**Table 1**: Representative high-influence training samples

**Table 2**: Changes in induction head scores across various model sizes after augmenting the training set with Pythia-14M-guided synthetic patterns

**Figure 4**: Distributional properties of high influence samples

**Figure 5**: Validating the functional role of induction heads in ICL via data intervention

**Figure 6**: Impact of MDA augmentation on general capabilities and different augmentation strategies

查看结构化数据

任务	指标	本文	基线	提升
归纳头分数提升 Pythia-14M	Induction Head Score	0.485	0.432	+12.3%
归纳头分数提升 Pythia-31M	Induction Head Score	0.523	0.472	+10.8%
归纳头分数提升 Pythia-70M	Induction Head Score	0.352	0.304	+15.8%
归纳头分数提升 Pythia-160M	Induction Head Score	0.558	0.508	+9.84%
跨模型泛化 14M 到 160M	Induction Head Score	0.521	0.508	+2.56%

局限与改进

本文存在若干局限性需要讨论。首先，归因范围问题：MDA 假设对可解释单元行为的归因严格只考虑训练数据对该特定单元的直接影响，但实际上训练数据可以通过调节单元计算图中的上游节点产生间接影响。例如，当一个前一个 token 头作为归纳头的上游组件时，基于路径的聚合无法解耦计算出的影响是源于对前一个 token 头的直接影响还是通过归纳头传播的间接影响。其次，机制与可解释单元的等同问题：术语机制经常与被识别为其基础的特定可解释单元互换使用，但最近研究表明单一抽象机制可以由多个结构不同的电路实现。第三，实验规模限制：虽然 MDA 在 Pythia 14M-160M 上进行了验证，并对 OLMo-2-1B 和 OLMo-2-7B 进行了定性分析，但由于全规模重训练的巨大计算成本，更大模型的因果验证仍有限。第四，合成数据在大规模时出现多样性耗尽问题，导致性能不如自然高影响样本。

独立分析的弱点

本文存在几个值得关注的弱点。第一，MDA 依赖于 EK-FAC 近似来计算逆 Hessian 矩阵，这种近似可能引入系统性偏差，特别是在参数子空间较小或损失景观高度非凸的情况下，近似精度可能下降。改进方向是探索更精确的二阶近似方法或自适应选择近似策略。第二，数据增强管道的第二步依赖于大型语言模型来提取结构模式，这引入了对额外模型的依赖和潜在的模式提取偏差。可以考虑开发更轻量级、更可控的模式提取方法。第三，Figure 6b 显示合成数据在样本量达到 100000 时出现性能反转，说明当前的程序化合成方法在多样性方面存在瓶颈。改进方向包括增加合成模板的变异度或混合使用多种模式生成策略。第四，虽然实验涵盖了 4 个 Pythia 模型规模，但这些模型共享相似的架构和训练数据，对不同架构的泛化性尚未充分验证。

未来方向

本文作者提出了几个重要的未来研究方向。首先，需要系统性地解耦训练数据的直接影响和通过上游组件传播的级联影响，功能分解方法可能有助于确定归因的范围。其次，需要更深入地探索抽象机制与具体结构实现之间的关系，特别是当单一机制可以由多个电路实现时，什么构成规范因果中介仍是一个开放问题。基于本文成果可延伸的方向包括：将 MDA 应用于更多类型的可解释单元如 SAE 特征和知识神经元；开发更高效的影响分数计算方法以支持更大规模模型的全量归因；利用 MDA 指导机制对齐，即精确地引导或遗忘特定模型行为；探索数据催化剂的特性以优化预训练效率。

复现评估

本文的复现条件相对良好。实验严格遵循 Pythia 仓库的配置，包括超参数、数据排序和随机种子，这有助于确保可重复性。Pythia 模型套件和相关数据集都是公开可用的。然而，MDA 的完整计算需要访问训练过程中的模型检查点和梯度信息，这通常只有在自定义训练设置中才可行。EK-FAC 近似的实现也需要特定的基础设施支持。作者提到附录中提供了详细的实验配置和所有生成阶段的提示，这有助于复现数据增强管道。总体而言，对于拥有足够计算资源的研究者，核心实验是可复现的，但完整的 MDA 计算管道的部署仍需要显著的工程投入。