解决扩散语言模型训练与推理不一致的问题,提升生成效果。
MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models
- 将去噪过程建模为序列决策问题,用强化学习优化训练策略。
- 仅用60分之一的梯度更新达到当前最优性能,数学题准确率提升9.6%。
- 提出无需训练的重掩码方法,可配合使用进一步提升效果。
扩散语言模型作为传统自回归模型的替代方案,能实现更快生成和更强的双向上下文条件控制。然而,其在训练与推理阶段存在关键差异:推理时模型逐步揭示生成序列结构,逐次解码更少的掩码项,而训练时掩码随机分布,忽略这一结构。这种差异导致性能下降,但以往工作未予重视。为此,本文将有效去噪轨迹的学习建模为序列决策问题,引入强化学习框架,提出新型掩码扩散策略优化(MDPO),利用扩散模型的马尔可夫特性,在与推理一致的渐进细化流程下训练模型。实验表明,MDPO仅需前人SOTA方法60倍少的梯度更新即可达到同等性能,并在MATH500上平均提升9.6%,在Countdown任务上提升54.2%。此外,提出无需训练的运行置信度重掩码(RCR)策略,作为推理阶段的即插即用改进,增强灵活性并带来额外增益。研究揭示了弥合扩散语言模型预训练与推理间差距的巨大潜力。
原文摘要 · Abstract (English)
Diffusion language models, as a promising alternative to traditional autoregressive (AR) models, enable faster generation and richer conditioning on bidirectional context. However, they suffer from a key discrepancy between training and inference: during inference, MDLMs progressively reveal the structure of the generated sequence by producing fewer and fewer masked tokens, whereas this structure is ignored in training as tokens are masked at random. Although this discrepancy between training and inference can lead to suboptimal performance, it has been largely overlooked by previous works, leaving closing this gap between the two stages an open problem. To address this, we frame the problem of learning effective denoising trajectories as a sequential decision-making problem and use the resulting framework to apply reinforcement learning. We propose a novel Masked Diffusion Policy Optimization (MDPO) to exploit the Markov property diffusion possesses and explicitly train the model under the same progressive refining schedule used at inference. MDPO matches the performance of the previous state-of-the-art (SOTA) method with 60x fewer gradient updates, while achieving average improvements of 9.6% on MATH500 and 54.2% on Countdown over SOTA when trained within the same number of weight updates. Additionally, we improve the remasking strategy of MDLMs as a plug-in inference replacement to overcome the limitation that the model cannot refine tokens flexibly. This training-free method, termed Running Confidence Remasking (RCR), consistently enhances performance and provides further improvements when used with MDPO. Our findings establish great potential for investigating the discrepancy between pre-training and inference of MDLMs. Code: https://github.com/autonomousvision/mdpo. Project Page: https://cli212.github.io/MDPO/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。