扩散语言模型中用于Token编辑的自生成错误训练

# 扩散语言模型中用于词元编辑的自我生成误差训练
来源：https://arxiv.org/html/2606.17175
林 尧¹,² ¹上海交通大学计算机科学与技术系，上海，200240，中国 ²中关村创新发展研究院，北京，100097，中国 lin\.yao@sjtu\.edu\.cn

###### 摘要

Token\-to\-token（T2T）编辑允许 LLaDA2\.1 在分块扩散解码过程中修正已生成的词元。官方发布的训练方案使用随机词汇损坏来训练这一编辑器，但在推理时，编辑器看到的是模型自身流畅、高置信度的草稿误差。我们研究了这种训练与推理之间的不匹配问题，并提出了**自我生成 T2T**（Self\-Generated T2T），该方法执行一次无梯度的草稿生成过程，用预测词元填充被遮罩的位置，并在第二次基于这些自我生成损坏的监督过程中进行恢复。我们以 LLaDA2\.1\-mini 上的短持续预训练 LoRA 更新来实现该方法，并在官方 Q\-Mode T2T 流程、推理参数不变的情况下，在多个基准上进行评估。该方法在提升准确率的同时降低了 T2T 编辑强度，缓解了某些失败模式，例如在推理正确后出现最终数字转录错误，以及在简短事实性答案前过度自我修正等问题。

## 1 引言

扩散语言模型通过对部分观测序列进行去噪来生成文本，而非严格遵循从左到右的轨迹。这使得它们具有两个吸引人的特性：能够并行更新多个位置，以及能够重新审视之前的承诺。LLaDA2\.1 通过 token\-to\-token（T2T）编辑实现了第二个特性 [4]。每个块首先通过标准的 mask\-to\-token（M2T）填充进行解码，随后进入编辑阶段，在此阶段中，可见词元会被重新评分，并在新候选词的置信度超过编辑阈值时被覆盖。这条编辑路径是 LLaDA2\.1 质量与速度权衡的核心：它允许模型提前接受词元，并在之后进行修复。这也是必要的，因为并行去噪可能在同一不完整上下文下提交多个词元：同一一步预测出的词元无法相互条件作用，因为相应的位置在预测时仍然被遮罩，因此局部看似合理的承诺在组合到一起后可能变得不一致。

这个编辑器背后的训练方案与部署时的实际情况不太一致。训练期间，T2T 流通过选择干净词元、用随机词汇词元替换它们并应用交叉熵损失来恢复原始词元而构建。推理期间，需要编辑的词元并不是从词汇表中随机采样的。它们是模型自身早期的预测结果：推导中看似合理的数字、局部语法正确但错误的实体、过早的答案词元，或者语义相近的替代项。这些错误比随机损坏更难处理，因为它们作为语言而言属于分布内，并能够主动引导上下文。

本文探讨 T2T 编辑器是否应该在其实际会遇到的误差上进行训练。我们提出了**自我生成 T2T**（Self\-Generated T2T），这是一种针对 LLaDA2\.1 T2T 编辑的训练侧对齐方法。我们不直接从均匀词汇分布中生成可见的损坏，而是先让当前模型填充被遮罩的位置，然后将这些预测出的词元作为可见输入用于第二次监督过程：错误的预测被训练为编辑回干净词元，而正确的预测则被训练为保持不变。训练目标仍然是监督去噪目标；推理过程保持原始的 LLaDA2\.1 T2T 流程和官方 Q\-Mode 参数不变。唯一的变化是 T2T 损坏的来源。

![参见标题](图1：T2T 编辑的训练-推理对齐。(a) 随机词元训练用无关词汇损坏可见词元（例如，*capital*→\to*banana*），产生流形外误差。(b) 自我生成 T2T 使用自我生成的词元草稿填充被遮罩位置，并在来自同一分布的上下文相关误差上进行训练。(c) 推理阶段编辑相同的流形内草稿；自我生成 T2T 的持续预训练（CPT）使训练与部署对齐，而不改变解码过程。)
我们在官方 Q\-Mode T2T 推理协议（第 3.4 节）下，仅改变检查点权重，保持推理过程和参数不变，在 CMATH、TriviaQA、PIQA 和 AIME 2025 上评估自我生成 T2T。表 2 报告了这四个基准上的准确率和生成轨迹。在 CMATH、TriviaQA 和 PIQA 上，我们的方法在降低 T2T 编辑强度的同时提升了准确率；在 AIME 2025 上，准确率保持不变（9/30），但每个问题的编辑次数从 130.2 降至 86.0。CMATH 的提升最大（准确率提升 +5.47 个百分点），尽管平均输出长度有所增加，这与更好的最终数字承诺而非过早的简短回答一致。这是一个在便利的 FineWeb 子集上进行的最小 LoRA CPT 研究（第 3.3 节），旨在作为机制压力测试，而非完全优化的方案。

