1 Introduction↩︎

Score-based diffusion models, which was originally designed for image generation in continuous space [1], [2], have recently emerged as a powerful paradigm for discrete data generation in the context of LLMs [3]–[6]. For discrete diffusion models with the forward process formulated by a continuous-time Markov chain (CTMC), one learns the discrete score function [6]–[8], i.e., the probability ratio of the forward marginals, to generate the reverse process. A notable example is masked diffusion models [9], [10] that corrupt data into fully masked sequences and train a denoiser to predict the clean data from the masked tokens. As was shown by [9], the score function reduces to the denoiser probability vector on the vocabulary space for masked positions, making the model formulation, training, and inference much simpler by directly predicting the clean data at each denoising step (\(x_0\)-prediction). This simplified structure gains dominant popularity, and further enhances the rapid development of mask-based diffusion large language models [11]–[17]. DLLMs mark a promising generative modeling paradigm for text generation because they enable more flexible parallel and any-order token decoding, in contrast to the traditional left-to-right sequence decoding in autoregressive large language models (ARMs).

Fine-tuning discrete diffusion models via reinforcement learning (RL) is challenging because sampling from the discrete categorical distribution is non-differentiable, so standard gradient methods are not applicable. Specific to language models, the ARM fine-tuning process is naturally constructed as a Markov Decision Process (MDP) thanks to the left-to-right decoding, in which the action is to predict the next token with the state being the current decoded sequence, leading to a factorized policy likelihood by the chain rule. However, it is difficult to scale and enhance the reasoning capabilities of dLLMs by online RL that has been well-developed in ARM post-training, because of the intractable likelihood of the generated sequences in any-order decoding. Recent methods, e.g., d1 [18], UniGRPO [19], wd1 [20], SPG [21], and d2 [22], focused on estimating the trajectory log-likelihood, either by heuristic estimators lacking a clear probabilistic interpretation or by tractable evidence lower and upper bounds. While these approaches have achieved promising empirical performance, they often suffer from coarse likelihood approximations and limited control over the estimation errors. More fundamentally, existing GRPO methods are not grounded in a unified MDP formulation, leaving the connection between iterative diffusion denoising for rollout generation and RL policy optimization largely heuristic and inconsistent. Moreover, most GRPO-based algorithms rely on reward signals only at the final completion, resulting in coarse credit assignment for the intermediate token-unmask decisions that shape the generation process. Although a recent work [23] introduced branching-based resampling at selected intermediate steps to provide intermediate feedback, the rewards are still being derived from the terminal evaluation, and are potentially computation-expensive due to branching.

The purpose of this work is to tackle the aforementioned issues by developing an RL framework for CTMC models, which enables us to propose a principled approach for fine-tuning any score-based discrete diffusion model via continuous-time reinforcement learning (CTRL). To see the motivation, [24], [25] and [26] first proposed to fine-tune diffusion models in continuous spaces via CTRL. In that setting, the score function is the only unknown component in the generative process, so one can treat it as a control/action/policy, and hence, the reverse process is a state-controlled stochastic differential equation (SDE). Then a score-related stochastic policy is designed for exploration, and the corresponding policy gradient and estimation can be obtained by the well-developed stochastic control and RL theory in continuous time and space [27]–[33]. Since the state dynamics of diffusion models are inherently formulated as continuous-time processes (either SDEs or CTMCs), fine-tuning them via CRTL is a natural choice, which not only exhibits robustness to time discretization [25], but also allows us to carry out all clean theoretical derivations in continuous time, while retain great flexibility for discretizing the continuous-time processes at the final algorithmic stage.

However, RL in continuous time and discrete spaces has yet been developed, and it remains unclear whether controlled-based RL fine-tuning carries over to discrete diffusion models. In this work, we address these problems affirmatively by formulating a CTRL framework for discrete state spaces with potentially arbitrary action spaces. On one hand, this framework opens a new gate for research at the intersection of control theory and RL in discrete spaces, where policy evaluation, policy gradient theory, and practical algorithms remain underexplored. On the other hand, we demonstrate its application to fine-tuning discrete diffusion models by treating the denoiser probability vector as the action, which offers flexibility for policy parameterization design. This perspective naturally induces a controlled CTMC that explores the output distribution of the denoiser, and yields a clean analytical expression for the policy log-likelihood. Consequently, it bridges the gap between autoregressive-style policy optimization objectives and diffusion-based generative modeling, providing a principled RL framework for fine-tuning discrete diffusion models.

1.1 Contributions↩︎

1. RL theory and algorithms. We first formally propose a CTRL framework in discrete spaces, which to our best knowledge, is novel. Analogous to classical stochastic control and RL in continuous time and space [27], [30], [31], we define controlled CTMC state dynamics and the (optimal) value function, and derive the corresponding analytical forms of Hamilton-Jacobi-Bellman (HJB) equation, optimal control, instantaneous advantage rate (\(q\)-function), and policy gradient. By utilizing these, we propose PPO and GRPO algorithms for general CTRL in discrete spaces.

2. Fine-tuning discrete diffusion models. Inspired by the works [25], [26], we view the sampling process of discrete diffusion models as a stochastic control problem, where the discrete score function is viewed as the control/action/policy. So the problem of fine-tuning discrete diffusion models naturally fits into the CTRL in discrete spaces, thereby solving it via the aforementioned RL algorithms. To this end, we propose several policy parameterization methods (Dirichlet policy, temperature softmax policy, and logistic normal policy), and derive the corresponding policy probability ratios. Our framework is flexible with any predefined process reward models that assign credits to intermediate partially denoised states. We also adopt trajectory subsampling to reduce the number of forward passes during training.

3. Experiments. We conduct experiments on both low-dimensional synthetic data [34], and LLaDA [11], an open-sourced 8B dLLM, for reasoning and coding tasks. For the low-dimensional synthetic example, we illustrate that the PPO algorithm outperforms its alternatives (d1 [18] and DAM [34]) in both generation performance and convergence. For LLaDA experiments, we train our CTRL on the GRPO algorithms, demonstrating strong effectiveness in both mathematical reasoning tasks (Sudoku, GSM8K and MATH500), and coding tasks (HumanEval and MBPP) compared to the state of the arts algorithms – d1 [18], d2 [22] and SPG [21]. For instance, we achieve \(88.2\%\) accuracy on Sudoku with 0-shot prompting training.

1.2 Related Works↩︎

Continuous-Time RL for Fine-Tuning Diffusion Models. [24], [25] introduced a continuous-time RL pipeline to fine-tune continuous diffusion models, where the action is parameterized by Gaussian with mean as the score network, and proposed an actor-critic type PPO algorithm by leveraging \(q\)-function. [26] used a similar idea but enabled pretraining a diffusion model without knowing any prior of score function or learning the score. They theoretically proved that the optimal control is Gaussian distribution, and developed an actor–critic type \(q\)-learning algorithm to solve the continuous-time RL problem.

