arXiv:2602.11590cs.LG2026-02被引 13

让扩散模型学会自我纠错,生成更高质量的文本。

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

  • 训练模型在解码时同时完成去掩码和修正错误。
  • 采样速度最快提升4倍,质量最高提高1.2倍。
  • 适合需要高精度文本生成的场景,如内容创作。

掩码扩散模型(MDMs)作为自回归模型的替代方案,实现了并行标记生成并达到良好性能。然而,其根本缺陷在于一旦标记被解码便无法修改,导致错误累积并降低样本质量。为此,我们提出一种框架,让模型同时学习解码与纠错。通过复用MDM去噪网络的输出作为纠错器的输入,训练模型从潜在错误中恢复。生成时,在解码步骤之间加入额外的修正步骤,可调整已生成的标记以提升输出质量。我们称该训练与采样方法为渐进式自我修正(ProSeCo),因其具备对整个序列进行迭代优化的能力。我们在多个条件与无条件任务上进行了广泛实验,结果表明, extit{method} 在质量-效率权衡上表现更优(采样速度最高快4倍),并支持推理时计算资源扩展,进一步提升样本质量(基准测试中最高提升1.2倍)。

原文摘要 · Abstract (English)

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 \method~yields better quality-efficiency trade-offs (up to ~4x faster sampling) and enables inference-time compute scaling to further increase sample quality beyond standard MDMs (up to ~1.2x improvement on benchmarks).

扩散模型文本生成自我修正高效采样

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。