双向归纳：掩码扩散语言模型中上下文学习的机制分析

# 双向归纳：掩码扩散语言模型中上下文学习的机制分析

来源：https://arxiv.org/html/2607.15893
Andy Catruna Emilian Radoi 布加勒斯特理工大学  \{andy\_eduard\.catruna, emilian\.radoi\}@upb\.ro

###### 摘要

虽然自回归（AR）Transformer 的内部机制已被广泛研究，但对于扩散语言模型（DLM）这种通过迭代去噪生成文本的新兴替代方案，我们知之甚少。在这项工作中，我们研究了 DLM 如何实现归纳——上下文学习背后的一种机制，即模型找到重复的上下文并复制其后跟随的 token。我们的分析比较了仅注意力的 AR 模型和具有匹配架构的吸收掩码 DLM。我们发现 DLM 学习了一个双向归纳电路，其中前 token 头和后 token 头将局部上下文写入残差流，后续的归纳头利用这些信息从匹配的源位置找到并复制答案。该电路是方向对称的，无论源出现在过去还是未来，都能正常工作。当仅能看见左侧上下文（与 AR 模型所见相同）时，DLM 在归纳能力上并未超越其 AR 对应模型。然而，我们观察到，当掩码 token 的两侧上下文都可见时，DLM 表现出更强的归纳能力，这表明其优势源于双向上下文访问，而非更强的单侧机制。除归纳外，我们还提供了因果证据，表明 DLM 计算了掩码 token 的全局比例，并将其用作隐式时间步，即使它们没有明确的 timestep embedding。

# 双向归纳：掩码扩散语言模型中上下文学习的机制分析

Andy Catruna Emilian Radoi
布加勒斯特理工大学
\{andy\_eduard\.catruna, emilian\.radoi\}@upb\.ro

## 1 引言

基于自回归（AR）Transformer 的大型语言模型（LLM）已成为写作、编程和科学工作的核心工具（Noy and Zhang, 2023 (https://arxiv.org/html/2607.15893#bib.bib34); Lee et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib38); Peng et al., 2023 (https://arxiv.org/html/2607.15893#bib.bib35); Boiko et al., 2023 (https://arxiv.org/html/2607.15893#bib.bib36)），这使得其内部机制成为一个重要的研究对象（Bereska and Gavves, 2024 (https://arxiv.org/html/2607.15893#bib.bib37)）。然而，AR Transformer 并非唯一重要的语言模型类别。扩散语言模型（DLM）是一种不断发展的替代方案，它通过反复对序列进行去噪来生成文本，而不是仅预测下一个 token（Austin et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib5); Sahoo et al., 2024 (https://arxiv.org/html/2607.15893#bib.bib14)）。这赋予了它们不同的计算结构以及并行解码的潜力，正如最近诸如 LLaDA（Nie et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib8)）、DiffusionGemma（O'Donoghue and Flennerhag, 2026 (https://arxiv.org/html/2607.15893#bib.bib28)）和 Mercury（Khanna et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib9)）等系统所展示的那样。

![图1：DLM 中识别出的双向归纳电路。相邻头写入前 token 和后 token 线索。后续的归纳头使用其 QK 权重将这些线索与源答案进行匹配，并通过其 OV 路径提高答案 token 在掩码位置上的概率。当源答案位于未来时，同样适用该电路。](https://arxiv.org/html/2607.15893/extracted/6287535/figures/dlm_induction.png)

AR 语言模型在机制可解释性方面受到了广泛关注（Bereska and Gavves, 2024 (https://arxiv.org/html/2607.15893#bib.bib37); Elhage et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib2); Olsson et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib1); Wang et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib3); Conmy et al., 2023 (https://arxiv.org/html/2607.15893#bib.bib26)），但关于 DLM 内部机制我们知之甚少（Wang et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib10); Dai et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib11); Kong et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib12)）。由于 DLM 使用了不同的训练目标、双向注意力以及迭代去噪过程，在 AR 模型中发现的机制可能无法直接迁移（Kong et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib12)）。

在 AR 模型中，归纳头是上下文学习的典范电路（Olsson et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib1)）。它们使模型能够利用重复的上下文来复制之前紧跟在同一上下文之后的 token，为在提示内通过示例进行学习提供了一种基本机制。理解 DLM 是否以相同的方式实现归纳，是比较其与 AR 模型上下文学习机制的第一步。

DLM 改变了归纳必须运作的环境。它们从双向上下文预测掩码 token，但预测位置的真实 token 被掩码 token `[M]`（通常写作 `[MASK]`）隐藏（Austin et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib5); Sahoo et al., 2024 (https://arxiv.org/html/2607.15893#bib.bib14)）。因此，DLM 可以使用相同的 AR 归纳电路，对称地扩展它以从过去和未来的上下文中复制，或者使用不同的机制。

在这项工作中，我们在一个受控环境中研究此问题，重点关注吸收掩码 DLM。我们训练了匹配的仅注意力 AR Transformer 和 DLM，在相同的重复 token 归纳任务上进行评估，并使用因果消融和权重分解分析 DLM 电路。这使我们能够在相同的架构和数据分布下比较行为和机制。

遵循机制可解释性的标准实践（Elhage et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib2); Olsson et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib1)），我们使用小型模型以实现电路级因果分析，为理解大规模 DLM 的机制奠定基础。