Fine-Tuning Discrete Diffusion Models via RL. In the early stage, policy gradient methods for fine-tuning score-based discrete diffusion models were proposed to tackle the non-differentiability of categorical distribution inherent in discrete diffusion model sampling or non-differentiability of reward models when maximizing the entropy-regularized reward objective. DRAKES [35] made the originally non-differentiable trajectories differentiable using the Gumbel-Softmax trick for the general CTMC-based discrete diffusion models, but the reward function is required to be differentiable. Similar to our work, they also derived the optimal value function and HJB equation specified to CTMC-based diffusion model dynamics, while our derivation is for the more general controlled CTMC diffusion processes. Score Entropy Policy Optimization [36] fine-tuned general discrete diffusion models over non-differentiable rewards by deriving importance sampling gradient. Later, the rising of mask-based dLLMs [11], [12] drived a line of work focusing on enhancing their reasoning abilities via RL post-training. Due to fundamental modeling differences between the auto-regressive and diffusion generative paradigms for LLMs, dLLMs enable parallel or any-order decoding at the cost of intractable or computationally prohibitive trajectory likelihood. D1 [18] was the first to scale the reasoning ability of dLLMs with GRPO by using mean-field approximation to estimate the intractable policy log probability in a single forward pass. SPG [21] utilized the advantage-weighted GRPO framework and estimated the log likelihood by sandwiching its evidence lower bound (ELBO) and deriving an evidence upper bound (EUBO), where the evidence bounds are estimated via Monte Carlo sampling so multiple forward passes are needed. WD1 [20] adopted the probability ratio estimation from d1 but integrated the estimation of the old and reference policy log likelihoods into an exponential-weighted objective. UniGRPO [19] estimated the log likelihoods via either a one-step unmasking from d1 or Monte Carlo estimation using the ELBO. DIFFPO [37] proposed a two-times mean-field approximation by conditioning on one additional latent at a randomly sampled timestep at each optimization step, yielding a better surrogate policy and a loss contained two importance sampling probability ratios. D2 [22] approximated trajectory likelihood by products of tractable factors over blocks, reducing the number of forward passes from the number of denoising steps to the number of blocks. ESPO [38] modeled entire-sequence generation as a single action and uses the ELBO as a tractable approximation to the sequence-level likelihood. LFPO [39] overcame the likelihood intractability by directly optimizing denoising logits via contrastive positive/negative trajectories with significantly faster inference. In order to exploit the information of intermediate sequence, DISPO [23] proposed to branch from an intermediate sequence by resampling the currently masked positions from rollout-cached logits, scores the resulting completions by the reward model, and updates only the newly filled tokens. TraceRL [40] introduced a diffusion-based value model that enhances training stability and accommodates the process rewards, providing stronger reward supervision. DTRPO [41] adapted direct policy optimization (DPO) [42] to dLLM policy optimization and proposed to estimate the trajectory log likelihood by leveraging subsampling denoising time steps and block-attention mask. Although this estimator can be implemented in a single forward pass, dTRPO is an offline RL method via constructed preference set; estimating the trajectory likelihood in online RL methods like GRPO in a single forward pass is still challenging. LLaDOU [43] defined action at each step to be first determine the set of tokens to unmask and second predict the values of these tokens to obtain the next-time sequence, where some score is predicted to rank masked tokens under the current state, which can be viewed as a specific action policy in our framework. Although JustGRPO [44] indicated that any-order generation during RL training may limit the reasoning capabilities of dLLMs, we still focus on maintaining consistency between rollout generation and policy optimization to align with the continuous-time RL dynamics, as our fine-tuning framework is designed to be general rather than specific to dLLMs.

Organization of the paper: The remainder of the paper is organized as follows. In Section 2, we formulate a general RL framework in continuous time and discrete spaces. Section 3 focuses on the application to fine-tuning discrete diffusion models, with the special case of masked diffusion models discussed in Section 4. We report the experimental results in Section 5. Concluding remarks are provided in Section 6.

2 An RL Framework in Continuous Time and Discrete Spaces↩︎

2.1 Model Formulation↩︎

We formulate CTRL for CTMCs with discrete state spaces. Let \({\cal S}\) be a discrete state space, and \({\cal A}\) be an action space. With an abuse of notation, for continuous-time dynamics (and theoretical development), \(t\in [0,T]\) where \(T\in{\mathbb{R}}_+\) is a finite time horizon; while in the final algorithmic stage where the continuous time range is discretized, \(t\in [T]\) denotes the time step with \(T\in {\mathbb{N}}_+\) being the number of time discretization steps.

We first introduce a controlled transition-rate map \(R: [0,T]\times {\cal S}\times {\cal A}\to {\mathbb{R}}^{|{\cal S}|}\), such that for each fixed \((t,x,a)\in [0,T]\times{\cal S}\times{\cal A}\), it holds that \[\label{rate} R(t,x,a)_y\geq 0 \quad \text{for}\;y \neq x,\; y\in {\cal S}, \quad \text{and}\quad R(t,x,a)_x=-\sum_{y\neq x} R(t,x,a)_y.\tag{1}\]

If we ignore the control \(a\) in 1 , \(R\) is the rate matrix of a time-inhomogeneous CTMC, with the \((x,y)\)-th entry given by \(R(t,x)_y\) at a fixed time \(t\in[0,T]\) (see e.g., [45]). Thus, 1 can be understood as an action-controlled rate matrix for any fixed \(t\). For the realization of such an \(R\), we give an example that is linear in \(a\), encompassing a rich family of transition-rate maps.

Next, given an action sequence \({\mathbf{a}}=(a_t)_{t\in [0,T]}\in {\cal A}^{[0,T]}\), we define the controlled state dynamics \((X_t^{\mathbf{a}})_{t\in [0,T]}\) on \({\cal S}\) by a CTMC with infinitesimal transition probability: \[\label{action} p_{t+\Delta t\mid t}(y\mid x_t^{\mathbf{a}})=\delta\{x_t^{\mathbf{a}},y\}+R(t,x_t^{\mathbf{a}},a_t)_y\cdot\Delta t+o(\Delta t), \quad y \in{\cal S}, \;x_0^{\mathbf{a}}\sim \rho,\tag{2}\] where \(\delta\{\cdot, \cdot\}\) denotes the Kronecker delta, \(x_t^{\mathbf{a}}\) is a realization of \(X_t^{\mathbf{a}}\), \(a_t\) stands for the action/control at time \(t\), and \(\rho\) is the initial state distribution. The zero-sum condition in 1 ensures that 2 is a valid infinitesimal transition probability, and 2 means that \(R(t,x_t^{\mathbf{a}},a_t)_y\) is the rate at which the probability mass moves from state \(x_t^{\mathbf{a}}\) to \(y\) under the control \(a_t\).

