DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning
This work addresses the efficiency and performance bottlenecks in diffusion-based language models for tasks like math and planning, representing an incremental advancement in optimizing inference-time compute.
The authors tackled the problem of improving the reasoning speed and accuracy of diffusion large language models (dLLMs) by proposing DiFFPO, a reinforcement learning framework that trains surrogate policies and joint samplers, resulting in better task performance with lower computational cost, such as improved Pareto frontiers in inference-time compute.
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.