用强化学习让扩散语言模型并行生成更快更准
dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learning
- 通过在线强化学习优化每轮解码的去掩码策略
- 单次前向传播可生成超过5个词,效率远超基线
- 适合需要高速生成的代码和数学推理场景
掩码扩散语言模型(MDLMs)具备并行生成潜力,但现有开源模型每次前向传播生成的词数通常少于5个,严重限制了并行性。现有加速方法依赖固定置信度启发式或基于蒸馏的微调,后者在微调过程中可能产生偏离策略的问题,且性能受限于基础模型生成质量。本文提出dUltra,一种基于组相对策略优化(GRPO)的在线强化学习框架,用于学习高效的去掩码策略。dUltra引入去掩码规划头,预测每个词的独立伯努利去掩码概率。通过联合优化基础扩散语言模型与去掩码顺序规划器,使用包含可验证奖励、蒸馏奖励及去掩码步数的综合奖励信号。在数学推理和代码生成任务中,dUltra相比最先进的启发式方法(Fast-dLLM)和蒸馏基线(d3LLM, dParallel),实现了更优的准确率-效率权衡,证明了在线强化学习所学去掩码轨迹能更好利用MDLM的并行生成能力。代码与模型检查点已公开于https://github.com/chinsengi/dUltra-os。
原文摘要 · Abstract (English)
Masked diffusion language models (MDLMs) offer the potential for parallel token generation, but most open-source MDLMs decode fewer than 5 tokens per model forward pass even with sophisticated sampling strategies, limiting their parallel generation potential. Existing acceleration methods either rely on fixed confidence-based heuristics or use distillation-based approaches that finetune MDLMs on trajectories generated by a base model, which can become off-policy during finetuning and restrict performance to the quality of the base model's samples. We propose \texttt{dUltra}, an on-policy reinforcement learning framework based on Group Relative Policy Optimization (GRPO) that learns unmasking strategies for efficient parallel decoding. dUltra introduces an unmasking planner head that predicts per-token unmasking likelihoods under independent Bernoulli distributions. We jointly optimize the base diffusion LLM and the unmasking order planner using reward signals combining verifiable reward, distillation reward, and the number of unmasking steps. Across mathematical reasoning and code generation tasks, dUltra achieves superior accuracy-efficiency trade-offs compared to state-of-the-art heuristic (Fast-dLLM) and distillation baselines (d3LLM, dParallel), demonstrating that learned unmasking trajectories through on-policy RL enable better exploitation of parallel generation in MDLMs. Code and checkpoints are released at https://github.com/chinsengi/dUltra-os.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。