Commonly, we consider stochastic policy in RL by exploration to interact with and learn the unknown environment through trial and error. At each time \(t\) with the current state \(x_t\), an action \(a_t\) is generated or sampled from the distribution \(\pi(\cdot\mid t,x_t)\). A policy \(\pi(\cdot\mid t,x)\) is a probability distribution over \({\cal A}\) given the current time-state pair \((t,x)\). Formally, \(\pi(\cdot\mid\cdot,\cdot)\) is a function \(\pi:[0,T]\times {\cal S}\to \Delta({\cal A})\). Denote by \((X_t^\pi)_{t\in [0,T]}\) the state dynamics governed by a feedback policy \(\pi(\cdot\mid \cdot, \cdot)\), which is a CTMC with infinitesimal transition probability: \[\label{pidynamics} p_{t+\Delta t\mid t}(y\mid x_t^\pi)=\delta\{x_t^\pi,y\}+R(t,x_t^\pi,a_t^\pi)_y\cdot\Delta t+o(\Delta t), \quad y \in{\cal S}, \;x_0^\pi\sim \rho,\tag{3}\] where \(x_t^\pi\) is a realization of \(X_t^\pi\) and \(a_t^\pi\sim \pi(\cdot\mid t,x_t^\pi)\). To evaluate such a policy \(\pi\), we consider an exploratory version of state dynamics \((\tilde{X}_t^\pi)_{t\in [0,T]}\) in the same spirit as RL in continuous time and space [27]. Define its transition mechanism by \[\label{RLprocess} p_{t+\Delta t\mid t}(y\mid \tilde{x}_t^\pi)=\delta\{\tilde{x}_t^\pi,y\}+\tilde{R}(t,\tilde{x}_t^\pi;\pi(\cdot\mid t,\tilde{x}_t^\pi))_y\cdot\Delta t+o(\Delta t), \quad y \in{\cal S}, \;\tilde{x}_0^\pi\sim \rho,\tag{4}\] where \(\tilde{x}_t^\pi\) is a realization of \(\tilde{X}_t^\pi\) and \(\tilde{R}(t,x;\pi(\cdot)):=\int_{\cal A}R(t,x,a)\pi(a)\,\mathrm{d}a=\mathop{\mathrm{\mathbb{E}}}_{a\sim \pi} R(t,x,a)\). Note that the external randomization of action \(a\) is removed by taking the expectation, leaving a chain governed by the averaged transition rate \(\tilde{R}\). Importantly, we have the following proposition that allows us to perform theoretical analysis with the action-averaged exploratory state dynamics 4 , while observing real trajectories by simulating 3 . The proof is deferred to Appendix 7.1.

Performance Metric. The goal is to find the optimal feedback policy \(\pi^*\) that maximize the expected reward: \[\label{RLobjective} \pi^*=\mathop{\mathrm{arg\,max}}_{\pi}\mathop{\mathrm{\mathbb{E}}}\left[\int_0^T r(t,X_t^\pi,a_t^\pi)\,\mathrm{d}t + h(X_T^\pi) \mid X_0^\pi\sim\rho \right],\tag{5}\] where \(r: [0,T]\times {\cal S}\times {\cal A}\to {\mathbb{R}}\) is the running reward function, and \(h:{\cal S}\to{\mathbb{R}}\) is the terminal reward function. Also denote the value function of a policy \(\pi\) by \[\begin{align} V(t,x;\pi)={}&\mathop{\mathrm{\mathbb{E}}}\left[\int_t^T r(s,X_s^\pi,a_s^\pi)\,\mathrm{d}s + h(X_T^\pi) \mid X_t^\pi=x \right]\\ ={}&\mathop{\mathrm{\mathbb{E}}}\left[\int_s^T \tilde{r}(s,\tilde{X}_s^\pi;\pi(\cdot\mid s,\tilde{X}_s^\pi))\,\mathrm{d}s + h(\tilde{X}_T^\pi) \mid \tilde{X}_t^\pi=x \right], \end{align}\] where \(\tilde{r}(t,x;\pi(\cdot)):=\int_{\cal A}r(t,x,a)\pi(a)\,\mathrm{d}a=\mathop{\mathrm{\mathbb{E}}}_{a\sim \pi} r(t,x,a)\) is the exploratory reward for any policy \(\pi\in\Delta({\cal A})\).

2.2 Feynman-Kac Formula↩︎

Given \(a\in{\cal A}\), let \({\cal L}^a\) be the infinitesimal generator associated with the process \((X_t^a)_{t\in [0,T]}\) defined by 2 (see e.g., [8]): \[{\cal L}^a f(t,x):=\frac{\partial f}{\partial t}(t,x)+\langle R(t,x,a), f(t,\cdot)\rangle,\quad f:[0,T]\times {\cal S}\to {\mathbb{R}}.\] Here \(\langle\cdot,\cdot\rangle\) denotes inner product between two vectors, i.e., \(\langle R(t,x,a), f(t,\cdot)\rangle=\sum_{y\in{\cal S}} R(t,x,a)_y f(t,y)\).

With the action-aware infinitesimal generator in place, we now state the Feynman-Kac formula for discrete state spaces. The proof is deferred to Appendix 7.3.

2.3 Policy Gradient↩︎

In practice, the policy is parameterized by a function family \(\{\pi^\theta(\cdot\mid\cdot.\cdot)\}_\theta\). We write \((X_t^\theta)_{t\in [0,T]}\) for the process \((X_t^{\pi^\theta})_{t\in [0,T]}\) governed by the policy \(\pi^\theta\). Its transition probability is: \[p^\theta_{t+\Delta t\mid t}(y\mid x^\theta_t)=\delta\{x^\theta_t,y\}+\tilde{R}(t,x^\theta_t;\pi^\theta(\cdot\mid t,x^\theta_t))_y\cdot\Delta t+o(\Delta t), \quad y \in{\cal S}, \;x_0^\theta\sim \rho,\] where \(x_t^\theta\) is a realization of \(X_t^\theta\), and we denote by \(p_t^\theta\) its distribution (a probability mass function). Also denote its corresponding value function by \(V^\theta(t,x)=V(t,x;\pi^\theta(\cdot\mid t,x))\) and \(V^\theta=\mathop{\mathrm{\mathbb{E}}}_{x\sim\rho}V^\theta(0,x)\). The following theorem computes \(\nabla_\theta V^\theta\), whose proof is given in Appendix 7.4.

Based on the derived \(q\)-function, we present a PPO algorithm in Algorithm 1, where the value and policy networks are optimized iteratively. Note that we subtract the action-independent term \(\partial V/\partial t\) when computing \(q\)-function \(q_{b,t}\). We also present a GRPO algorithm that provides process supervision [47], [48] in Algorithm 2, eliminating the need for a value network by directly leveraging the intermediate (running) reward \(r\) and the terminal reward \(h\). While both algorithms employ a uniform time discretization by default, the continuous-time formulation allows for arbitrary discretization schemes. Consequently, more flexible time discretizations in the final algorithmic stage can be adopted for efficiency or accuracy.

3 A General Framework for Fine-Tuning Score-Based Discrete Diffusion Models↩︎

3.1 Discrete Diffusion Models↩︎