我们的分析表明，DLM 学习了归纳并以双向形式实现，如图 1 (https://arxiv.org/html/2607.15893#S1.F1) 所示。归纳突现，需要至少两层，并且从过去或未来上下文中的有效性大致相同。

本工作做出以下贡献：

- •我们识别并提供了因果证据，证明存在一个双路径电路：前 token 和后 token 头写入局部上下文，后续的归纳头利用此信息从过去和未来上下文中复制。
- •我们在匹配条件下比较了 AR 模型和 DLM 中的归纳，并表明仅当 DLM 能够利用双向上下文时，其归纳性能才优于 AR 模型；当限制于 AR 模型使用的左侧上下文设置时，则并非如此。
- •我们展示了 DLM 在没有显式 timestep embedding 的情况下，将全局掩码比例编码为隐式时间步。一个线性探针可以在第一注意力层之后从单个位置的残差流中恢复全局掩码比例，并且修补掩码率方向会因果性地改变预测熵。

## 2 相关工作

Transformer 语言模型的机制可解释性。自回归 Transformer 一直是语言模型机制可解释性的主要焦点。先前的工作引入了分析 Transformer 电路的工具，包括残差流分解、注意力头电路以及 QK/OV 分解（Elhage et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib2)）。这一研究方向将归纳头确定为 AR 模型中上下文学习的核心机制，其中较早的 token 提供了从重复上下文中复制所需的特征（Olsson et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib1)）。

机制可解释性还开发了因果和表示级方法来研究模型行为（Mueller et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib42)）。电路分析和路径修补（Wang et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib3); Goldowsky-Dill et al., 2023 (https://arxiv.org/html/2607.15893#bib.bib4); Conmy et al., 2023 (https://arxiv.org/html/2607.15893#bib.bib26)）测试哪些组件对于某种行为是必要的，而探针和稀疏特征方法（Belinkov, 2022 (https://arxiv.org/html/2607.15893#bib.bib39); Huben et al., 2024 (https://arxiv.org/html/2607.15893#bib.bib25); Templeton et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib24)）则研究残差流中表示了什么信息。关于叠加的工作进一步表明为什么单个神经元可能无法干净地对应单个特征（Elhage et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib23)）。我们的 DLM 分析借鉴了其中几种方法，结合了均值消融、QK/OV 分解和线性探针。

扩散语言模型。DLM 通过逐渐去噪损坏的序列来生成文本，而不是从左到右逐个生成 token。早期工作开发了离散扩散目标，包括吸收态和更一般的转移过程，后来工作将这些目标调整到语言建模（Austin et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib5); Gulrajani and Hashimoto, 2023 (https://arxiv.org/html/2607.15893#bib.bib6); Lou et al., 2023 (https://arxiv.org/html/2607.15893#bib.bib7); Sahoo et al., 2024 (https://arxiv.org/html/2607.15893#bib.bib14); Shi et al., 2024 (https://arxiv.org/html/2607.15893#bib.bib15)）。我们在本文中将 DLM 用作广义术语，我们的实验集中在吸收掩码 DLM 上，先前的工作通常称之为掩码扩散模型（MDM）（Ou et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib16); Zheng et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib17)）。

最近的工作通过规模化掩码扩散模型、将 AR 模型适配为扩散模型以及结合自回归和扩散式生成，使 DLM 更具竞争力（Nie et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib18); Gong et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib19); Arriola et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib21)）。大型 DLM 进一步表明，基于扩散的生成正逐渐成为标准 AR 解码的实用替代方案（Nie et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib8); Ye et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib20); Khanna et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib9); O'Donoghue and Flennerhag, 2026 (https://arxiv.org/html/2607.15893#bib.bib28)）。这些模型通常是为了更快或更并行的解码而提出的，但它们双向去噪的目标（Artetxe et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib40); Patel et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib41)）也改变了上下文学习的内部计算（Kim et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib22)）。

