用强化学习训练扩散语言模型，又快又准。
DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning
- 通过近似策略与重要性采样，提升强化学习采样效率。
- 联合训练采样器使推理次数减少，准确率更高。
- 适合追求高效推理的模型优化研究者。
我们提出DiFFPO，一种统一框架，用于训练掩码扩散大语言模型（dLLMs），使其在推理时不仅更准确（furious），而且更快。首先，我们统一现有基线方法，提出通过离策略强化学习训练代理策略，其似然性更易计算，作为真实dLLM策略的近似。这自然引出一种更精确、信息量更高的两阶段似然近似结合重要性采样校正，从而获得样本效率更高、任务性能更优的广义强化学习算法。其次，我们提出一种新方向：联合训练dLLMs的高效采样器/控制器。通过强化学习，激励模型自适应地为每个提示分配推理阈值。联合训练采样器后，相比仅训练模型，显著降低函数评估次数（NFEs），同时提升准确率，在推理时计算开销的帕累托前沿上达到最优表现。我们在基准数学和规划任务上，对开源扩散语言模型进行了验证，展示了该流水线的有效性。
原文摘要 · Abstract (English)
We propose DiFFPO, Diffusion Fast and Furious Policy Optimization, a unified framework for training masked diffusion large language models (dLLMs) to reason not only better (furious), but also faster via reinforcement learning (RL). We first unify the existing baseline approach such as d1 by proposing to train surrogate policies via off-policy RL, whose likelihood is much more tractable as an approximation to the true dLLM policy. This naturally motivates a more accurate and informative two-stage likelihood approximation combined with importance sampling correction, which leads to generalized RL algorithms with better sample efficiency and superior task performance. Second, we propose a new direction of joint training efficient samplers/controllers of dLLMs policy. Via RL, we incentivize dLLMs' natural multi-token prediction capabilities by letting the model learn to adaptively allocate an inference threshold for each prompt. By jointly training the sampler, we yield better accuracies with lower number of function evaluations (NFEs) compared to training the model only, obtaining the best performance in improving the Pareto frontier of the inference-time compute of dLLMs. We showcase the effectiveness of our pipeline by training open source large diffusion language models over benchmark math and planning tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。