In score-based discrete diffusion models [4]–[6], the forward and reverse processes are formulated as CTMCs. Let \(Q_t\in{\mathbb{R}}^{|{\cal S}|\times |{\cal S}|}\) be the transition rate matrix of the forward process \((Y_t)_{t\in [0,T]}\) at time \(t\in [0,T]\) such that \(\mathrm{Law}(Y_0)=p_\mathrm{data}\) and \(\mathrm{Law}(Y_T)\approx p_\mathrm{noise}\). The reverse process \((X_t)_{t\in [0,T]}\) is the exact time reversal of \((Y_t)_{t\in [0,T]}\) [49], that is, \(X_t \overset{\text{d}}{=} Y_{T-t}\), with the rate matrix \(Q_t^\leftarrow\) given by \[Q_t^\leftarrow (x,y)=Q_{T-t}(y,x)\cdot s(T-t, x)_y\quad \text{for}\;y\neq x,\;y\in {\cal S},\;\text{and}\;Q_t^\leftarrow(x,x)=-\sum_{y\neq x}Q_t^\leftarrow(x,y),\;\forall x\in{\cal S}.\] Here \(s(t,x):=(q_t(y)/q_t(x))_{y\neq x}\in {\mathbb{R}}_+^{|{\cal S}|-1}\) is called discrete score function [6], [7], which is a collection of probability ratios of forward marginal distributions. Here, \(q_t\) is the forward marginal at time \(t\), i.e., the distribution (a probability mass function) of \(Y_t\).

In practice, the score \(s(t,x)\) is parameterized by a function family \(\{s_\theta(t,x)\}_\theta\), and is then learned by score matching methods, such as denoising score entropy [6], [8]. Once we have a pretrained estimator \(s_{\theta_{\mathrm{pre}}}(\cdot,\cdot)\), we can perform generative modeling by discretizing the continuous-time reverse process: \[\label{sampling} p^{\theta_\mathrm{pre}}_{t+\Delta t\mid t}(y\mid x_t)=\delta\{x_t,y\}+Q_{T-t}(y,x_t)s_{\theta_\mathrm{pre}}(T-t,x_t)_y\cdot \Delta t+o(\Delta t),\quad y\neq x_t,\;y\in{\cal S},\;x_0\sim p_\mathrm{noise},\tag{6}\] where \(x_t\) is a realization of \(X_t\) and \(p_t^{\theta_\mathrm{pre}}\) is the marginal distribution of the pretraining sampling process with score \(s_{\theta_{\mathrm{pre}}}(\cdot,\cdot)\) at time \(t\).

3.2 Score as Action↩︎

To draw connections between CTRL and discrete diffusion models, we compare the RL process 4 with the diffusion sampling process 6 . Given the forward diffusion rate matrix \((Q_t)_{t\in [0,T]}\), we choose the transition-rate map \(R\) in the RL dynamics 4 as \[\label{diffusionrate} R(t,x,a)_y=Q_{T-t}(y,x)\cdot a_y\quad \text{for}\;y\neq x,\;y\in{\cal S},\tag{7}\] for any \((t,x,a)\in[0,T]\times {\cal S}\times {\cal A}\). Here the action space is specified as \({\cal A}={\mathbb{R}}_+^{|{\cal S}|-1}\), because we are only concerned with the off-diagnal entries. The following proposition shows that the rate map 7 belongs to the family of rate maps in Example 1. We give the proof in Appendix 7.6.

Suppose that we have access to a pretrained score function \(s_{\theta_\mathrm{pre}}(\cdot,\cdot)\), and hence a pretrained model. Let \(\pi^{\theta_\mathrm{pre}}(\cdot\mid t,x)\) be a deterministic policy \(\delta_{s_{\theta_\mathrm{pre}}(T-t,x)}\); that is, if \(a_t^{\theta_\mathrm{pre}}\sim \pi^{\theta_\mathrm{pre}}(\cdot\mid t,x_t)\), then \(a_t^{\theta_\mathrm{pre}}=s_{\theta_\mathrm{pre}}(T-t,x_t)\). So the reverse CTMC 6 becomes \[p^{\theta_\mathrm{pre}}_{t+\Delta t\mid t}(y\mid x_t)=\delta\{x_t,y\}+R(t,x_t,a_t^{\theta_\mathrm{pre}})_y\cdot \Delta t+o(\Delta t)\quad \text{for} \;y\in{\cal S}.\] As a result, analogous to [25], the score function is replaced by the action, and the problem of finding the optimal score becomes finding the optimal action/policy, which can be tackled by policy optimization.

3.3 Post-Training↩︎

Generally, for a parameterized policy \(\pi^\theta(\cdot\mid \cdot,\cdot)\), we denote the control \(a_t^\theta\sim \pi^\theta(\cdot\mid t,X_t^\theta)\), where \(X_t^\theta\) is short for the reverse process \(X_t^{\pi^\theta}\) driven by \(\pi^\theta\): \[p^{\theta}_{t+\Delta t\mid t}(y\mid x^\theta_t)=\delta\{x^\theta_t,y\}+R(t,x^\theta_t,a_t^{\theta})_y\cdot \Delta t+o(\Delta t)\quad \text{for} \;y\in{\cal S}, \;x_0^\theta\sim p_\mathrm{noise},\] where \(x_t^\theta\) is a realization of \(X_t^\theta\) and \(p_t^\theta\) is the distribution of \(X_t^\theta\). The post-training objective is: \[\label{ptobjective} \max_\theta \mathop{\mathrm{\mathbb{E}}}\left[\mathrm{TRF}(X_T^\theta)+\int_0^T \mathrm{IRF}(X_t^\theta)\,\mathrm{d}t-\beta D_\mathrm{KL}({\mathbb{P}}^\theta\,\Vert\, {\mathbb{P}}^{\theta_\mathrm{pre}}) \right].\tag{8}\] Here \(\mathrm{TRF}(\cdot)\) denotes the terminal reward function for the final output \(X_T^\theta\), \(\mathrm{IRF}(\cdot)\) denotes the intermediate reward function (possibly time-dependent) for the intermediate state \(X_t^\theta\), \({\mathbb{P}}^\theta\) and \({\mathbb{P}}^{\theta_\mathrm{pre}}\) are the path measures of the processes \((X_t^\theta)_{t\in [0,T]}\) and \((X_t^{\theta_\mathrm{pre}})_{t\in [0,T]}\), respectively, and \(\beta>0\) is a penalty hyperparameter. The following theorem connects the post-training objective 8 with the RL objective 5 . The proof will be given in Appendix 7.7.

By Theorem 2, the objective in 8 is \[\label{ELBO} \mathop{\mathrm{\mathbb{E}}}_{X_t^\theta\sim p_t^\theta}\left[\underbrace{\mathrm{TRF}(X_T^\theta)}_{h(X_T^\theta)} + \int_0^T \left(\underbrace{\mathrm{IRF}(X_t^\theta)-\beta \sum_{y\neq X_t^\theta} Q_{T-t}(y,X_t^\theta)D_I(\mathop{\mathrm{\mathbb{E}}}_{a_t^\theta\sim \pi^\theta(\cdot\mid t,X_t^\theta)}(a_t^\theta)_y \,\Vert\, s_{\theta_\mathrm{pre}}(T-t,X_t^\theta)_y)}_{\tilde{r}(t,X_t^\theta;\pi^\theta(\cdot\mid t,X_t^\theta))}\right)\mathrm{d}t \right].\tag{9}\] The objective 9 coincides with the RL objective 5 . With this expression, for any \((t,x)\in [0,T]\times {\cal S}\), we define the value function

