Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

📄 arXiv: 2602.11590v1 📥 PDF

作者: Yair Schiff, Omer Belhasin, Roy Uziel, Guanghan Wang, Marianne Arriola, Gilad Turok, Michael Elad, Volodymyr Kuleshov

分类: cs.LG

发布日期: 2026-02-12

💡 一句话要点

提出ProSeCo，通过自校正机制提升Masked扩散模型的生成质量与效率。

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: Masked扩散模型 自校正 序列生成 迭代优化 文本生成

📋 核心要点

  1. Masked扩散模型存在误差累积问题，已生成的token无法被修正，影响最终生成质量。
  2. ProSeCo通过训练模型进行unmask和校正，迭代地细化整个序列，包括已生成的tokens。
  3. 实验表明，ProSeCo在多个任务上实现了更好的质量-效率权衡，并能通过推理时计算扩展进一步提升样本质量。

📝 摘要（中文）

Masked扩散模型(MDM)作为自回归模型的替代方案，在实现并行token生成的同时，也具备了竞争性的性能。然而，MDM面临一个根本限制：一旦tokens被unmask，它们就会保持固定，导致误差累积并最终降低样本质量。本文提出了一种训练模型执行unmask和校正的框架。通过复用MDM去噪网络的输出作为校正器训练的输入，训练模型以从潜在错误中恢复。在生成过程中，我们在unmask步骤之间应用额外的校正细化步骤，以改变解码的tokens并改进输出。我们将训练和采样方法命名为渐进式自校正(ProSeCo)，因为它具有迭代细化整个序列（包括已生成tokens）的独特能力。我们在多个条件和无条件任务上进行了广泛的实验验证，表明ProSeCo产生了更好的质量-效率权衡（采样速度提高约2-3倍），并支持推理时计算扩展，以进一步提高样本质量（在基准测试中提升高达约1.3倍）。

🔬 方法详解

问题定义：Masked扩散模型（MDM）虽然能够并行生成token，但其生成过程存在误差累积的问题。一旦某个token被unmask，后续的生成过程将基于这个固定的token进行，如果该token存在错误，则会影响后续token的生成，最终导致整体样本质量下降。现有方法缺乏对已生成token的纠错机制，无法有效解决这个问题。

核心思路：ProSeCo的核心思路是引入一个自校正机制，允许模型在生成过程中对已生成的token进行迭代修正。通过训练一个校正器，利用MDM去噪网络的输出作为输入，学习如何从潜在的错误中恢复。在生成过程中，交替执行unmask和校正步骤，从而实现对整个序列的渐进式优化。

技术框架：ProSeCo的整体框架包含两个主要部分：MDM去噪网络和校正器。MDM去噪网络负责执行标准的masked扩散过程，即根据masked的输入预测未masked的token。校正器则以MDM去噪网络的输出作为输入，学习如何修正这些token中的错误。在训练阶段，校正器的目标是最小化修正后的token与ground truth之间的差异。在推理阶段，交替执行MDM去噪和校正步骤，直到生成完整的序列。

关键创新：ProSeCo最重要的创新点在于引入了自校正机制，允许模型在生成过程中对已生成的token进行迭代修正。这与传统的MDM方法不同，传统MDM方法一旦unmask某个token，该token就保持固定不变。ProSeCo通过校正器实现了对整个序列的渐进式优化，从而提高了生成质量。

关键设计：校正器的网络结构可以根据具体任务进行选择，例如可以使用Transformer网络。校正器的损失函数通常采用交叉熵损失或均方误差损失，用于衡量修正后的token与ground truth之间的差异。在推理阶段，需要仔细调整unmask和校正步骤的比例，以实现最佳的质量-效率权衡。论文中提到，复用MDM去噪网络的输出作为校正器训练的输入是一个关键设计，使得校正器能够有效地学习如何从MDM的输出中恢复错误。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ProSeCo在多个条件和无条件任务上都取得了显著的性能提升。例如，在文本生成任务中，ProSeCo的采样速度比标准MDM快2-3倍，并且在基准测试中，样本质量提升高达约1.3倍。这些结果表明，ProSeCo在提高生成质量和效率方面具有显著优势。

🎯 应用场景

ProSeCo具有广泛的应用前景，可以应用于文本生成、图像生成、音频生成等领域。例如，在机器翻译中，ProSeCo可以用于提高翻译的准确性和流畅性。在图像生成中，ProSeCo可以用于生成更高质量、更逼真的图像。此外，ProSeCo还可以应用于数据修复、序列预测等任务，具有重要的实际价值和未来影响。

📄 摘要（原文）

Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models, enabling parallel token generation while achieving competitive performance. Despite these advantages, MDMs face a fundamental limitation: once tokens are unmasked, they remain fixed, leading to error accumulation and ultimately degrading sample quality. We address this by proposing a framework that trains a model to perform both unmasking and correction. By reusing outputs from the MDM denoising network as inputs for corrector training, we train a model to recover from potential mistakes. During generation we apply additional corrective refinement steps between unmasking ones in order to change decoded tokens and improve outputs. We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence, including already generated tokens. We conduct extensive experimental validation across multiple conditional and unconditional tasks, demonstrating that ProSeCo yields better quality-efficiency trade-offs (up to ~2-3x faster sampling) and enables inference-time compute scaling to further increase sample quality beyond standard MDMs (up to ~1.3x improvement on benchmarks).