我们的贡献如下：

- • 我们揭示了 LLaDA2\.1 T2T 编辑中的一个训练-推理不匹配问题：随机词元训练损坏与结构化的模型生成推理误差不匹配。
- • 我们提出了自我生成 T2T（Self\-Generated T2T），这是一种两轮持续训练目标，在自我生成的草稿词元上训练原生 T2T 编辑器，无需添加检测器、修正头、强化学习阶段或推理时机制。
- • 我们提供了一个在 LLaDA2\.1\-mini 上可复现的 LoRA 实现，包括在便利的 FineWeb\-Edu CPT 设置下的损坏参数、损失权重、优化器设置，以及使用官方 Q\-Mode T2T 推理默认值的评估。
- • 我们在相同的官方推理流程和参数下，报告了四个基准上的结果：CMATH（准确率提升 +5.47 个百分点）、TriviaQA 和 PIQA 在编辑强度降低的同时表现提升，而 AIME 2025 与基线得分持平，但所需的 T2T 编辑次数更少。

## 2 相关工作

#### 离散扩散语言模型与并行承诺。

离散扩散模型将去噪扩散扩展到分类状态空间，吸收-遮罩损坏为掩码语言建模提供了自然的桥梁 [2]。针对语言的变体将这一观点细化为实际的文本生成目标和架构，包括 SEDD、MDLM、MD4 和 RADD [13, 17, 19, 15]。最近的系统以不同但相关的方式扩展了这一家族。有些直接训练大规模掩码扩散 LM，有些从自回归检查点初始化并作为扩散 LM 继续训练，有些通过分块扩散改变解码架构以获得更高的吞吐量；代表性例子包括 LLaDA、Dream、DiffuLLaMA、Mercury、BD3\-LM 和 LLaDA2\.1 [14, 23, 7, 9, 1, 4]。这些模型的核心优势在于并行去噪：在获得严格从左到右的上下文之前就可以预测多个位置，最近的分析研究了由此产生的速度、生成顺序与质量之间的权衡 [11]。同样的优势也造成了本文研究的修正问题。在同一并行步骤中预测出的词元共享同一部分观测到的上下文，但它们无法条件作用于此步骤中正在预测的其他词元，因为这些位置在预测时仍然被遮罩。因此，模型可能会承诺那些在共享上下文下单独看来合理但组合后相互不一致的词元；它们可能在语义上不一致，重复不相容的角色，或形成局部合理但全局无效的片段。一旦被承诺，这些独立预测的词元就会成为后续去噪步骤的可见上下文，因此局部冲突可能传播到后续预测中。因此，高质量的并行去噪需要某种机制来修正早期的承诺。

#### 编辑与重新遮罩机制。

现有的修正机制主要区别在于它们对可疑的可见词元采取的行动。LLaDA2\.1 引入了 token\-to\-token（T2T）编辑：在标准的 mask\-to\-token（M2T）去噪之后，当新候选词的置信度足够高时，模型可以覆盖一个已经可见的词元 [4]。而重新遮罩方法则通过将不确定的可见词元重置为 [MASK]，让 M2T 流重新预测它们来撤销承诺；例子包括推理时的重新遮罩采样器、自我反思重新遮罩以及 token\-to\-mask 精炼 [20, 8, 22]。学习型的自我修正方法增加了另一个维度：ProSeCo 插入显式的修正精炼步骤，而 BackPlay 训练一个轻量级的修正器来处理冻结生成器产生的错误 [18, 12]。这些机制解决了并行承诺的相同后果，但它们在不同的点进行干预：改变推理动作、添加修正模块或添加修正阶段。它们与我们的工作互补，原则上可以与我们的工作结合：一个在真实可编辑词元分布上训练的模型在推理时仍然可以使用替换、重新遮罩或额外的修正模块。为了隔离训练分布的影响，我们在这篇论文中特意使用了最简单的修正动作。我们保持原生 LLaDA2\.1 的词元替换动作和推理算法固定，仅改变原生编辑器的训练分布。