\[\begin{align} V^\theta(t,x) ={}& \mathop{\mathrm{\mathbb{E}}}_{X^\theta_s\sim p_s^\theta} \Bigg[ \mathrm{TRF}(X_T^\theta) + \int_t^T \Bigg( \mathrm{IRF}(X_s^\theta) \notag\\ &-\beta \sum_{y\neq X_s^\theta} Q_{T-s}(y,X_s^\theta) D_I\!\left( \mathop{\mathrm{\mathbb{E}}}_{a_s^\theta\sim \pi^\theta(\cdot\,\mid\, s,X_s^\theta)} (a_s^\theta)_y \,\middle\Vert\, s_{\theta_{\mathrm{pre}}}(T-s,X_s^\theta)_y \right) \Bigg) \,\mathrm ds \,\Bigg|\, X_t^\theta=x \Bigg]. \label{valuefn} \end{align}\tag{10}\]

To design the fine-tuning algorithms, it suffices to inject the functions \(R\), \(r\), and \(h\) in 7 and 10 into Algorithms 1 and 2, and replace the initial distribution \(\rho\) with \(p_\mathrm{noise}\) and the initial parameter \(\theta_0\) with \(\theta_\mathrm{pre}\). Notably, our framework naturally accommodates arbitrary intermediate rewards, does not require the reward function to be differentiable, and avoids the non-differentiability issue of sampling from a categorical distribution when maximizing the reward (see [35]).

4 Fine-Tuning Masked Diffusion Models↩︎

Although our framework applies to general high-dimensional CTMC-based diffusion models as discussed in Section 3, our focus will be on the widely adopted masked diffusion models (MDMs). In this setting, the score function has a direct correspondence with the denoiser, allowing the policy to be naturally parameterized as a probability distribution over the simplex.

Notations. We use lowercase letters to denote scalars, and boldface lowercase letters to denote vectors. The \(i\)-th entry of a vector \({\mathbf{x}}\) is denoted by \({\mathbf{x}}^i\), or simply \(x^i\) when the context is clear. We use \({\mathbf{x}}^{\backslash i}\) to refer to all dimensions of \({\mathbf{x}}\) except the \(i\)-th, and \({\mathbf{x}}^{\backslash i} \odot \hat{x}^i\) to denote a vector whose \(i\)-th dimension takes the value \(\hat{x}^i\), while the other dimensions remain as \({\mathbf{x}}^{\backslash i}\). For a positive integer \(n\), we denote \(\mathbf{1}_n \in \mathbb{R}^n\) as the column vector of ones, and \(I_n \in \mathbb{R}^{n \times n}\) as the identity matrix. The notation \({\mathbf{e}}_j\) refers to a one-hot column vector with a \(1\) in the \(j\)-th position. We use \(\mathrm{Cat}(\cdot;{\mathbf{p}})\) for a categorical distribution over a one-hot column vector with probabilities given by the column vector \({\mathbf{p}}\). \(\mathbf{1}\{\cdot\}\) is the indicator function.

4.1 Preliminaries on Masked Diffusion Models↩︎

Let \({\cal S}=({\cal V}\cup\{\mathtt{m}\})^L\) be the state space, where \({\cal V}:=\{1,\cdots,V\}\) is the vocabulary space, \(\mathtt{m}\) is the special “mask” token, and \(L\) is the sequence length. The continuous time horizon is set to \(1\). We use superscript \(i \in [L]\) for the sequence index, subscript \(j \in {\cal V}\) for the token vocabulary index, subscript \(t\) for the current time (continuous-time, \(t\in [0,1]\)) or time step (discrete-time, \(t\in [T]\)), and subscript \(g\in [G]\) for a specific sample trajectory.

