arXiv:2512.09106cs.LG2025-12中稿 · ICML被引 26

用强化学习训练扩散语言模型的解码策略,提升生成效率与质量。

Learning Unmasking Policies for Diffusion Language Models

  • 用强化学习自动学习解码时该解码哪些词,替代人工调参的启发式方法。
  • 在全扩散设置下,新策略生成质量优于现有启发式方法,且块大小越大优势越明显。
  • 适合关注高效生成、推理加速的开发者和研究者使用。

扩散(大)语言模型(dLLMs)在多项下游任务上已达到自回归模型的性能水平,同时具备更高的推理效率潜力。其关键设计之一是每一步扩散过程中选择解码哪些标记的采样过程。近期研究表明,基于置信度阈值等启发式策略相比随机解码,能提升样本质量和令牌吞吐量。但此类策略存在需手动调参、块大小增大时性能下降等问题。本文提出采用强化学习训练采样策略,将掩码扩散采样建模为马尔可夫决策过程,其中dLLM作为环境,并设计一种基于单层Transformer的轻量级策略,将令牌置信度映射为解码决策。实验表明,该训练策略在半自回归(块)生成中表现媲美当前最优启发式方法,在全扩散设置中则表现更优。

原文摘要 · Abstract (English)

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One critical design aspect of dLLMs is the sampling procedure that selects which tokens to unmask at each diffusion step. Indeed, recent work has found that heuristic strategies such as confidence thresholding improve both sample quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger block sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive (block) generation, while outperforming them in the full-diffusion setting.

扩散模型生成策略强化学习

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