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

2026-07-15

AI Summary

This work addresses the challenge that existing fine-tuning methods for discrete diffusion models struggle to leverage non-differentiable and intermediate reward signals for reinforcement learning optimization. It introduces continuous-time reinforcement learning to this domain for the first time, modeling state evolution via controlled continuous-time Markov chains and establishing a policy optimization framework compatible with arbitrary action spaces and intermediate rewards. Key contributions include an analytically tractable probability ratio parameterization tailored for mask-based diffusion models, a trajectory subsampling technique to reduce the computational cost of large language model fine-tuning, and continuous-time variants of PPO and GRPO. Experiments demonstrate significant improvements in reward-driven fine-tuning performance on low-dimensional entropy-regularized problems as well as complex tasks such as mathematical reasoning and code generation.

Abstract

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.

Problem

• continuous-time reinforcement learning
• discrete diffusion models
• reward-driven optimization
• masked diffusion models
• trajectory likelihood estimation

Innovation

• continuous-time reinforcement learning
• discrete diffusion models
• controlled CTMC
• intermediate reward incorporation
• trajectory subsampling