Forward Process. The forward process is factorized as \(p_{t|0}({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\prod_{i=1}^L p_{t|0}(x_t^i\mid x_0^i)\), where \[p_{t|0}(x_t^i\mid x_0^i)=\mathrm{Cat}(x_t^i;\alpha_t {\mathbf{e}}_{x_0^i}+(1-\alpha_t){\mathbf{e}}_\mathtt{m}),\quad t\in[0,1].\] \(\alpha_t\) is the noise scheduled to satisfy \(\alpha_0\approx 1\) and \(\alpha_1\approx 0\). A common choice is \(\alpha_t=1-t\) [11], which we adopt throughout this work. The forward process interpolates the clean data and pure mask noise. The resulting forward rate matrix for each dimension is \[Q_t^\mathrm{tok}=-\frac{1}{\alpha_t}\frac{\partial \alpha_t}{\partial t}(\mathbf{1}_{V+1}\cdot {\mathbf{e}}_\mathtt{m}^\top-I_{V+1})\in{\mathbb{R}}^{(V+1)\times (V+1)},\] and the full-dimensional rate matrix \(Q_t=\oplus_{i=1}^L Q_t^\mathrm{tok} \in {\mathbb{R}}^{(V+1)^L\times (V+1)^L}\) is the Kronecker sum of \(L\) dimension-wise rate matrices, with non-diagnal entries given by \[Q_t({\mathbf{x}},{\mathbf{x}}^{\backslash i}\odot \mathtt{m})=-\frac{1}{\alpha_t}\frac{\partial \alpha_t}{\partial t}=\frac{1}{1-t},\quad x^i\neq \mathtt{m}.\]

Reverse Process. Given \({\mathbf{x}}_0\sim p_\mathrm{data}\), for \(0\leq s\leq t\), the reverse transition probability is \[p_{s|t,0}(x_s^i\mid {\mathbf{x}}_t,{\mathbf{x}}_0)=\begin{cases} \mathrm{Cat}(x_s^i;{\mathbf{e}}_{x_t^i}), & x_t^i\neq \mathtt{m};\\ \mathrm{Cat}(x_s^i;\frac{1-\alpha_s}{1-\alpha_t}{\mathbf{e}}_\mathtt{m}+\frac{\alpha_s-\alpha_t}{1-\alpha_t}{\mathbf{e}}_{x_0^i}),& x_t^i=\mathtt{m}. \end{cases}\] Note that conditional on \({\mathbf{x}}_0\), if \(x_t^i=\mathtt{m}\), then with probability \(\frac{\alpha_s-\alpha_t}{1-\alpha_t}\), \(x_t^i\) will jump to \(x_0^i\) at time \(s\); once \(x_t^i\) is unmasked, it remains the same until \(t=0\). Thus, one can parameterize a denoiser \({\mathbf{p}}_\theta({\mathbf{x}}_t)\), which takes all the dimensions of \({\mathbf{x}}_t\) as the input and predicts masked tokens (\(x_t^i=\mathtt{m},\;i\in [L]\)) simultaneously. To be more specific, for any \(i\in [L]\), a neural network \({\mathbf{p}}_\theta^{(i)}({\mathbf{x}}_t)\), the \(i\)-th output of \({\mathbf{p}}_\theta({\mathbf{x}}_t)\), outputs a probability distribution of \(x_0^i\): \[{\mathbf{p}}_\theta^{(i)}({\mathbf{x}}_t)=\begin{cases} \mathrm{softmax}({\mathbf{f}}_\theta^{(i)}({\mathbf{x}}_t)), &i\in[L]:x_t^i=\mathtt{m},\\ {\mathbf{e}}_{x_t^i},&i\in[L]:x_t^i\neq \mathtt{m} \end{cases}\in\Delta^{V-1},\] where \(\Delta^{V-1}:=\{{\mathbf{x}}\in{\mathbb{R}}_+^V:\sum_{i=1}^V {\mathbf{x}}^i=1\}\subset {\mathbb{R}}^V\) is the \((V-1)\)-dimensional simplex. For \(i:x_t^i=\mathtt{m}\), \({\mathbf{f}}_\theta^{(i)}({\mathbf{x}}_t)\in{\mathbb{R}}^V\) is the model logits and \({\mathbf{p}}_\theta^{(i)}({\mathbf{x}}_t)\in \Delta^{V-1}\) is a probability vector with the \(j\)-th component, denoted by \(p^{(i)}_\theta({\mathbf{x}}_t)_j\) or \(p_\theta^{(i)}(j\mid {\mathbf{x}}_t)\), predicting the conditional probability \(p_{0|t}(x_0^i=j\mid {\mathbf{x}}_t)\) for all \(j\in {\cal V}\). Notably, by the identity [9], [51], [52]: \[s(t,{\mathbf{x}}_t)_{i,j}:=\frac{p_t({\mathbf{x}}_t^{\backslash i}\odot j)}{p_t({\mathbf{x}}_t)}=\delta\{\mathtt{m},j\}+\frac{\alpha_t}{1-\alpha_t}p_{0|t}(x_0^i=j\mid {\mathbf{x}}_t),\quad x_t^i=\mathtt{m},\] parameterizing the denoising probability is equivalent to parametrizing the score function up to some scalar: \[s_\theta(t,{\mathbf{x}}_t)_{i,j}=\frac{\alpha_t}{1-\alpha_t}p_\theta^{(i)}({\mathbf{x}}_t)_j,\quad x_t^i=\mathtt{m}, \;j\in {\cal V}.\] Here the denoiser network does not depend on the time \(t\) because the partially unmasked sequence \({\mathbf{x}}_t\) implicitly contained the time information [9], [11], [53]. In the remainder of this work, we always consider the denoiser parameterization \({\mathbf{p}}_\theta({\mathbf{x}}_t)\) (which will be replaced by actions) instead of the score parameterization. Also, the reverse rate matrix is \[Q^\leftarrow_{1-t}({\mathbf{x}},{\mathbf{x}}^{\backslash i}\odot j)=-\frac{1}{\alpha_t}\frac{\partial \alpha_t}{\partial t}s(t,{\mathbf{x}})_{i,j}=-\frac{1}{1-\alpha_t}\frac{\partial \alpha_t}{\partial t}p_{0|t}(x_0^i=j\mid {\mathbf{x}})=\frac{1}{t}p_{0|t}(x_0^i=j\mid {\mathbf{x}}),\quad x^i=\mathtt{m},\;j\in{\cal V},\] Hence, the reverse process parametrized by \({\mathbf{p}}_\theta\) has rate matrix \[\label{rate:reverse} Q^\theta_t({\mathbf{x}},{\mathbf{x}}^{\backslash i}\odot j)=\frac{1}{1-t}p_\theta^{(i)}({\mathbf{x}})_j,\quad x^i=\mathtt{m},\;j\in {\cal V}.\tag{11}\]

4.2 Post-Training↩︎

Policy Parameterization. In RL, an action \({\mathbf{a}}_t\) is sampled from a parameterized stochastic policy to encourage exploration: \({\mathbf{a}}_t\sim \pi^\theta(\cdot\mid t,{\mathbf{x}}_t)\) given the current time-state pair \((t,{\mathbf{x}}_t)\). Since the reverse masked diffusion dynamics 11 is fully determined by the denoiser probability vector, we consider the action space \(\mathcal{A}=(\Delta^{V-1})^L\). Then \(\pi(\cdot\mid t,{\mathbf{x}})\) is a distribution over the product simplex \((\Delta^{V-1})^L\). For simplicity, we let the parameterized policy \(\pi^\theta\) factorize over all dimensions: \[\pi^\theta({\mathbf{a}}_t\mid t,{\mathbf{x}}_t)=\prod_{i=1}^L \pi^{\theta,(i)}({\mathbf{a}}_t^{(i)}\mid t,{\mathbf{x}}_t),\] where \(\pi^{\theta,(i)}(\cdot\mid t,{\mathbf{x}}_t)\) is a distribution over \(\Delta^{V-1}\), and \({\mathbf{a}}_t^{(i)}\in\Delta^{V-1}\) is a probability vector over \({\cal V}\) for all \(i\in [L]\). At time \(t\), we sample from \(\mathrm{Cat}(\cdot;{\mathbf{a}}_t^{(i)})\), instead of \(\mathrm{Cat}(\cdot;{\mathbf{p}}_\theta^{(i)}({\mathbf{x}}_t))\), to predict the clean token at the \(i\)-th position. So \({\mathbf{a}}_t^{(i)}\) can be viewed as performing exploration over the denoiser \({\mathbf{p}}_\theta^{(i)}({\mathbf{x}}_t)\). As a result, we need to specify the dimension-wise policies \(\pi^{\theta,(i)}(\cdot\mid t,{\mathbf{x}})\) for all \(i\).

We emphasize that the policy parameterization must be chosen carefully. In particular, \(\pi^\theta(\cdot\mid t,{\mathbf{x}})\) should have a tractable likelihood, and the resulting probability ratio \(\rho_t^\theta := \pi^\theta({\mathbf{a}}_t\mid t,{\mathbf{x}}_t)/\pi^{\theta_{\mathrm{old}}}({\mathbf{a}}_t\mid t,{\mathbf{x}}_t)\) should be amenable to stable optimization. Several policy parameterizations that perform full-logit exploration and induce distributions supported on the full simplex, including Dirichlet, temperature softmax, and logistic-normal policies, are summarized in Appendix 8.1, together with their practical implementation details. However, for masked dLLMs with a large vocabulary size \((|{\cal V}|\approx 1.2\times10^5)\), it is inefficient to explore the full simplex: the learned denoiser distribution \({\mathbf{p}}_\theta\) is typically highly concentrated, with probabilities spanning several orders of magnitude and effectively residing on a much lower-dimensional region of the simplex, often near a small number of vertices. As discussed in Appendix 8.1, applying full-simplex exploration in this regime amounts to exploring the entire vocabulary, which can lead to unstable optimization due to the large vocabulary size. Motivated by the empirical success of d1 [18], we consider a policy family whose support is concentrated on the \(V\) vertices of the simplex \(\Delta^{V-1}\), assigning probability mass directly to the vertices rather than exploring its interior through a continuous density. Nevertheless, the CTRL perspective suggests that broader exploration of the denoiser distribution still be beneficial. Consequently, while vertex-supported policies are used for optimization, we retain simplex-level exploration during rollout generation in our experiments.

Following prior works [18], [21], [23], [37], [40], we focus on the GRPO algorithm for fine-tuning masked dLLMs to avoid training an expensive value network in PPO, though we also describe a possible value function approximation without training a network in Appendix 9. Let \(({\mathbf{x}}_t)_{t=0}^T\) be a decoding trajectory, where \({\mathbf{x}}_0\) is a prompt concatenated with mask tokens of the completion length \(L\), \({\mathbf{x}}_T\) is the final output corresponding to the prompt, and \(T \in{\mathbb{Z}}_+\) is the number of denoising steps. The trajectory can be generated by any inference strategy, while we adopt a semi-autoregressive strategy with low-confidence remasking and random sampling [11], [18]. Specifically, at time step \(t\), we first select the \(k\) masked positions with the highest confidence score in the current block \(B_t\subset [L]\), denoted by \[\mathcal{M}_t^{\mathrm{TopkConf}} := \left\{ i\in B_t : x_t^i=\mathtt{m},\; c_t^i \in \operatorname{Topk} \left( \{c_t^i\}_{i\in B_t:x_t^i=\mathtt{m}} \right) \right\},\quad c_t^i := \max_{j\in\mathcal{V}} p_{\theta}^{(i)}(\mathbf{x}_t)_j,\] to unmask, and then sample according to \(\mathrm{Cat}(\cdot;{\mathbf{a}}_t^{(i)})\) for all \(i\in {\cal M}_t^{\mathrm{TopkConf}}\): \[\label{our95sample} {\mathbf{a}}_t^{(i)}\sim \pi^{\theta,(i)}(\cdot\mid t,{\mathbf{x}}_t)\quad \text{then}\quad {\mathbf{a}}_t^{(i)}= \begin{cases} {\mathbf{e}}_j \quad \text{w.p.}\;p_\theta^{(i)}({\mathbf{x}}_t)_j,\;j\in{\cal V}, &\text{for}\;i\in{\cal M}_t^{\mathrm{TopkConf}},\\ {\mathbf{e}}_{x_t^i}, &\text{for}\;i\notin {\cal M}_t^{\mathrm{TopkConf}} \end{cases}\in \Delta^{V-1}.\tag{12}\]

The corresponding intermediate and terminal reward functions are given by 9 : \[\label{ourreward} \tilde{r}(t,{\mathbf{x}}_t;\pi^\theta(\cdot\mid t,{\mathbf{x}}_t))=\mathrm{IRF}({\mathbf{x}}_t)-\frac{\beta}{1-t}\sum_{i:x_t^i=\mathtt{m}} D_\mathrm{KL}({\mathbf{p}}_\theta^{(i)}({\mathbf{x}}_t) \,\Vert\, {\mathbf{p}}_{\theta_\mathrm{pre}}^{(i)}({\mathbf{x}}_t))\quad \text{and}\quad h({\mathbf{x}}_1)=\mathrm{TRF}({\mathbf{x}}_1),\tag{13}\] and the probability ratio is \[\label{ourratio} \rho_t^\theta=\frac{\pi^\theta({\mathbf{a}}_t\mid t,{\mathbf{x}}_t)}{\pi^{\theta_\mathrm{old}}({\mathbf{a}}_t\mid t,{\mathbf{x}}_t)}=\frac{\prod_{i\in{\cal M}_t^\mathrm{TopkConf}}\pi^{\theta,(i)}({\mathbf{a}}_t^{(i)}\mid t,{\mathbf{x}}_t)}{\prod_{i\in{\cal M}_t^\mathrm{TopkConf}}\pi^{\theta_\mathrm{old},(i)}({\mathbf{a}}_t^{(i)}\mid t,{\mathbf{x}}_t)}=\prod_{i\in{\cal M}_t^\mathrm{TopkConf}}\frac{p_\theta^{(i)}(x_T^i\mid {\mathbf{x}}_t)}{p_{\theta_\mathrm{old}}^{(i)}(x_T^i\mid {\mathbf{x}}_t)}.\tag{14}\] Thus, the GRPO loss is

\[\label{ourloss} {\cal L}(\theta)=\mathop{\mathrm{\mathbb{E}}}_{({\mathbf{x}}_{g,0},{\mathbf{x}}_{g,1},\cdots,{\mathbf{x}}_{g,T})\}_{g=1}^G\sim \pi^{\theta_\mathrm{old}}}\left[ \frac{1}{G}\sum_{g=1}^G \frac{1}{T} \sum_{t=1}^T \min\left( \rho_{g,t}^\theta A_{g,t}, \mathrm{clip}\left(\rho_{g,t}^\theta,1-\epsilon,1+\epsilon\right) A_{g,t} \right) \right],\tag{15}\] where \(T=L/k\), \(\rho_{g,t}^\theta=\prod_{i\in{\cal M}_{g,t}^\mathrm{TopkConf}}p_\theta^{(i)}(x_{g,T}^i\mid {\mathbf{x}}_{g,t})/p_{\theta_\mathrm{old}}^{(i)}(x_{g,T}^i\mid {\mathbf{x}}_{g,t})\), and the computation of \(A_{g,t}\) is implied by 13 and Algorithm 2.

Practical Considerations. Our CTRL-based GRPO loss in 15 cannot be evaluated in a single causal forward pass over an entire denoising trajectory due to the evaluation of \(\{\rho_{g,t}^\theta\}_{t\in [T]}\). In online GRPO-based methods for dLLMs, this leads to a fundamental trade-off: on one hand, accurate likelihood computation under a CTRL formulation and the incorporation of intermediate rewards to optimize intermediate steps both require conditioning on the intermediate state \({\mathbf{x}}_t\) at each time step \(t\); on the other hand, this necessitates multiple forward passes along the trajectory, incurring computational and memory costs proportional to the prompt batch size, group size \(G\), and number of denoising steps \(T\), as all intermediate activations must be stored for backpropagation. Existing approaches make different compromises. For instance, d1 [18] computed the loss using a single forward pass by conditioning only on the perturbed prompt. DTRPO [41] estimated the trajectory probability via a single forward pass of a re-masked final state; however, it is designed for offline DPO-style objectives and is difficult to extend to online GRPO settings. D2 [22] utilized block composite likelihood as an estimator to reduce the number of forward passes.

In practice, to reduce the number of required forward passes in our framework, we propose trajectory subsampling, which yields an unbiased estimator of the full GRPO loss in 15 . At each optimizer update we draw a subset \({\cal T}\subset [T]\) of size \(N=|{\cal T}|\) uniformly at random without replacement, the same subset being reused for the policy, old, and reference passes so the ratio 14 is consistent. We then form the subsampled estimator \[\hat{{\cal L}}_{\cal T}(\theta)=\mathop{\mathrm{\mathbb{E}}}_{({\mathbf{x}}_{g,0},{\mathbf{x}}_{g,1},\cdots,{\mathbf{x}}_{g,T})\}_{g=1}^G\sim \pi^{\theta_\mathrm{old}}}\left[ \frac{1}{G}\sum_{g=1}^G \frac{1}{N} \sum_{t\in {\cal T}} \min\left( \rho_{g,t}^\theta A_{g,t}, \mathrm{clip}\left(\rho_{g,t}^\theta,1-\epsilon,1+\epsilon\right) A_{g,t} \right) \right],\] which reduces the number of forward passes from \(T\) to \(N\) per rollout. One can easily see that this estimator is unbiased by the fact that \({\mathbb{P}}(t\in {\cal T})=N/T\): \(\mathop{\mathrm{\mathbb{E}}}_{\cal T}[\hat{{\cal L}}_{\cal T}(\theta)]={\cal L}(\theta)\). When \(N\) is too small, the model fails to reliably converge to competitive performance due to inaccurate loss estimation. Conversely, increasing \(N\) substantially raises the number of forward passes, thereby reducing training efficiency. Our experiments indicate that \(N=8\) provides a favorable trade-off, achieving performance comparable to \(N=16\) and \(N=32\) while requiring substantially fewer function evaluations. More theoretical derivation and practical implementation details including the ablation study on \(N\) are provided in Appendix 8.2.