#### T2T 损坏不匹配。

LLaDA2\.1 中的 T2T 编辑器通过从随机替换的可见词元中恢复干净词元来进行训练 [4]。这是一个便利的监督目标，但它与推理时遇到的可编辑词元不匹配。在推理时，应该被编辑的词元本身就是由早期模型预测选择的。由于编辑器是在随机替换上训练的，它对于哪些可见词元值得编辑以及应该编辑成什么内容的估计，在面对流畅的模型生成误差时可能校准不良。这类误差与均匀采样的词汇词元有本质区别：它们是上下文相关的、语义合理的，并且能够引导周围上下文。最近的重新遮罩工作注意到了同样的不匹配，并通过将可疑词元移回 [MASK] 来避免替换 [22]。我们的回应与替换或重新遮罩的选择是正交的。无论系统最终是编辑词元还是重新遮罩它，训练分布都应该让模型接触真实的模型生成误差；在这项工作中，我们保留原生 LLaDA2\.1 的替换动作，并使其 T2T 训练分布更接近部署情况。

#### 在模型诱导状态上训练。

一种减少训练-测试差异的常见方法是在模型自身诱导的状态上进行训练，但相关的状态取决于生成范式。在自回归序列建模中，调度采试用模型采样代替部分前缀中的黄金词元，使下一个词元预测器暴露在模型生成的前缀而非仅教师强制的前缀中 [3]。在模仿学习中，DAgger 执行当前策略，在策略实际访问的状态上查询专家，并将这些策略诱导的状态聚合到训练集中 [16]。在扩散模型中，自条件化将模型自身的中间去噪预测作为后续去噪的额外输入，从而使训练更好地匹配采样时使用的迭代结构 [6]。这些方法共享一个原则，但状态空间不同：自回归模型诱导前缀，模仿策略诱导顺序决策状态，而扩散去噪器诱导干净数据的中间估计。对于离散扩散 LM 中的 T2T 编辑，诱导的状态又有所不同：一个部分去噪的词元块，包含早期去噪步骤产生的可见草稿词元，其中一些看似合理但却是错误的。LLaDA2\.1 在其编辑器上训练随机词元替换，而推理则要求编辑器修复模型生成的可见词元错误。**自我生成 T2T** 通过使用模型自身来创建用于 T2T 监督的损坏可见词元，从而为这一设定实现了自我生成的 T2T 训练。

## 3 方法

### 3\.1 基础 LLaDA2\.1 训练与推理

#### 基础 LLaDA2\.1 训练。

训练时，设 x = (x₁, ..., xₙ) 为干净序列，其中 i ∈ {1, ..., n} 表示词元位置索引。损坏的训练输入 z 包含三种位置。M 中的位置是输入位置被填充为 [MASK] 的遮罩输入位置。W 中的位置是可见但错误的词元位置。G 中的位置是可见的黄金词元位置。模型读取完整的损坏输入 z，包括所有遮罩、错误词元和黄金词元位置，并为每个位置生成一个词元分布：

(p₁, ..., pₙ) = f_θ(z)。 (1)

对于所有监督位置，目标是干净词元 x_i。因此，三种位置类型对应的损失为：

L_m2t = Σ_{i∈M} CE(p_i, x_i), L_edit = Σ_{i∈W} CE(p_i, x_i), L_clean = Σ_{i∈G} CE(p_i, x_i)。 (2)

从 x 开始，LLaDA2\.1 训练通过首先采样一个遮罩比例并将相应位置替换为 [MASK] 来构建 z。在剩余的可见位置中，选择一部分子集进行词元级损坏，并用不同于干净词元的随机词汇词元替换。未损坏的可见位置保留其黄金词元。M2T 在遮罩输入位置上训练：给定完整的 z 作为上下文，模型预测 M 中位置的干净词元，对应于 L_m2t。T2T 在可见位置上训练：W 中的错误可见词元应编辑回干净词元，而 G 中已正确的可见词元应保持不变，对应于 L_edit + L_clean。

#### LLaDA2\.1 T2T 推理。

推理时，状态 z^t 同样包含遮罩输入和可见输入位置，但可见词元是由早期模型去噪步骤产生的，而非随机替换产生。M2T 根据选定的遮罩输入位置填充