Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization
This work addresses the problem of improving reasoning performance for diffusion LLMs, which is incremental as it develops a specialized RL method for an emerging model type.
The paper tackles the challenge of enhancing reasoning capabilities in diffusion large language models (dLLMs) by proposing Distribution Matching Policy Optimization (DMPO), a reinforcement learning fine-tuning method, resulting in accuracy improvements of up to 42.9% over previous state-of-the-art baselines and 55.8% over the base model on reasoning benchmarks.
Diffusion large language models (dLLMs) are promising alternatives to autoregressive large language models (AR-LLMs), as they potentially allow higher inference throughput. Reinforcement learning (RL) is a crucial component for dLLMs to achieve comparable performance with AR-LLMs on important tasks, such as reasoning. However, RL algorithms that are well-suited for dLLMs' unique characteristics have yet to be developed. This paper proposes Distribution Matching Policy Optimization (DMPO), a principled and theoretically grounded RL fine-tuning method specifically designed to enhance the reasoning capabilities of dLLMs by matching the dLLM policy distribution to the optimal, reward-tilted one through cross-entropy optimization. We identify a key challenge in the implementation with a small training batch size and propose several effective solutions through a novel weight baseline subtraction technique. DMPO exhibits superior performance on multiple reasoning benchmarks without supervised fine-tuning, with an accuracy improvement of up to $42.9\%$ over previously SOTA baselines and $55.8\%$ over the base model, underscoring the effectiveness of the distribution matching framework. Our code is available at https://github.com/yuchen-zhu-zyc/DMPO.