5 Experiments↩︎

We showcase our method for fine-tuning MDMs on two benchmarks: \(2\)-dimensional synthetic example with PPO, and higher-dimensional reasoning and coding dataset for dLLMs with GRPO.

5.1 Synthetic Example↩︎

Setup. We consider a synthetic example described in Discrete Adjoint Matching (DAM) [34], where we fine-tune an MDM on a \(90\times 90\) discrete grid so that its terminal sampling distribution matches a prescribed checkerboard target (the first subfigure in Figure 3). The target is the entropy-regularized optimum \(p^*\propto p^\mathrm{base}\,e^{h/\beta}\). We solve the task with our proposed PPO algorithm.

We use a two-dimensional discrete space \({\cal X}\;=\; \{\mathtt{m},1,2,\dots,90\}^2\), and a state is a pair \({\mathbf{x}}=(x^1,x^2)\) in which each coordinate is either masked or holds a token. The MDM begins at the fully masked state \(\mathbf{X}_0=(\mathtt{m},\mathtt{m})\) and unmasks exactly one token per step, so there are \(T=2\) denoising steps. The base/pretrained denoiser is uniform: at every masked position it predicts a uniform distribution over the \(V\) tokens, so the base terminal distribution \(p^\mathrm{base}_1\) is uniform on the grid and all target structure is induced by the reward. We use the softmax exponential-temperature exploration policy described in Appendix 8.1.2. The PPO advantage is the value difference across one unmasking step. Because only the states reachable in the two-step process matter, we store an explicit logit table indexed by the context (position, other coordinate), where the other coordinate is a revealed token in \(\{1,\cdots,90\}\) or the mask token, and we evaluate the value \(V^\theta\) exactly by computing the surrogate’s expectation over all \(V\) tokens for each visited context in closed-form. Further implementation details are provided in Appendix 10.1.