掩码扩散模型的可解释性。关于掩码扩散模型的机制工作仍在涌现。最近的研究分析了 DLM 中的稀疏特征、去噪过程中的注意力浮动行为、当 AR 模型被后训练为掩码扩散模型时的机制变化，以及掩码扩散与学习顺序自回归之间的理论联系（Wang et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib10); Dai et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib11); Kong et al., 2026 (https://arxiv.org/html/2607.15893#bib.bib12); Garg et al., 2025 (https://arxiv.org/html/2607.15893#bib.bib13)）。这些论文表明，DLM 可以共享 AR 模型的一些行为，同时也会改变信息的路由方式。相反，我们的工作识别了 DLM 的归纳电路，并在匹配的 AR-DLM 设置中为其双路径结构提供了因果证据。

## 3 方法

### 3.1 匹配的自回归与扩散模型

我们训练了深度 L∈{1,2,3}L∈{1,2,3} 的匹配仅注意力 AR 和 DLM Transformer。AR 模型使用因果注意力和下一个 token 预测，而 DLM 使用双向注意力和吸收掩码扩散目标（Austin et al., 2021 (https://arxiv.org/html/2607.15893#bib.bib5); Sahoo et al., 2024 (https://arxiv.org/html/2607.15893#bib.bib14)）。除了目标和注意力掩码外，两类模型的架构和训练方案是相同的。对于主要的行为比较，我们为每个模型家族和深度训练三个独立的随机种子，从而允许我们跨模型家族比较归纳的涌现和最终归纳性能。

### 3.2 测量归纳

我们在由随机 token 构建的序列上评估归纳能力，这些序列中只能利用上下文信息来预测正确的 token（Olsson et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib1)）。每个序列长度为 512，前 256 个 token 均匀随机采样并作为后半部分重复。令 `W X A Y R` 表示一个包含答案 token `A` 的五 token 源窗口，令 `W X [M] Y R` 表示相应的查询窗口，其中 `A` 被掩码 token `[M]` 替换。在前向归纳中，源位于第一个副本，查询位于第二个副本；而在反向归纳中，它们的顺序互换。总共，我们使用 128 个这样的序列来评估每个模型。我们掩码一个固定的、确定性的非邻近位置集，前向评估中评分 3,328 个掩码 token，反向评估中评分 3,328 个掩码 token。

我们使用随机 token，以便语义和自然语言频率效应无法解释性能，模型必须利用上下文信息来识别掩码 token。我们通过重复对模型预测正确 token 的帮助程度来评分归纳：对于每个掩码位置，我们将模型在重复序列中正确 token 的对数概率与该 mask 布局相同但非重复的控制序列中相同 token 的对数概率进行比较。归纳分数是此差值（以 nats 为单位），在所有掩码位置上取平均。

AR 模型没有掩码 token `[M]`，因此我们使用对相同序列的下一个 token 预测来评估它们：模型在没有任何掩码的情况下读取序列，并从其左侧的 token 预测每个评分位置，使用相同位置匹配的控制。这使得 AR 分数可以直接与 DLM 前向分数进行比较，而反向设置没有 AR 对应版本。

### 3.3 电路分析

为了定位电路，我们利用两种互补的测量方法。首先，我们计算结构性的注意力指纹，例如一个头从掩码位置 i 注意至 i-1、i+1 或另一个副本中无掩码 token 的源答案位置的程度。其次，我们执行均值消融：对于每个头，我们将其输出替换为在非重复控制提示上的每个位置均值，并重新计算归纳分数（Wang et al., 2022 (https://arxiv.org/html/2607.15893#bib.bib3)）。这种方法帮助我们识别哪些头注意到了相关信息，以及哪些头实际上对归纳具有因果重要性。

为了进行权重级分析，我们使用没有 LayerNorm 的 DLM（通过转换训练时带有 LayerNorm 的模型获得），以便残差流可以分解为可加组件。对于 QK 分析，我们将归纳头从 `[M]` 到源答案的注意力 logit 分解为按查询组件和按键组件的项。遵循 Elhage 等人（2022）的工作，...（此处原文被截断，但翻译保留此形式）