Search papers, labs, and topics across Lattice.
This paper introduces a continuous-time reinforcement learning framework that fine-tunes discrete diffusion models using a stochastic control approach, allowing for policy optimization with both terminal and intermediate rewards. The authors derive continuous-time variants of existing policy gradient methods, specifically proximal policy optimization (PPO) and group relative policy optimization (GRPO), which enhance the flexibility of reward signals during the denoising process. Key results demonstrate the framework's effectiveness in optimizing masked diffusion large language models (dLLMs) for tasks such as mathematical reasoning and coding, while also reducing computational costs through trajectory subsampling techniques.
Reward-driven optimization in continuous-time RL can significantly enhance the fine-tuning of discrete diffusion models, even with non-differentiable reward signals.
We formulate reinforcement learning (RL) in continuous time with discrete state spaces and possibly arbitrary action spaces via a stochastic control approach, where the state dynamics are modeled as a controlled continuous-time Markov chain (CTMC). We consider policy optimization problems and derive the corresponding policy gradient methods, leading to continuous-time variants of proximal policy optimization (PPO) and group relative policy optimization (GRPO). As a primary application, we develop a complete continuous-time RL framework for fine-tuning score-based discrete diffusion models. The proposed framework enables reward-driven optimization without requiring differentiability on the reward signals. In contrast to the existing GRPO-based approaches that only rely on terminal rewards, our formulation allows intermediate reward or advantage signals to be incorporated throughout the denoising trajectory. Importantly, when specialized to masked diffusion models (MDMs), our framework encompasses a rich class of policy parameterizations over the vocabulary simplex with analytically tractable probability ratios, providing a unified perspective on exploration and policy optimization in MDMs. For masked diffusion large language models (dLLMs), we further propose trajectory subsampling techniques to efficiently estimate computationally prohibitive trajectory likelihoods, reducing the computational cost of computing per-position probability ratios. We showcase the effectiveness of our methods on both low-dimensional entropy-regularized optimization problems and RL post-training of dLLMs on mathematical reasoning and coding tasks.
