A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models

Fine-tuning discrete diffusion models (dDMs) for complex tasks like mathematical reasoning and coding is challenging because standard gradient-based reinforcement learning (RL) methods struggle with the non-differentiable nature of discrete sampling and the intractability of trajectory likelihoods in any-order decoding. This paper asks: can we formulate a principled RL framework in continuous time for discrete state spaces to optimize these models effectively?

The authors propose a Continuous-Time Reinforcement Learning (CTRL) framework. By modeling the state dynamics as a controlled continuous-time Markov chain (CTMC), they derive the corresponding Hamilton-Jacobi-Bellman (HJB) equations and policy gradient methods, including continuous-time variants of Proximal Policy Optimization (PPO) and Group Relative Policy Optimization (GRPO). They treat the denoiser probability vector of the diffusion model as the control variable (policy). This allows them to define a clear analytical expression for policy log-likelihoods and incorporate intermediate reward signals, which are often ignored by terminal-reward-only methods.

The framework provides a mathematically rigorous way to fine-tune score-based discrete diffusion models. By treating the denoiser as a policy, the authors achieve flexible parameterization over the vocabulary simplex. To address the computational burden of calculating likelihoods in large language models (dLLMs), they introduce trajectory subsampling techniques. Experimental results on reasoning (e.g., Sudoku, GSM8K) and coding tasks show that this CTRL-based approach outperforms existing state-of-the-art methods like d1, d2, and SPG, demonstrating both better generation performance and training stability.

This work bridges the gap between control theory and discrete generative modeling. By moving the RL formulation into continuous time, the authors provide a unified perspective that is robust to time discretization and flexible enough to handle various reward structures. This offers a scalable path for improving the reasoning capabilities of non-autoregressive language models, which are increasingly important for parallel and flexible token decoding.