Title: Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization
Primary area: reinforcement learning
Keywords: fine-tuning, diffusion large language model, policy optimization

Abstract: Diffusion large language models (dLLMs) are promising alternatives to autoregressive large language models (AR-LLMs), as they potentially allow higher inference throughput. Reinforcement learning 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 up to 18\% over previously SOTA baselines and 31\% over the base model, underscoring the effectiveness of the distribution matching framework.

TL;DR: We propose a framework to fine-tune a diffusion LLM through distribution matching policy optimization.