Results. We select DAM [34] and d1 [18] as the baselines. The experimental results are presented in Figure 3. PPO aligns most closely with \(p^*\), both visually and numerically. In this \(L=2\) tabular setting, PPO uses an exact critic, so it converges faster and tighter than DAM’s sample-based adjoint. d1 was designed for large-vocabulary reasoning with informative prompts, so this synthetic distribution-matching task isolates the objective rather than d1’s intended use case.

5.2 Mathematical Reasoning and Coding Tasks↩︎

Setup. We evaluate our GRPO algorithm on four standard mathematical reasoning tasks: GSM8K [54], which consists of multi-step grade-school math word problems; MATH500 [55], a collection of high-school competition mathematics problems; \(3\)-number Countdown, a combinatorial arithmetic game in which models must reach a target number using basic arithmetic operations on three given numbers; and \(4\times 4\) Sudoku, which requires constraint satisfaction and systematic reasoning to complete the grid. We employ LLaDA-8B-Instruct [11] without supervised fine-tuning as the base model, and adopt the experiment setting in d1 [18]. We choose d1 [18], SPG with mixture [21], and d2 [22] as the baselines: d1 is the first GRPO-based fine-tuning algorithm for dLLMs, and SPG and d2 are included because they are representative methods that outperform d1, while requiring multiple forward passes during training. For each method, we fine-tune a separate model on each task with completion lengths of \(128\) and \(256\) in \(1000\) steps (due to the time cost), and evaluate each resulting model at completion lengths of \(128\) and \(256\), accordingly.

Results. The correctness reward curves during the RL training with a completion length of \(128\) are depicted in Figure 4. For the Sudoku experiments, we note that SPG uses \(3\)-shot prompting during training, whereas all four methods in our experiments are trained with \(0\)-shot prompting. Despite this more challenging setting, CTRL consistently converges to a correctness reward of \(0.92\) on Sudoku in \(1000\) training steps, whereas d2 converges much more slowly, and both d1 and SPG plateau at around \(0.25\). We note that our success on Sudoku can be attributed to using the intermediate reward as process supervision, which guides the model to generate valid answers (see Figure [fig:sudoku95interreward] in Appendix). Further experimental details including the specification of intermediate reward functions are provided in Appendix 10.2. Table 1 reports the best evaluation results achieved by each method over all checkpoints. For fair comparison, all evaluations are conducted with \(0\)-shot prompting, using the evaluation code from d1 with a batch size of \(8\). Across nearly all experimental settings except Countdown, CTRL achieves comparable or superior accuracy to all baselines, demonstrating its effectiveness for fine-tuning moderate-sized dLLMs with moderate sequence lengths. In particular, CTRL substantially outperforms the competing methods on Sudoku with a completion length of \(128\); all methods collapse when the completion length is increased to \(256\), suggesting that shorter sequence lengths are better suited to this task. By contrast, the performance gains on GSM8K and MATH500 are more modest across different baseline methods, which likely reflects the limited pretraining capability of the base model.

Extension to Coding. We further extend CTRL to coding tasks, fine-tuning the LLaDA-8B-Instruct base model on the KodCodeLight-RL-10K dataset [56] and evaluating on HumanEval and MBPP benchmarks. As shown in Table 2, CTRL consistently improves the accuracy on all benchmarks over the baselines across different generation lengths, demonstrating its strong ability in the coding domain.

6 Conclusion and Future Work↩︎

We introduce a theoretical framework of reinforcement learning in continuous time and discrete spaces, with applications to fine-tuning score-based discrete diffusion models, especially masked diffusion models by setting the score/denoiser as the action. Our framework naturally accommodates intermediate rewards into the RL objective, and leads to algorithms that do not require differentiable rewards. We demonstrate the effectiveness of our methods by both solving low-dimensional entropy-regularized optimization problems and fine-tuning diffusion large language models on reasoning and coding tasks. Future work will focus on developing refined techniques to reduce the number of function evaluations associated with the computation of importance sampling probability ratios under the continuous-time RL framework, and improve the computational efficiency of PPO-based RL post-training for large-scale problems.

Acknowledgement↩︎

Tang is supported by NSF CAREER Award DMS-2538791 and the Tang Family Assistant Professorship. Sheng and Zhang are supported by NSF Grant DMS-2206038. This research is part of a Columbia-CityU/HK collaborative project that is supported by InnoHK Initiative, The Government of the HKSAR and the AIFT Lab.

References↩︎

---

1. Department of Industrial Engineering and Operations Research, Columbia University, New York, NY 10027. Emails: zz3367@columbia.edu, js6646@columbia.edu, yao@columbia.edu, wt2319@columbia.edu.↩︎

2. https://github.com/dllm-reasoning/d1↩︎