Random Remasking Scales Masked Diffusion Inference

📄 arXiv: 2610.30744v1 📥 PDF

作者: Zhaoyi Zhou, Frank Wu, Andrea Zanette

分类: cs.LG

发布日期: 2026-10-22

备注: 18 pages, 4 figures, 8 tables

💡 一句话要点

提出随机重掩码方法以提升掩码扩散推理效果

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

关键词: 掩码扩散模型 去噪 随机重掩码 轨迹修正训练 逻辑推理 自然语言处理 模型优化

📋 核心要点

  1. 现有的掩码扩散模型在解码过程中引入的错误通常是不可逆的，导致生成质量下降。
  2. 本文提出随机重掩码-去噪（R²D）推理方法，通过随机重引入掩码来改善去噪过程，分析去噪轨迹。
  3. 在逻辑任务中，提出的方法在测试时间上实现了近对数线性扩展，数独和3-SAT的准确率分别超过97%和99%。

📝 摘要（中文）

掩码扩散模型（MDMs）通过逐步去掩码生成序列，但在解码过程中引入的错误通常是不可逆的。尽管迭代修正已被证明可以提高样本质量，但现有理论未能解释这一现象。本文将MDM推理重新表述为在干净序列空间上的马尔可夫链，定义了一个随机重掩码-去噪内核，随机重新引入掩码并进行去噪。这一分析促使我们提出了随机重掩码-去噪（R²D）推理方法，这是一种完全在测试时、无微调、无奖励的过程。为了进一步提升性能，我们提出了轨迹修正训练（TCT），该过程使模型暴露于自身生成轨迹中的错误。我们的方案在逻辑任务上表现出显著的近对数线性测试时间扩展，在困难的数独谜题上达到了97%以上的准确率，在3-SAT上超过99%的准确率。该方法还适用于预训练模型在实际语言任务中的应用，超越了迭代修正基线和LLaDA2.1-mini在HumanEval+上的官方Q模式准确率。

🔬 方法详解

问题定义：本文旨在解决掩码扩散模型在解码过程中引入的不可逆错误问题，导致生成序列的质量下降。现有方法未能有效处理这些错误，影响了模型的性能。

核心思路：论文提出将MDM推理视为在干净序列空间上的马尔可夫链，通过随机重掩码和去噪的内核设计，改善去噪过程并分析生成轨迹。

技术框架：整体架构包括随机重掩码-去噪内核和轨迹修正训练（TCT）模块。随机重掩码模块在生成过程中随机引入掩码，而去噪模块则负责去除噪声并生成高质量的序列。TCT则通过暴露模型于自身错误，进一步提升去噪效果。

关键创新：最重要的创新在于随机重掩码-去噪推理方法（R²D），与现有方法相比，它通过动态引入掩码来改善去噪过程，提供了一种新的视角来分析去噪轨迹。

关键设计：在模型训练中，采用了特定的损失函数来优化去噪效果，同时在TCT过程中，模型被设计为能够识别并修正自身的错误，确保生成的序列更接近真实数据分布。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，提出的方法在逻辑任务上实现了显著的性能提升，数独谜题的准确率超过97%，而3-SAT的准确率超过99%。此外，该方法在预训练模型上也超越了迭代修正基线和LLaDA2.1-mini在HumanEval+上的官方Q模式准确率，展现了良好的通用性。

🎯 应用场景

该研究的潜在应用领域包括逻辑推理、自然语言处理和其他需要序列生成的任务。通过提升掩码扩散模型的性能，能够为复杂问题的求解提供更高效的工具，未来可能在智能助手、自动编程等领域产生深远影响。

📄 摘要（原文）

Masked diffusion models (MDMs) generate sequences by progressively unmasking tokens, but errors introduced during decoding are typically irreversible. Recent work has shown that iterative correction can improve sample quality, yet existing theory fails to explain this given a learned, imperfect denoiser model. We re-formulate MDM inference as a Markov chain on the clean sequence space, defined by a remask-denoise kernel that randomly re-introduces masks and then denoises. This helps us analyze the actual denoising trajectory towards the data distribution. The analysis motivates Random Remask-Denoise (R$^2$D) inference, a purely test-time, finetuning-free, reward-free procedure. To further improve performance, we propose Trajectory Correction Training (TCT), a denoiser-training procedure that exposes the model to its own errors along the generation trajectory. On logical tasks, our approach exhibits a striking near-log-linear test-time scaling, achieving over 97% accuracy on hard Sudoku puzzles and over 99% accuracy on 3-SAT. Our approach also applies to pretrained models on actual language tasks, beating iterative correction baselines and the official Q-mode accuracy of LLaDA2.1-mini on HumanEval+.