掩码生成到「再掩码」训练：RemeDi让扩散语言模型学会纠正与反思

近期，扩散语言模型备受瞩目，提供了一种不同于自回归模型的文本生成解决方案。为使模型能够在生成过程中持续修正与优化中间结果，西湖大学 MAPLE 实验室齐国君教授团队成功训练了具有「再掩码」能力的扩散语言模型（Remasking-enabledDiffusion Language Model,RemeDi9B）。在扩散去噪的多步过程中，通过进行再掩码 SFT 和 RL 训练，为每个 token 输出一个去掩码置信度，RemeDi 能够从序列中已经生成的内容中识别无法确定的位置进行再掩码（remask），从而修正错误内容并提升文本质量，在各方面都超越了现有的扩散语言模型。该模型还具有可变长生成（variable-length generation）能力，打破了现有中大规模扩散语言模型仅支持定长生成的限制，提高了模式能力的灵活性。

• 论文地址：https://arxiv.org/abs/2509.23653
• 代码与模型地址：https://github.com/maple-research-lab/RemeDi

背景

扩散语言模型已成为自回归语言模型的有力替代方案。这一类方法首先定义了一个将文本逐步破坏为噪声的前向过程，然后让模型学习从噪声中恢复出干净文本的逆向过程。在这一类方法中，当前最主流的是基于掩码的扩散语言模型。该方案要求模型在训练中学习恢复被掩码的 token，而已经被恢复的 token 则在之后的生成步骤中保持不变，直到生成结束。这其中蕴含了一则假设：每一步中预测的 token 都必然是正确的，无需修正，直接可以当作最后的生成内容。这一假设显然过于理想 —— 生成过程中，模型不可避免地会产生预测错误，而我们应当赋予模型通过自我反思发现并修正这些错误的能力。

为解决这一问题，提出一种面向扩散语言模型的自我反思式生成范式 —— 再掩码（remask），并基于这一范式训练了具有「再掩码」能力的扩散语言模型 RemeDi。如图所示，RemeDi 具备发现错误 token，并通过再掩码将其修正的能力：模型首先生成了 “left”，但随后在生成完整句子的语义表示时，发现 “left for the pies” 这一表述与实际含义不符，因此，将 “left” 一词再掩码，修改为更合适的 “used”。可以看出，通过再掩码，模型能利用在后续步骤中生成的上下文信息，识别较早步骤中存在的错误，将其改正，并基于更丰富的上下文信息进行更精确的预测。

用置信度识别「再掩码」目标

为了让 RemeDi 能够通过再掩码修改已经生成的文本内容，一个核心的挑战是让模型能够找到需要修改的 token，执行再掩码操作。为此，我们对网络结构进行了修改，让其在预测序列中每个 token 输出分布的同时，能够为每个 token 额外预测一个置信度分数。整个模型采用了一种双流协同的模型结构：

此外，在语言生成任务中，许多场景下的输出并非固定长度。如果模型只能在固定长度下生成，将导致资源浪费或生成结果被压缩、截断。因此，使扩散语言模型具备灵活的不定长生成能力（variable-length generation）是必要的。在 RemeDi 中，我们采用分块自回归生成的方法实现这一点：模型每次会通过一个完整的反向扩散过程生成一段长为 L=32 的序列。完成后，如果该序列中没有生成结束符，则将已生成的这一段序列拼接在上下文中，继续往后生成下一段长为 L=32 的序列，如此重复直到生成结束符为止。与自回归模型类似，我们采用分块因果注意力掩码机制，确保在生成时，每个 token 能看到自己所在的 block 内的其他 token，和之前已生成 block 内的 token，而无法看到未来将要生成的 block。

在实验中，我们基于 LLaDA 的权重继续训练，将其改造成一个具有不定长生成能力的分块扩散模型。上面表 4 中的 baseline 模型即展示了不定长生成模型在经过再掩码训练前的性能。

两阶段训练，赋予「再掩码」能力

1.Remask SFT（监督微调阶段）

传统的掩码扩散语言模型通常通过在输入序列上随机掩码进行有监督微调（SFT）。与之不同的是，RemeDi 在反向扩散过程中还需要能够找到潜在的不正确 token 并再掩码。我们在 SFT 过程中将这类不正确 token 视为除掩码 token 之后的第二类噪声。因此，在 SFT 阶段，我们不仅要训练模型从掩码 token 恢复原文本的能力，同时也需要训练识别那些需要再掩码的不正确 token。

由于在反向扩散过程中，噪声水平（定义为 mask token 的数量）应当单调递减。由于在 SFT 设计中，长度为 L 的输入序列中，所有不正确 token 都必须被重新掩码，因此需要满足以下不等式约束：

以确保输出中掩码位置的数量单调减少。若该不等式不成立，则在下一步重新掩码所有不正确 token 会增加总的掩码数量，从而违反扩散过程中掩码比例应逐步减少的基本原则。

整个再掩码微调算法流程如下图：

2.Remask RL（强化学习阶段）

在完成 Remask SFT 训练后，我们进一步通过基于结果的强化学习对模型进行微调。根据实验室先前的研究，反向扩散过程中的每一步中间结果都可以视为大模型的一个「思考」步骤，而基于结果的强化学习可以优化整个生成轨迹，提升模型生成正确最终答案的概率。这种面向扩散语言模型的大模型推理范式称为扩散式「发散思维链」，在机器之心的往期报道中已有详细阐述。（与Gemini Diffusion共振！首个扩散式「发散思维链」来了）

实验结果

在同规模与相近计算预算下，RemeDi 在数学推理、代码生成与通用问答三类任务上均取得稳定提升。其中，仅采用 Remask SFT 带来显著增益；在此基础上加入 Remask RL，多数基准再获得进一步提升。

我们在不同类型的任务上对再掩码次数进行了统计，可以看出：对输出约束更强的任务（如代码生成）会更频繁触发再掩码。

而具体的生成示例也表明，通过再掩码机制，RemeDi 可以实现纠错、插入、删除等多种文本修改手段。

总结

这篇文章介绍了由西湖大学 MAPLE 实验室推出的，具有再掩码反思机制的扩散语言模型，RemeDi。基于额外的置信度预测，RemeDi 能够识别生成过程中的错误，并通过「再掩码」机制重新预测，从而做到生成过程中的自我反思与优化。针对「再掩码」机制设计的有监督训练与强化学习算法确保了这一机制的有效性。实验结果表明 RemeDi 在数学推理、代码生成、通用知识问答等多个任务上都取得了超越其他扩散语言模型的性能。这些结果说明「再掩码」能有效提升扩散语言模型的文本生成质量，值得进一步探讨。