Mask-Aware Policy Gradients for Diffusion Language Models
July 16, 2026
Abstract
Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Optimizing both terms achieves state-of-the-art results across mathematical reasoning and coding benchmarks, reaching 87.1% on GSM8K and 53.4% on MBPP.1
1 Introduction↩︎
Reinforcement learning has emerged as a powerful tool for improving large language models [1], [2], with applications spanning instruction following [3], alignment [4], and complex reasoning [1], [2]. At the core of these methods, policy gradient algorithms sample group rollouts autoregressively and estimate their log-likelihoods under the policy. The key to their success is the tractability and simplicity of the autoregressive log-likelihood: the generation trajectory is a sequence of token prediction actions, and the log-likelihood is their product.
Masked Diffusion Language Models (MDLMs) [5]–[7] offer an alternative to autoregressive LLMs, replacing the fixed left-to-right generation order with iterative unmasking from a fully masked sequence [8], [9]. This enables parallel decoding and flexible generation orderings, and MDLMs have proven effective across text generation, code generation, mathematical reasoning, and protein sequence generation [5], [10], [11]. However, this flexibility comes at a cost: unlike autoregressive models, MDLMs lack a tractable log-likelihood estimation and are instead trained via an evidence lower bound (ELBO). This makes adapting policy gradient algorithms to MDLMs significantly more challenging.
Existing methods approximate the log-likelihood using ELBO variants [12]–[14] or generation trajectories [15], focusing on the token prediction component. However, the MDLM generation trajectory has a richer structure than its autoregressive counterpart. While an autoregressive trajectory is a sequence of token prediction actions at fixed positions, an MDLM trajectory involves two decisions at each step: what token to place at the masked positions and which of these to remask.3 The trajectory log-likelihood therefore decomposes into a product of both token prediction and remasking probabilities. Since the policy determines the unmasking order, this richer trajectory structure provides an additional signal for policy optimization that prior ELBO- and trajectory-based methods do not exploit. See 2 for an overview.
We formalize the unmasking order as part of the policy. We model MDLM generation as a two-stage action MDP: at each step, the model selects a token for each masked position and then decides which positions to remask. This factorization yields a policy gradient that decomposes into two terms, one reinforcing token predictions and one reinforcing the ordering. The position term is computed from the model’s own logits at no cost beyond the token-likelihood computation, and is compatible with any standard policy gradient algorithm. From a theoretical perspective, we show that ignoring the position component of the policy gradient can miss directions that improve expected return. Empirically, jointly optimizing both terms leads to consistent improvements over prior methods, achieving gains of up to 4.0% on MATH500, 2.9% on HumanEval, and 2.5% on GSM8K and MBPP (1).
Our contributions are as follows:
We formalize MDLM generation as a two-stage action MDP and show that the policy gradient decomposes into a token term and a masking term, providing theoretical grounding for optimizing the unmasking order.
We derive position log-probabilities as a softmax over unmasking scores from the model’s own logits, requiring no additional parameters or architectural changes, and no forward passes beyond those already used for the token likelihood.
We achieve state-of-the-art results across mathematical reasoning and code generation benchmarks.
2 Preliminaries↩︎
A language model \(\pi_{\boldsymbol{\theta}}\) produces sequences of tokens \(\boldsymbol{x}_{1:n} \sim \pi_{\boldsymbol{\theta}}(\cdot \!\mid\! \boldsymbol{c})\) conditioned on an input \(\boldsymbol{c}\). The joint distribution over sequences is exponentially large. Autoregressive language models [1]–[3], [16], decompose the distribution into a product of per-token conditionals: \[\pi_{\boldsymbol{\theta}}(\boldsymbol{x}_{1:n} \mid \boldsymbol{c}) = \prod_{k=1}^{n} \pi_{\boldsymbol{\theta}}(x_k \mid \boldsymbol{x}_{<k}, \boldsymbol{c}). \label{eq:ar95logprob}\tag{1}\] This makes likelihood estimation tractable, and simplifies both training and generation, see 3.
2.0.0.1 Masked Diffusion Language Models (MDLMs)
[5]–[7] formulate text generation as an iterative denoising process. The forward process corrupts a clean sequence \(\boldsymbol{x}_{1:n}\) into \(\boldsymbol{z}^t_{1:n}\) by independently replacing each token with a special \(\texttt{[MASK]}\) token with probability \(\alpha_t \in [0, 1]\): \[\boldsymbol{z}^t_{k} = \begin{cases} \boldsymbol{x}_k & \text{with probability } 1 - \alpha_t \\ \texttt{[MASK]}& \text{with probability } \alpha_t \end{cases}\] Here \(t=1\ldots T\) is the diffusion step, and \(\alpha_t\) is the corresponding masking probability. We write the forward process as \(\boldsymbol{z}^t_{1:n} \sim q_{t}(\cdot \mid \boldsymbol{x}_{1:n})\). At the final step \(t=T\), the sequence \(\boldsymbol{z}^T=\boldsymbol{x}\) remains uncorrupted (\(\alpha_T=0\)). At the first step \(t=0\), \(\boldsymbol{z}^0\) is fully masked (\(\alpha_0=1\)).
A neural network \(\pi_{\boldsymbol{\theta}}\) learns to reverse this process by predicting the original tokens at masked positions from a corrupted input \(\boldsymbol{z}\), maximizing the Evidence Lower Bound (ELBO) of the data log-likelihood: \[\mathcal{L}_{\mathrm{ELBO}}(\boldsymbol{x};\boldsymbol{\theta}) = \mathop{\mathrm{\mathbb{E}}}_{\substack{t \sim \mathcal{U}\{1,\ldots,T\} \\ \boldsymbol{z}^t \sim q_{t}(\cdot \mid \boldsymbol{x})}} \bigg[\sum_{k=1}^{n} \frac{1}{\alpha_t} \mathbb{1}(\boldsymbol{z}^t_{k} = \texttt{[MASK]}) \log \pi_{\boldsymbol{\theta}}(\boldsymbol{x}_k \mid \boldsymbol{z}^t) \bigg], \label{eq:elbo}\tag{2}\] In essence, the model learns to predict the original token at each masked position, with unmasked positions left unchanged.
2.0.0.2 Generation in MDLMs
proceeds as described in [alg:mdlm]: We start from a completely masked sequence \({\boldsymbol{z}^0 = [\texttt{[MASK]}]^n}\) of length \(n\) and iteratively denoise the sequence over \(T\) steps with masking schedule \(1 = \alpha_0 > \alpha_1 > \ldots > \alpha_T = 0\). Each step \(t\) consists of a prediction followed by a remasking. First, the model predicts tokens at all masked positions: \({\hat{\boldsymbol{z}}^{t} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{z}^{t-1})}\), with \(\hat{z}^{t}_k = z^{t-1}_k\) for already-unmasked positions (\(z^{t-1}_k \ne \texttt{[MASK]}\)). Then, a subset \({\mathcal{U}_t \subseteq \{k : z^{t-1}_k = \texttt{[MASK]}\}}\) of newly predicted positions is selected to remain unmasked, and the rest are remasked: \[z^{t}_k = \begin{cases} \hat{z}^{t}_k & \text{if } k \in \mathcal{U}_t \text{ or } z^{t-1}_k \ne \texttt{[MASK]}, \\ \texttt{[MASK]}& \text{otherwise.} \end{cases} \label{eq:remask}\tag{3}\] This produces a trajectory \(\mathbf{z} = (\boldsymbol{z}^0, \boldsymbol{z}^1, \ldots, \boldsymbol{z}^T)\) where \(\boldsymbol{z}^T = \boldsymbol{x}_{1:n}\) is the final output. We write \(\mathbf{z} \sim \pi_{\boldsymbol{\theta}}(\cdot \!\mid\! \boldsymbol{c})\) for the distribution over trajectories induced by the model \(\pi_{\boldsymbol{\theta}}\) when generating from prompt \(\boldsymbol{c}\).
In theory, remasking should be random to mimic the training process. In practice however, a deterministic strategy that selects positions to unmask based on confidence scores derived from the model’s logits [9], [17] often performs much better. For example, the model may select the top-\(k\) positions with the highest maximum logit values to unmask at each step. This allows the model to focus on positions it is most confident about, potentially improving generation quality and convergence speed.
2.1 Reinforcement learning for language models↩︎
RL optimizes the expected return \[J(\boldsymbol{\theta}) = \mathop{\mathrm{\mathbb{E}}}_{\boldsymbol{x} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c})}[R(\boldsymbol{c}, \boldsymbol{x})]. \label{eq:pg}\tag{4}\] In the context of language models, \(\boldsymbol{x}\) is the output produced by the model and \(R\) is the reward obtained via a rule-based verifier or a learned reward model. Most RL algorithms for language models rely on a policy gradient estimate: \({\nabla_{\boldsymbol{\theta}} J = \mathbb{E}_{\boldsymbol{x} \sim \pi_{\boldsymbol{\theta}}} \big[ R(\boldsymbol{c}, \boldsymbol{x}) \, \nabla_{\boldsymbol{\theta}} \log \pi_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{c}) \big]}\).
In reinforcement learning, \(\boldsymbol{x} \sim \pi_{\boldsymbol{\theta}}\) is referred to as the action. For autoregressive models, this action is a sequence of token predictions, and \(\log \pi_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{c})\) is directly available from 1 : the chain rule decomposition gives a sum of per-token log-probabilities, computed in a single forward pass. This makes the policy gradient fully tractable.
For MDLMs, however, \(\log \pi_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{c})\) is intractable to compute, as it is defined by the marginal over all trajectories that produce \(\boldsymbol{x}\): \[\pi_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{c}) = \sum_{\mathbf{z} : \boldsymbol{z}^T = \boldsymbol{x}} \pi_{\boldsymbol{\theta}}(\mathbf{z} \mid \boldsymbol{c}).\] Existing approximations fall into two families:
2.1.0.1 Evidence bound methods
[12]–[14] estimate the log-likelihood by evaluating the model on randomly masked versions of the completed sequence, typically computing some variant of the ELBO objective (Eq. 2 ) used for training.
2.1.0.2 Trajectory-based methods
[15] rewrite the expected return as \[J(\boldsymbol{\theta}) =\!\!\!\mathop{\mathrm{\mathbb{E}}}_{\boldsymbol{x} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c})}\!\!\![R(\boldsymbol{c}, \boldsymbol{x})] = \sum_{\boldsymbol{x}}\pi_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{c})R(\boldsymbol{c}, \boldsymbol{x}) = \sum_{\mathbf{z}} \pi_{\boldsymbol{\theta}}(\mathbf{z} \mid \boldsymbol{c})R(\boldsymbol{c}, \boldsymbol{z}^T) =\!\!\!\!\mathop{\mathrm{\mathbb{E}}}_{\mathbf{z} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c})}\!\!\![R(\boldsymbol{c}, \boldsymbol{z}^T)]. \label{eq:exact95traj}\tag{5}\] The policy gradient estimate of the rephrased return is given by \[\begin{align} \nabla_{\theta} & J(\boldsymbol{\theta}) \!=\! \sum_{\mathbf{z}} R(\boldsymbol{c}, \boldsymbol{z}^T) \frac{ \nabla_{\theta} \pi_{\boldsymbol{\theta}}(\mathbf{z} \!\mid\! \boldsymbol{c}) }{ \pi_{\boldsymbol{\theta}}(\mathbf{z} \!\mid\! \boldsymbol{c}) } \pi_{\boldsymbol{\theta}}(\mathbf{z} \!\mid\! \boldsymbol{c}) = \mathop{\mathrm{\mathbb{E}}}_{\mathbf{z} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c})}[ R(\boldsymbol{c}, \boldsymbol{z}^T) \nabla_{\theta} \log \pi_{\boldsymbol{\theta}}(\mathbf{z} \!\mid\! \boldsymbol{c}) ] \label{eq:pg95trajectory} \end{align}\tag{6}\] Unlike ELBO methods, trajectory-based methods compute an exact policy gradient estimate given we can compute \(\log \pi_{\boldsymbol{\theta}}(\mathbf{z} \!\mid\! \boldsymbol{c})\) for any sampled trajectory \(\mathbf{z}\): \[\begin{align} \pi_{\theta}(\mathbf{z} \!\mid\! \boldsymbol{c}) = \pi_{\theta}( \boldsymbol{z}^1, \ldots, \boldsymbol{z}^T \!\mid\! \boldsymbol{c}) = \prod_{t=1}^{T} \pi_{\theta}( \boldsymbol{z}^{t} \!\mid\! \boldsymbol{c}, \boldsymbol{z}^{\leq t-1}) = \prod_{t=1}^{T} \pi_{\theta}( \boldsymbol{z}^{t} \!\mid\! \boldsymbol{c}, \boldsymbol{z}^{t-1}), \label{eq:traj95likelihood} \end{align}\tag{7}\] where the last step follows since the diffusion process is Markovian. [15] approximate \({\log \pi_{\boldsymbol{\theta}}(\mathbf{z} \!\mid\! \boldsymbol{c}) \approx \sum_{t} \log\pi(\boldsymbol{z}^{t} | \boldsymbol{z}^{t-1}, c)}\), but ignore the remasking step.
We show that a proper treatment of this full probability estimate leads to a position term in the policy gradient estimate, which in practice greatly improves performance.
3 Method↩︎
We derive a policy gradient estimator for MDLMs by modeling the probability of trajectories using both token prediction and remasking in the trajectory likelihood. Denoising from sequence \(\boldsymbol{z}^{t-1}\) to \(\boldsymbol{z}^{t}\) can be viewed as a two-stage process: predicting mask tokens \(\hat{\boldsymbol{z}}^t \sim \pi_\theta(\cdot\mid\boldsymbol{z}^{t-1})\), and selecting positions to unmask \(\mathcal{U}_t\). A masking function \(\boldsymbol{z}^t = f_\text{remask}(\hat{\boldsymbol{z}}^t, \mathcal{U}_t \cup \boldsymbol{z}^{t-1} \ne \texttt{[MASK]})\) remasks all remaining tokens. 4 shows an overview.
The main issue with the above formulation is that the standard greedy top-\(K\) mask \(\mathcal{U}_t\) depends on \(\pi_\theta\), but is not differentiable, and thus a policy gradient estimate is not possible. We instead replace the greedy remasking with a probabilistic version both during generation, and reinforcement learning. This leads to a proper gradient estimate, and better empirical performance.
3.1 Probabilistic remasking↩︎
Let \(v^t_k = \log \pi(\hat{z}^t_k | \boldsymbol{z}^{t-1})\) be the log-likelihood of the sampled token \(\hat{z}^t_k\) at position \(k\) and step \(t\). In standard top-\(K\) remasking, the \(K\) positions with the highest scores \(v^t_k\) are deterministically selected for unmasking. This hard selection is not differentiable with respect to the model parameters \(\boldsymbol{\theta}\), preventing gradient-based optimization of the remasking. We replace this with a probabilistic variant: instead of selecting the top-\(K\) positions deterministically, we sample \(K\) positions without replacement from a categorical distribution whose probabilities are proportional to \(\exp(v^t_k / \tau)\) over all masked positions, where \(\tau > 0\) is a temperature parameter. As \(\tau \to 0\), this recovers the greedy top-\(K\) selection; for \(\tau > 0\), it defines a proper distribution over position subsets whose log-probability can be included in the trajectory likelihood and optimized via policy gradients. Specifically, we sample the unmasking set \(\mathcal{U}_t = \{u_1, u_2, \ldots, u_K\}\) sequentially without replacement, following the Plackett–Luce model [18], [19]: \({p_\text{unmask}(\mathcal{U}_t\mid \boldsymbol{z}^{t-1}, \hat{\boldsymbol{z}}^{t}) = \prod_{i=1}^K p_\text{unmask}(u_i=x | u_{<i}, \boldsymbol{z}^{t-1}, \hat{\boldsymbol{z}}^{t})}\) with \[\begin{align} p_\text{unmask}(u_i=x | u_{<i}, \boldsymbol{z}^{t-1}, \hat{\boldsymbol{z}}^{t}) = \frac{1_{[\boldsymbol{z}^{t-1}_{x}=\texttt{[MASK]}\wedge x \notin u_{<i}]} \exp(v^t_x / \tau)}{\sum_{j=1}^n 1_{[\boldsymbol{z}^{t-1}_{j}=\texttt{[MASK]}\wedge j \notin u_{<i}]} \exp(v^t_j / \tau)}, \label{eq:pos95prob} \end{align}\tag{8}\] where each distribution models a softmax over currently masked tokens. Masked tokens with a higher log-likelihood \(v^t_k\) have a higher chance of being unmasked. Tokens with a low log-likelihood will likely be remasked.
The remasking probability \(p_\text{remask}\) is then \[\begin{align} p_\text{remask}(\boldsymbol{z}^t | \mathcal{U}_t, \boldsymbol{z}^{t-1}, \hat{\boldsymbol{z}}^{t}) = 1_{[\boldsymbol{z}^t = f_\text{remask}(\hat{\boldsymbol{z}}^{t}, \mathcal{U}_t \cup \boldsymbol{z}^{t-1}=\texttt{[MASK]})]} \end{align}\] where \(f_\text{remask}\) is the remasking operator that applies the mask \(\mathcal{U}_t\) to the sequence \(\boldsymbol{z}^t\): \[\begin{align} f_\text{remask}(\hat{\boldsymbol{z}}^{t}, S)_k = \begin{cases}\hat{\boldsymbol{z}}^{t}_k&\text{if k \in S}\\\texttt{[MASK]}&\text{otherwise}\end{cases} \end{align}\] Unlike greedy top-\(K\) selection, this probabilistic formulation defines a proper distribution \(p(\mathcal{U}_t)\) over position subsets whose log-probability is differentiable with respect to the model parameters. The remasking operator \(f_\text{remask}\) and \(p_\text{remask}\) are deterministic and do not directly depend on the model parameters.
3.2 Reinforcement learning with probabilistic remasking↩︎
The diffusion process can be viewed as producing an extended sequence \({\hat{\boldsymbol{z}} = \{\boldsymbol{z}^0, \hat{\boldsymbol{z}}^1, \mathcal{U}^1, \boldsymbol{z}^1, \ldots, \boldsymbol{z}^T\}}\) over a distribution: \[\pi_{\boldsymbol{\theta}}(\hat{\boldsymbol{z}}|\boldsymbol{c}) = \prod_{t=1}^T \pi(\hat{\boldsymbol{z}}^t | \boldsymbol{z}^{t-1}, \boldsymbol{c}) p_\text{unmask}(\mathcal{U}_t | \hat{\boldsymbol{z}}^t, \boldsymbol{z}^{t-1}, \boldsymbol{c}) p_\text{remask}(\boldsymbol{z}^t | \mathcal{U}_t, \hat{\boldsymbol{z}}^t, \boldsymbol{z}^{t-1}, \boldsymbol{c}). \label{eq:pos95val95factorization}\tag{9}\] This interpretation extends the trajectory-based view of diffusion in [15], and further considers the remasking process as part of the generation. It leads to a similarly simple and direct policy gradient estimate over extended trajectories \(\hat{\boldsymbol{z}}\): \[\nabla_{\theta} J(\boldsymbol{\theta}) \!=\!\!\!\!\! \mathop{\mathrm{\mathbb{E}}}_{\hat{\boldsymbol{z}} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c})}\!\! \bigg[\! R(\boldsymbol{c}, \boldsymbol{z}^T)\!\sum_{t=1}^{T}\!\Big( \underbrace{ \nabla_{\theta} \log p_\text{unmask}(\mathcal{U}_t | \hat{\boldsymbol{z}}^t, \boldsymbol{z}^{t-1}, \boldsymbol{c}) }_{\text{unmasking gradient}} + \underbrace{ \nabla_{\theta} \log \pi_\theta(\hat{\boldsymbol{z}}^t \!\mid\! \boldsymbol{c}, \boldsymbol{z}^{t-1} ) }_{\text{token gradient}} \Big)\!\bigg]. \label{eq:pg95decomp}\tag{10}\] Here, the remasking is deterministic and does not depend on \(\boldsymbol{\theta}\) and thus falls out of the gradient estimate. Both terms in the decomposition (Eq. 10 ) are functions of the same policy \(\pi_{\boldsymbol{\theta}}\). Therefore, we can optimize them jointly with any standard policy gradient algorithm. We use GSPO [20] and find it most effective.
Finally, we note that the unmasking gradient is necessary for a correct gradient estimate: optimizing only the token component, as in [15], can miss directions that improve the objective \(J(\boldsymbol{\theta})\). Intuitively, changes to the unmasking order can affect the final reward even when they leave the token probabilities unchanged, so a token-only gradient provides no signal along such directions. We make this precise in Appendix 13 with a minimal counterexample: a finite-horizon MDLM problem with a linear policy in which the token-only gradient is identically zero along a direction \(\boldsymbol{v}\) that strictly increases the true objective, while the full position-aware gradient is not.
4 Experiments↩︎
| GSM8K | MATH500 | HumanEval | MBPP | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 2-4 (lr)5-7 (lr)8-10 (lr)11-13 Model / Seq Len | 128 | 256 | 512 | 128 | 256 | 512 | 128 | 256 | 512 | 128 | 256 | 512 |
| LLaDA-8B-Inst. | 69.5 | 77.2 | 79.8 | 28.2 | 32.4 | 34.6 | 27.4 | 35.5 | 37.8 | 36.2 | 41.2 | 40.4 |
| LLaDA-1.5 | 70.4 | 80.5 | 81.9 | 26.8 | 32.2 | 35.8 | — | — | — | — | — | — |
| d1 | 72.2 | 80.6 | 81.3 | 31.4 | 36.0 | 39.4 | 29.3 | 39.0 | 34.8 | 42.0 | 45.5 | 41.6 |
| wd1 | 74.6 | 81.5 | 83.0 | 31.0 | 37.4 | 39.0 | — | — | — | — | — | — |
| UniGRPO | 74.9 | 82.5 | 82.7 | 32.4 | 37.4 | 39.4 | — | — | — | — | — | — |
| GDPO | 78.4 | 82.8 | 85.0 | 33.2 | 39.6 | 41.4 | 26.2 | 39.6 | 39.0 | 43.6 | 50.6 | 47.1 |
| SPG w/ Mixture | 78.5 | 83.9\(^\dagger\) | 84.5 | 33.4 | 40.0 | 41.8 | 31.0 | 40.6 | 41.1 | 44.3 | 50.6 | 50.8 |
| StepMerge | 78.0 | 83.3 | 84.8 | 33.1 | 39.1 | 41.2 | 30.5 | 39.9 | 40.7 | 44.9 | 50.1 | 50.9 |
| Ours | ||||||||||||
| +2.5 | ||||||||||||
| +2.0 | ||||||||||||
| +2.1 | ||||||||||||
| +4.0 | ||||||||||||
| +2.2 | ||||||||||||
| +2.4 | ||||||||||||
| +2.2 | ||||||||||||
| +2.5 | ||||||||||||
| +2.9 | ||||||||||||
| +2.2 | ||||||||||||
| +2.2 | ||||||||||||
| +2.5 | ||||||||||||
Following prior work [12]–[14], [21], we use LLaDA-8B-Instruct [11] as the base model for RL fine-tuning. We employ Low-Rank Adaptation (LoRA) with rank \(r=128\), scaling factor \(\alpha=64\), and dropout \(0.05\), targeting all attention and MLP projection layers. We use 4-bit quantization (NF4) and Flash Attention 2 for memory efficiency, with all training in bfloat16.
4.0.0.1 Position log-prob computation.
We estimate both token and position log-probabilities using the StepMerge approximation [15], which groups \(T\) denoising steps into \(N\) contiguous segments and treats all tokens unmasked within a segment as if they were unmasked at the same step. See Appendix 12 for more details on the method and the theoretical analysis of approximation error induced by using StepMerge to evaluate the trajectory log-probability. We subsample \(K\) boundaries for evaluation. The unmasking score at each masked position is the maximum logit from the model’s output. Since we compute both log-probability terms from the same forward passes, the position term adds no overhead. For our main results, we set \(N=32\) and \(K=12\) (see Appendix 11). We optimize with GSPO [20], clipping the token and position importance ratios separately.
4.0.0.2 Datasets and evaluation.
We evaluate on GSM8K [22] and MATH500 [23] for mathematical reasoning, and HumanEval [24] and MBPP [25] for code generation. For math, we follow the train-test splitting, reward functions, and evaluation protocol of d1 [12], wd1 [21], and SPG [14]. For code, we follow GDPO [13], using KodCode-Light-RL-10K [26] as the training dataset. All evaluations are zero-shot. For RL rollouts, we use a sequence length of 256 tokens, 128 diffusion steps, and block size 32. We report accuracy at generation lengths 128, 256, and 512.
4.0.0.3 Baselines.
We compare against LLaDA-8B-Instruct [11] (base model), LLaDA-1.5 [27], d1 [12], wd1 [21], UniGRPO [28], GDPO [13], SPG w/ Mixture [14], and StepMerge [15]. Math results for prior methods are taken from their original papers. SPG does not report code generation results, so we evaluate their official code on HumanEval and MBPP ourselves. For StepMerge, we reimplement the algorithm with the same \(K\) subsampling as our method, making it a direct ablation that differs only in the absence of the masking term.
4.1 Experimental Results↩︎
4.1.0.1 Main results.
Table 1 presents our main results across four benchmarks. Our method achieves the best results in every configuration. The comparison against StepMerge is particularly informative: both methods share the same trajectory-based likelihood framework and \(K\) subsampling, differing only in the masking log-probability term. The consistent 2-4 point improvement across all configurations demonstrates that the masking gradient provides a meaningful training signal beyond what token-level optimization alone captures. Against evidence bound methods such as GDPO and SPG, our method also shows consistent gains, particularly on MATH500 (+4.0 at length 128, +2.4 at length 512) and code generation benchmarks. We further show that our method generalizes to a second base model, Dream-7B (Appendix 9), and to two planning tasks, Sudoku and Countdown (Appendix 8).
4.1.0.2 Training efficiency.
We compare the training efficiency of our method against SPG [14], the strongest ELBO-based baseline in our setting, on GSM8K. Both methods run on the same hardware (8 H100 GPUs) from the same base model, with identical generation, optimization, and evaluation configurations; the only difference is the log-probability estimator.
We compare the wall-clock time to reach fixed accuracy targets both for our method and SPG in Table [tbl:tab:eff95speedup]. Our method reaches \(74\%\) accuracy \(1.9\times\) faster than SPG and reaches SPG’s final accuracy in roughly 15 hours, compared with 18 hours for SPG. It also reaches \(79\%\) and \(80\%\) accuracy, which SPG does not reach. Table [tbl:tab:eff95summary] gives the efficiency summary: our per-step throughput is lower because the trajectory-based estimator requires more forward passes per optimizer step, but the position-aware gradient converges in fewer steps and reaches a higher final accuracy (81.0 vs.). Figure 5 plots accuracy against wall-clock time: our method starts marginally behind SPG, overtakes it within the first few hours, and the gap widens as training proceeds. Across three random seeds our final accuracy varies by at most \(0.2\) points, so these differences are well outside run-to-run variation.
0.44
| SPG | Ours | |
|---|---|---|
| Throughput (steps/h) | 360 | 290 |
| Convergence (steps) | 6500 | 6000 |
| Final accuracy (%) | 78.5 | 81.0 |
| Match SPG-final (h) | 18 | 15 |
0.52
| Target | SPG (h) | Ours (h) | Speedup |
|---|---|---|---|
| 74% | 8.6 | 4.5 | 1.9\(\times\) |
| 75% | 9.9 | 5.8 | 1.7\(\times\) |
| 76% | 13.4 | 10.0 | 1.3\(\times\) |
| 77% | 15.9 | 11.8 | 1.3\(\times\) |
| 78% | 17.2 | 14.0 | 1.2\(\times\) |
| 79% | — | 16.5 | — |
| 80% | — | 18.3 | — |
![Figure 5: Wall-clock training efficiency on GSM8K (8 H100 GPUs, generation length 128). Test accuracy versus wall-clock training time for our method and SPG [14]. Although our method runs fewer steps per hour, it overtakes SPG within the first few hours and converges to a higher accuracy (81.0 vs.). Stars mark where each method first reaches SPG’s final accuracy: \approx15 hours for ours versus \approx18 hours for SPG.](/article/2607.15200/figures/efficiency_convergence.png "fig:")
4.1.0.3 Inference block size.
Figure 6 compares inference block sizes (32, 64, and full sequence) at generation length 256 with confidence-based unmasking. We compare against SPG [14], a strong evidence-bound baseline, and StepMerge [15], which shares the same trajectory-based framework but without the position term. All methods are trained with block size 32. Our method outperforms both baselines at every block size, and the gap widens as block size increases. Against SPG, the margin grows from +2.0 to +4.3 on GSM8K and from +2.2 to +3.5 on MATH500 (block 32 to full sequence). Against StepMerge, the margin grows from +2.6 to +4.4 on GSM8K and from +3.1 to +3.7 on MATH500. We attribute this to the growing importance of position selection as the diffusion window expands: larger blocks present more masked positions to choose from at each step, increasing the variability in unmasking order and amplifying the benefit of mask-aware optimization. Full numerical results are provided in Table 5 in the appendix. We provide further ablations on the RL algorithm, the StepMerge configuration, the sampling temperatures, and the unmasking-set size \(|\mathcal{U}_t|\) in Appendix 11.
![Figure 6: Ablation over inference block size. All methods use confidence-based unmasking at generation length 256. Results for baselines are from SPG [14].](/article/2607.15200/figures/blocksize_ablation.png "fig:")
5 Related work↩︎
5.0.0.1 Diffusion Language Models.
Discrete diffusion models for text generation have developed along several lines. Early work applied diffusion in continuous embedding space [29], [30], while more recent approaches operate directly in discrete token space. Score Entropy Discrete Diffusion (SEDD) [31] and Masked Diffusion Language Models (MDLM) [5]–[7] established masked diffusion as a competitive paradigm for text generation, with simplified training objectives based on the ELBO. Scaling these models to larger sizes, LLaDA [11] and Dream [32] demonstrated that masked diffusion LLMs can approach autoregressive performance on standard benchmarks. More recently, diffusion language models have been extended to multimodal settings [28], [33]–[35], and several works have improved inference efficiency through caching [36], [37] and parallel decoding strategies [38], [39].
5.0.0.2 Reinforcement Learning for Language Models.
Reinforcement learning has become a central technique for post-training large language models, building on classical policy gradient methods [40], [41]. Early work used RL from human feedback (RLHF) with PPO [42] and TRPO [43] to align models with human preferences [4], [44]. More recently, RL has been applied to improve reasoning capabilities. DeepSeekMath [45] introduced Group Relative Policy Optimization (GRPO), a critic-free policy gradient method that estimates baselines from group rollouts, removing the need for a value network. DeepSeek-R1 [1] and Kimi K1.5 [16] scaled RL-based reasoning training, demonstrating that policy gradient methods with verifiable rewards can elicit strong mathematical and coding performance. Subsequent work has refined the GRPO objective, reducing bias [46], improving sample efficiency [47], and introducing sequence-level clipping [20].
5.0.0.3 Reinforcement Learning for Diffusion Language Models.
Reinforcement learning for MDLMs has been used for reward optimization [48], preference alignment [27], faster parallel decoding [49]–[52], and reasoning [12]–[15]. For reasoning, the central challenge of RL is estimating the intractable sequence log-likelihood. ELBO-based methods estimate log-probabilities via variants of the training objective: d1 [12] uses a mean-field approximation, UniGRPO [28] and ESPO [53] use sequence-level Monte Carlo estimators, GDPO [13] tightens the bound with Gaussian quadrature, wd1 [54] uses weighted optimization, and SPG [14] combines ELBO and EUBO bounds with block-wise masking. Trajectory-based methods instead decompose the log-likelihood over denoising steps, attempting to follow the generation trajectory: d2 [15] and [55] partition the trajectory into segments via the StepMerge estimator. However, these trajectory-based methods model the trajectory as a sequence of token prediction steps only, and in doing so ignore the position selection decision that MDLMs make at generation time. Our work addresses this gap.
5.0.0.4 Learning the unmasking order.
A line of work aims to learn or optimize the unmasking order of MDLMs, frequently to accelerate parallel decoding by reducing the number of forward passes. Jazbec et al. [50] learn an unmasking policy over the model’s confidences, dUltra [51] trains an unmasking planner head with reinforcement learning, and P2 [56] frames sampling as path planning to choose which positions to update; Where-to-Unmask [57] instead trains a separate planner in a supervised fashion to imitate an oracle order derived from ground-truth tokens. Unlike these methods, our goal is not to introduce an explicit model of the unmasking order; rather, we show that a position-selection term arises directly from the complete policy-gradient estimator, and we optimize it using the model’s own logits without any additional module. The most closely related method in this respect is LLaDOU/DCoLT [58], which also incorporates position selection into the RL policy through a Plackett–Luce distribution; unlike DCoLT, which trains a separate module to model the position logits, we reuse the base model’s own logits. We provide a detailed comparison with DCoLT in Appendix 10.
6 Conclusion↩︎
We propose a policy gradient method for masked diffusion language models that optimizes both token prediction and position selection. By replacing greedy top-\(K\) remasking with a probabilistic variant, we obtain a differentiable distribution over unmasking positions, yielding a gradient decomposition into a token term and a masking term. Extensive experiments across mathematical reasoning and code generation benchmarks show that jointly optimizing these terms in the policy gradient achieves significant improvements over existing methods and state-of-the-art results.
Acknowledgments↩︎
The experiments in this work were run on the Vista GPU Cluster through the Center for Generative AI (CGAI) and the Texas Advanced Computing Center (TACC) at The University of Texas at Austin.
7 Training details↩︎
7.0.0.1 Base model.
We use LLaDA-8B-Instruct [11] as the base model for all experiments. LLaDA is a masked diffusion language model that uses absorbing-state diffusion with a
special [MASK] token (ID 126336). We do not apply supervised fine-tuning before RL training; all results use the pretrained instruction-tuned checkpoint directly.
7.0.0.2 Parameter-efficient fine-tuning.
We use Low-Rank Adaptation (LoRA) for parameter-efficient fine-tuning. We set the LoRA rank to \(r=128\), scaling factor \(\alpha=64\), and dropout \(0.05\). LoRA adapters are applied to all attention projection layers (query, key, value, and output) and all MLP projection layers (gate, up, and down). The base model weights are frozen in 4-bit quantization (NF4 format) to reduce memory usage, with all LoRA parameters and computations in bfloat16. We use Flash Attention 2 [59] for memory-efficient attention computation.
7.0.0.3 Optimizer.
We use AdamW [60] with \(\beta_1 = 0.9\), \(\beta_2 = 0.99\), and weight decay \(0.1\). The learning rate is set to \(3 \times 10^{-6}\) with a constant schedule and a brief warmup (warmup ratio \(0.0001\)). We clip gradients to a maximum norm of \(0.2\).
7.0.0.4 Batch size and gradient accumulation.
We use a per-device training batch size of 3 with gradient accumulation over 2 steps. Training is distributed across 8 GPUs, yielding an effective batch size of \(3 \times 2 \times 8 = 48\) prompt-completion pairs per optimization step. For each prompt, we sample \(G=6\) completions for group-relative advantage estimation.
7.0.0.5 Policy optimization.
We optimize using GSPO [20] with separate clipping for the token and position importance ratios. GSPO clips the importance ratio with a lower bound of \(3 \times 10^{-4}\) and an upper bound of \(4 \times 10^{-4}\). No KL regularization term is used. We perform \(\mu = 12\) inner gradient updates per batch of generated completions. We synchronize the reference model every 64 optimization steps.
7.0.0.6 Generation parameters.
During RL rollouts, we generate completions with a maximum length of 256 tokens. Generation uses semi-autoregressive block-wise decoding with a block size of 32 tokens and 128 diffusion steps per block. At each denoising step, we use probabilistic remasking. The sampling temperature is set to 0.9 for token selection and 0.5 for remasking.
7.0.0.7 Datasets.
For mathematical reasoning, we train on the training split of GSM8K [22] for grade-school math and on the training split of the MATH dataset for competition-level math, evaluating on the MATH500 [23] test set. We follow the same train-test splitting as d1 [12] and SPG [14]. For code generation, we train on KodCode-Light-RL-10K [26], a dataset of coding problems at varying difficulty levels with synthetic unit tests for reward computation. We evaluate on HumanEval [24] (164 hand-crafted Python problems) and sanitized MBPP [25] (257 crowd-sourced Python tasks). We follow the same dataset setup as GDPO [13].
7.0.0.8 Reward functions.
We follow the reward functions of d1 [12] and SPG [14] for mathematical reasoning, and GDPO [13] for code generation.
7.0.0.8.1 GSM8K.
We use a composite reward with five additive components:
XML Structure Reward: \(+0.125\) per correctly placed opening or closing tag.
Soft Format Reward: \(+0.5\) if the output matches the pattern
<reasoning>...</reasoning><answer>...</answer>.Strict Format Reward: \(+0.5\) for exact formatting with correct line breaks.
Integer Answer Reward: \(+0.5\) if the extracted answer is a valid integer.
Correctness Reward: \(+2.0\) if the extracted answer matches the ground truth.
7.0.0.8.2 MATH500.
We use a two-component reward:
Format Reward: \(+1.0\) if answer tags are present with
\boxed{}inside; \(+0.75\) if answer tags are present without\boxed{}; \(+0.5\) if\boxed{}is present without answer tags; \(+0.25\) otherwise.Correctness Reward: \(+2.0\) if the answer in
\boxed{}matches the ground truth.
7.0.0.8.3 Code generation (HumanEval, MBPP).
We use a two-component reward:
Format Reward: \(+1.0\) if the output contains a valid Python code block with no syntax errors; \(+0.5\) if well-formatted but with syntax errors; \(+0.0\) otherwise.
Code Execution Reward: the fraction of unit tests passed. We block unsafe modules (os, sys, shutil, subprocess, socket, psutil, ctypes, pathlib, builtins, __import__) and assign a reward of \(0\) if any are used.
7.0.0.9 Training schedule.
We train each dataset independently. For GSM8K, we train for approximately 8,000 optimization steps. For MATH500, we train for approximately 8,000 steps. For code generation (KodCode), we train for approximately 6,000 steps. We save checkpoints every 100 steps and evaluate at each checkpoint.
All experiments are conducted on 8 NVIDIA H100-80GB GPUs with DeepSpeed ZeRO-2.
Figure 7: Training reward dynamics for our method on GSM8K, MATH500, and KodCode. Curves show EMA-smoothed reward with one standard deviation band..
8 Results on planning tasks↩︎
We additionally evaluate on two planning benchmarks, Sudoku [61] and Countdown [62], following the evaluation protocol of SPG [14] and d1 [12]. Following SPG, Sudoku uses 3 few-shot examples during rollouts and evaluation. Table 2 reports accuracy at generation lengths 128, 256, and 512; results for prior methods are taken from SPG [14]. Our method achieves the best results on both tasks at every generation length, outperforming the strongest baseline, SPG, with the largest gain on Sudoku (+3.1 at length 512). This is consistent with our finding that the benefit of position-aware optimization grows as the number of masked positions increases.
| Sudoku | Countdown | |||||
|---|---|---|---|---|---|---|
| 2-4 (lr)5-7 Model / Seq Len | 128 | 256 | 512 | 128 | 256 | 512 |
| LLaDA-8B-Inst. | 5.7 | 27.7 | 26.2 | 18.8 | 16.8 | 16.8 |
| LLaDA-1.5 | 7.4 | 26.9 | 29.0 | 21.9 | 21.1 | 21.5 |
| d1 | 7.2 | 32.5 | 29.3 | 30.9 | 30.9 | 34.4 |
| wd1 | 33.1 | 32.1 | 22.5 | 48.8 | 52.3 | 50.8 |
| UniGRPO | 59.0 | 67.0 | 62.9 | 44.5 | 43.0 | 57.0 |
| SPG w/ Mixture | 82.9 | 94.0 | 93.1 | 68.8 | 70.7 | 70.3 |
| Ours | 84.5 | 94.9 | 96.2 | 69.3 | 71.4 | 71.0 |
9 Generalization to Dream-7B↩︎
We additionally evaluate our method on a different base model, Dream-7B [32]. Table 3 reports GSM8K, MATH500, and Sudoku accuracy at generation lengths 128, 256, and 512; SPG results are obtained by running their publicly available code. Our method outperforms all baselines across all three benchmarks and generation lengths, with improvements of up to \(2.9\) points over the strongest baseline.
| GSM8K | MATH500 | Sudoku | |||||||
|---|---|---|---|---|---|---|---|---|---|
| 2-4 (lr)5-7 (lr)8-10 Model / Seq Len | 128 | 256 | 512 | 128 | 256 | 512 | 128 | 256 | 512 |
| Dream-7B | 75.8 | 81.3 | 80.7 | 38.2 | 45.7 | 48.0 | 9.3 | 2.1 | 14.0 |
| d1 | 77.0 | 81.9 | 81.7 | 39.4 | 46.9 | 48.9 | 64.4 | 69.7 | 51.1 |
| wd1 | 76.3 | 82.4 | 82.9 | 39.5 | 47.4 | 50.5 | 29.5 | 39.2 | 30.3 |
| ESPO | 79.6 | 82.3 | 82.0 | 40.3 | 47.4 | 50.3 | 71.7 | 72.3 | 71.3 |
| SPG w/ Mixture | 80.1 | 82.6 | 83.5 | 41.1 | 48.5 | 50.7 | 72.1 | 72.4 | 71.9 |
| Ours | 82.5 | 84.1 | 86.4 | 43.8 | 50.3 | 52.9 | 73.4 | 74.1 | 74.3 |
10 Comparison to LLaDOU↩︎
LLaDOU [58], trained with the Diffusion Chain of Lateral Thought (DCoLT) algorithm, shares our high-level insight that the unmasking order should be part of the policy, and arrives at a similar probabilistic framework based on Plackett–Luce position sampling. The key difference is architectural: LLaDOU introduces a separately trained head for position selection, whereas our method derives position scores directly from the model’s own token log-likelihoods, requiring no additional parameters or architectural changes. This has the practical advantage that our approach integrates directly with existing MDLM training infrastructure such as StepMerge [15], and applies to any off-the-shelf MDLM.
We compare the two approaches on GSM8K in Table 4, reporting accuracy, training throughput, and total training cost. When DCoLT uses full trajectory back-propagation (full replay), it reaches accuracy comparable to ours but at roughly \(6\times\) lower throughput and \(5\times\) higher total GPU cost. In a matched comparison where both methods use the StepMerge approximation, our method outperforms DCoLT by approximately \(5\%\) on average across generation lengths. This efficiency gain does not come solely from StepMerge: even with StepMerge, DCoLT still requires about \(5\times\) more GPU hours to train (\(\sim\)800 vs \(\sim\)160), because it must train a separate position-selection head from scratch. Our method instead derives position scores from the model’s existing token log-likelihoods, providing a useful signal from the start of training and converging much faster. Overall, our method matches or exceeds DCoLT’s accuracy while being substantially cheaper to train and applicable to any MDLM without architectural changes.
| Method | 128 | 256 | 512 | Throughput (steps/h) | GPU hours |
|---|---|---|---|---|---|
| DCoLT (full replay) | 80.4 | 84.3 | 86.2 | 48 | \(\sim\)800 |
| DCoLT (w/ StepMerge) | 77.1 | 80.1 | 82.3 | 280 | \(\sim\)800 |
| Ours | 81.0 | 85.9 | 87.1 | 290 | \(\sim\)160 |
11 Additional ablations↩︎
11.1 Inference block size↩︎
Table 5 reports the full numerical results for the inference block size ablation discussed in Section 4.1 and summarized in Figure 6.
| Block 32 | Block 64 | Full Sequence | ||||
|---|---|---|---|---|---|---|
| 2-3 (lr)4-5 (lr)6-7 Model | GSM8K | MATH500 | GSM8K | MATH500 | GSM8K | MATH500 |
| LLaDA-8B-Inst. | 77.2 | 32.4 | 78.6 | 33.2 | 23.9 | 17.8 |
| LLaDA-1.5 | 80.5 | 32.2 | 81.0 | 35.4 | 41.4 | 20.4 |
| d1 | 80.6 | 36.0 | 80.9 | 37.6 | 57.5 | 22.6 |
| wd1 | 81.5 | 37.4 | 82.5 | 37.4 | 56.7 | 25.0 |
| UniGRPO | 82.5 | 37.4 | 82.3 | 37.4 | 50.0 | 24.2 |
| SPG w/ Mixture | 83.9 | 40.0 | 84.1 | 41.1 | 58.8 | 24.0 |
| StepMerge | 83.3 | 39.1 | 83.7 | 40.2 | 58.7 | 23.8 |
| Ours | 85.9 | 42.2 | 86.3 | 43.6 | 63.1 | 27.5 |
11.2 RL algorithm↩︎
Our decomposition is compatible with any policy gradient algorithm. We compare three widely used algorithms, RLOO [63], GRPO, and GSPO, in Table ¿tbl:tab:rl95ablation? with \(N=32\) and \(K=12\), evaluating on GSM8K and MATH500 at generation length 128. GSPO performs the strongest on both benchmarks, which we attribute to its sequence-level clipping being better suited to handle the variance introduced by subsampling \(K\) boundaries.
0.36
| Algorithm | GSM8K | MATH500 |
|---|---|---|
| RLOO | 75.0 | 30.7 |
| GRPO | 79.5 | 36.4 |
| GSPO | 81.0 | 37.4 |
0.56
| N | K | GSM8K | MATH500 |
|---|---|---|---|
| 16 | 8 | 79.9 | 36.8 |
| 16 | 80.6 | 37.2 | |
| 32 | 1 | 78.0 | 35.4 |
| 4 | 79.4 | 36.3 | |
| 8 | 80.3 | 36.9 | |
| 12 | 81.0 | 37.4 | |
| 16 | 81.3 | 37.4 | |
| 64 | 8 | 79.8 | 36.8 |
| 16 | 80.2 | 37.0 |
11.3 StepMerge configuration↩︎
Table [tbl:tab:nk95ablation] ablates the number of segments \(N\) and subsampled boundaries \(K\), evaluating on GSM8K and MATH500 at generation length 128. Performance improves with larger \(N\), as finer segmentation reduces the approximation error. Performance degrades only modestly as \(K\) decreases (at \(N=32\), from 81.0 at \(K=12\) to 78.0 at \(K=1\) on GSM8K), indicating the estimator remains effective even with few sampled boundaries. However, increasing \(K\) beyond 12 provides only marginal gains: at \(N=32\), \(K=12\) matches \(K=16\) on MATH500 while requiring fewer forward passes. We use \(N=32\), \(K=12\) for our main results.
11.4 Sampling temperature↩︎
Our method uses two temperatures: a token sampling temperature \(\tau_{\mathrm{tok}}\) and a position selection temperature \(\tau_{\mathrm{pos}}\) from our probabilistic position selection (Eq. 8 ), set to \(\tau_{\mathrm{tok}}=0.9\) (matching prior work) and \(\tau_{\mathrm{pos}}=0.5\) by default. We ablate both in Table ¿tbl:tab:temp95ablation? on GSM8K and MATH500 at generation length 128. Performance degrades as \(\tau_{\mathrm{tok}}\) is lowered, since less diverse rollouts weaken the group-relative advantage signal, and as \(\tau_{\mathrm{pos}}\) is raised, since more random selection unmasks low-confidence positions. The defaults strike a good balance, with results remaining reasonably stable across the tested range.
0.46
| \(\bm{\tau_{\mathrm{tok}}}\) | GSM8K | MATH500 |
|---|---|---|
| 0.5 | 79.1 | 35.9 |
| 0.9 | 81.0 | 37.4 |
| 1.5 | 80.5 | 37.0 |
0.46
| \(\bm{\tau_{\mathrm{pos}}}\) | GSM8K | MATH500 |
|---|---|---|
| 0.1 | 80.7 | 37.3 |
| 0.5 | 81.0 | 37.4 |
| 1.0 | 79.0 | 35.7 |
11.5 Unmasking-set size↩︎
At each denoising step, our method selects a set \(\mathcal{U}_t\) of positions to unmask; for all main results we use \(|\mathcal{U}_t|=2\) per step, matching d1 [12], wd1 [21], and SPG [14]. The selection distribution \(p_{\mathrm{unmask}}(\mathcal{U}_t)\) in Eq. 8 is defined for any \(|\mathcal{U}_t|\), and Table 6 varies it from 2 to 8 on GSM8K with all other settings fixed. Our method remains robust across this range, degrading by at most 1.0 point from \(|\mathcal{U}_t|=2\) to \(8\). This modest drop is expected, since larger sets commit more positions per step and yield lower-quality rollouts during training.
| \(\bm{|\mathcal{U}_t|}\) | Gen Len 128 | Gen Len 256 | Gen Len 512 |
|---|---|---|---|
| 2 | 81.0 | 85.9 | 87.1 |
| 4 | 80.4 | 85.1 | 86.8 |
| 8 | 80.1 | 84.9 | 86.5 |
12 StepMerge approximation↩︎
In our experiments, we approximate the token component of the policy gradient objective using the conditional distribution over each unmasked token, \(\pi_{\theta}(\boldsymbol{z}_i^t \mid \boldsymbol{c}, \boldsymbol{z}^{t-1})\). Specifically, following [15], we approximate the log-likelihood term in the token gradient by \(\sum_{i \in \mathcal{U}_t} \log \pi_{\theta}(\boldsymbol{z}_i^t \mid \boldsymbol{c}, \boldsymbol{z}^{t-1})\).
Computing trajectory log-probability in the policy gradient using above approximation requires evaluating all \(T\) denoising transitions, which is computationally expensive. Following [15], we adopt the StepMerge approximation, which partitions the trajectory into \(N\) segments with \(N \ll T\). Let \({t_0=0 < t_1 < \cdots < t_N = T}\) denote segment boundaries. Instead of evaluating every transition, we approximate the trajectory likelihood using only the segment endpoints \(\log \pi_{\theta}(\boldsymbol{z}^{t_i} \mid \boldsymbol{z}^{t_{i-1}}, \boldsymbol{c})\). Intuitively, this approximation treats multiple denoising steps as a single macro-transition, reducing the number of required forward passes from \(T\) to \(N\) while preserving the overall trajectory structure.
12.1 StepMerge approximation error↩︎
We analyze the approximation error induced by using StepMerge to evaluate the trajectory log-probability. The error has three sources: within-segment unmasking timing information lost by the macro-stage approximation, token-probability error from using the segment boundary context, and unmasking-probability error from using the same boundary context.
Proposition 1 (Macro-StepMerge approximation error). Let the \(T\) denoising steps be partitioned into \(N\) equal-length StepMerge segments with boundaries \(0=t_0<t_1<\cdots<t_N=T\). Let \(I_m=\{t_{m-1}+1,\ldots,t_m\}\) be segment \(m\).
Suppose that, for every segment \(m\), every \(t\in I_m\), and every retained position \(i\in\mathcal{U}_t\), \[\log \frac{ \pi_{\boldsymbol{\theta}}(z_i^t\mid \boldsymbol{z}^{t-1},\boldsymbol{c}) }{ \pi_{\boldsymbol{\theta}}(z_i^t\mid \boldsymbol{z}^{t_{m-1}},\boldsymbol{c}) } \leq \epsilon_{\mathrm{tok}}. \label{eq:macro95tok95stability}\qquad{(1)}\] Suppose also that the unmasking distribution is stable within each segment: \[\log \frac{ p_{\mathrm{unmask}}(\mathcal{U}_t\mid \boldsymbol{z}^{t-1},\boldsymbol{c}) }{ p_{\mathrm{unmask}}(\mathcal{U}_t\mid \boldsymbol{z}^{t_{m-1}},\boldsymbol{c}) } \leq |\mathcal{U}_t|\epsilon_{\mathrm{unmask}}. \label{eq:macro95unmask95stability}\qquad{(2)}\] Then the macro-StepMerge approximation error satisfies \[D_N \leq n\log\!\left(\frac{T}{N}+1\right) + n\epsilon_{\mathrm{tok}} + n\epsilon_{\mathrm{unmask}}. \label{eq:macro95stepmerge95bound}\qquad{(3)}\]
Proof. Let the full length-\(T\) trajectory be \(\xi = (\boldsymbol{z}^0,\hat{\boldsymbol{z}}^1,\mathcal{U}_1,\boldsymbol{z}^1,\ldots,\boldsymbol{z}^{T-1},\hat{\boldsymbol{z}}^T,\mathcal{U}_T,\boldsymbol{z}^T).\) For segment \(m\), define the segment-level unmasking set \(\bar{\mathcal{U}}^{\,m} = \bigcup_{t\in I_m}\mathcal{U}_t.\) The macro-StepMerge block trajectory for the segment keeps only \(\xi_{\mathrm{SM}}^m = \left(\boldsymbol{z}^{t_{m-1}},\bar{\mathcal{U}}^{\,m},\boldsymbol{z}^{t_m}\right).\) It records the start state, the end state, and the positions unmasked somewhere in the segment, but not the exact within-segment step at which those positions were unmasked.
Let the skipped intermediate information be \(H^m = \left( \boldsymbol{z}^{t_{m-1}+1},\hat{\boldsymbol{z}}^{t_{m-1}+1}, \ldots, \boldsymbol{z}^{t_m-1},\hat{\boldsymbol{z}}^{t_m} \right),\) and let the within-segment timing variable be \(\Lambda^m = \left( \mathcal{U}_{t_{m-1}+1}, \mathcal{U}_{t_{m-1}+2}, \ldots, \mathcal{U}_{t_m} \right).\) Thus the full block can be represented as \((\xi_{\mathrm{SM}}^m,H^m,\Lambda^m)\).
Let \(\mathcal{F}_{m-1}\) denote the history up to the start of block \(m\). By the chain rule of probability, the exact block conditional distribution is \[\begin{align} p_m(\xi_{\mathrm{SM}}^m,H^m,\Lambda^m\mid \mathcal{F}_{m-1}) &= p_m(\xi_{\mathrm{SM}}^m\mid \mathcal{F}_{m-1}) p_m(H^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) \notag\\ &\quad \cdot p_m(\Lambda^m\mid H^m,\xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}). \label{eq:exact95block95factor} \end{align}\tag{11}\]
The StepMerge block distribution on the same variables is given by: \[\begin{align} q_m(\xi_{\mathrm{SM}}^m,H^m,\Lambda^m\mid \mathcal{F}_{m-1}) &= q_m(\xi_{\mathrm{SM}}^m\mid \mathcal{F}_{m-1}) p_m(H^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) \notag\\ &\quad \cdot p_m(\Lambda^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}). \label{eq:sm95block95factor} \end{align}\tag{12}\] The timing distribution in Eq. 12 does not condition on \(H^m\), because StepMerge only knows the coarsened block trajectory \(\xi_{\mathrm{SM}}^m\) when assigning within-segment timing and does not depend on \(H^m\).
The full and StepMerge approximation trajectory distributions factor over blocks: \[p(\xi) = \prod_{m=1}^N p_m(\xi_{\mathrm{SM}}^m,H^m,\Lambda^m\mid \mathcal{F}_{m-1}), \qquad q(\xi) = \prod_{m=1}^N q_m(\xi_{\mathrm{SM}}^m,H^m,\Lambda^m\mid \mathcal{F}_{m-1}).\] Therefore, the KL divergence between full and StepMerge trajectories is \[\begin{align} D_N &:= D_{\mathrm{KL}}(p(\xi)\,\|\,q(\xi)) \notag\\ &= \sum_{m=1}^N \mathbb{E}_{\mathcal{F}_{m-1}\sim p} \left[ D_{\mathrm{KL}}\!\left( p_m(\cdot\mid\mathcal{F}_{m-1}) \,\middle\|\, q_m(\cdot\mid\mathcal{F}_{m-1}) \right) \right]. \label{eq:block95kl95decomp} \end{align}\tag{13}\]
Using Eqs. 11 and 12 , the \(H^m\) terms cancel in the block log-ratio: \[\begin{align} \log \frac{p_m}{q_m} &= \log \frac{ p_m(\Lambda^m\mid H^m,\xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) }{ p_m(\Lambda^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) } + \log \frac{ p_m(\xi_{\mathrm{SM}}^m\mid \mathcal{F}_{m-1}) }{ q_m(\xi_{\mathrm{SM}}^m\mid \mathcal{F}_{m-1}) }. \end{align}\] Taking expectations gives \[D_N = D_{\mathrm{time}} + D_{\mathrm{coarse}},\] where \[D_{\mathrm{time}} = \sum_{m=1}^N I\!\left( \Lambda^m;H^m \mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1} \right), \label{eq:dtime95cmi}\tag{14}\] and \(D_{\mathrm{coarse}}\) is the remaining KL between exact and StepMerge macro-block likelihoods. We first bound the timing term. By the conditional mutual information identity, \[\begin{align} I(\Lambda^m;H^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) & = H(\Lambda^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) - H(\Lambda^m\mid H^m,\xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}). \\ & \leq H(\Lambda^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}), \end{align}\] where the inequality follows from the nonnegativity of conditional entropy. For each \(i\in \bar{\mathcal{U}}^{\,m}\), let \(R_i^m\) denote the exact within-segment unmasking time of position \(i\). Since \(|I_m|=T/N\), there are at most \(T/N+1\) possible timing choices under the macro-stage endpoint convention. Therefore, \[H(R_i^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) \leq \log\!\left(\frac{T}{N}+1\right).\] The variable \(\Lambda^m\) is determined by the collection of per-position timing variables, so by subadditivity of entropy, \[H(\Lambda^m\mid \xi_{\mathrm{SM}}^m,\mathcal{F}_{m-1}) \leq |\bar{\mathcal{U}}^{\,m}|\log\!\left(\frac{T}{N}+1\right).\] Therefore, \[D_{\mathrm{time}} \leq \sum_{m=1}^N |\bar{\mathcal{U}}^{\,m}|\log\!\left(\frac{T}{N}+1\right) = n\log\!\left(\frac{T}{N}+1\right), \label{eq:dtime95bound95macro}\tag{15}\] because every final token position is unmasked exactly once.
We now bound the coarse term. Marginalization cannot increase KL divergence so we have: \(D_{\mathrm{KL}}(p(X,Y)\|q(X,Y)) \geq D_{\mathrm{KL}}(p(X)\|q(X))\) for random variables \(X\) and \(Y\). Applying the KL inequality between \(\xi_{\mathrm{SM}}^m\) and \((\mathcal{U}_{t_{m-1}+1},\boldsymbol{z}^{t_{m-1}+1},\ldots,\mathcal{U}_{t_m})\), we have \(\xi_{\mathrm{SM}}^m\) and \(Y^m\) gives \[\begin{align} D_{\mathrm{coarse}} &= \sum_{m=1}^N \mathop{\mathrm{\mathbb{E}}}_{\mathcal{F}_{m-1}\sim p} \left[ D_{\mathrm{KL}}\!\left( p_m(\xi_{\mathrm{SM}}^m\mid \mathcal{F}_{m-1}) \,\middle\|\, q_m(\xi_{\mathrm{SM}}^m\mid \mathcal{F}_{m-1}) \right) \right] \\ &\leq \sum_{m=1}^N \mathop{\mathrm{\mathbb{E}}}_{\mathcal{F}_{m-1}\sim p} \Big[ D_{\mathrm{KL}}\!\Big( p_m( (\mathcal{U}_{t_{m-1}+1},\boldsymbol{z}^{t_{m-1}+1},\ldots,\mathcal{U}_{t_m},\boldsymbol{z}^{t_m}) \mid \mathcal{F}_{m-1}) \,\Big\|\, \\ & \qquad \qquad \qquad \qquad \qquad q_m( (\mathcal{U}_{t_{m-1}+1},\boldsymbol{z}^{t_{m-1}+1},\ldots,\mathcal{U}_{t_m},\boldsymbol{z}^{t_m}) \mid \mathcal{F}_{m-1}) \Big) \Big]. \end{align}\]
For a realized segment, the deterministic remasking factors cancel. The exact-to-StepMerge density ratio therefore decomposes into token and unmasking likelihood ratios: \[\begin{align} \log \frac{p_m(\cdot \mid \mathcal{F}_{m-1})}{q_m(\cdot \mid \mathcal{F}_{m-1})} &= \sum_{t \in I_m} \sum_{i \in \mathcal{U}_t} \log \frac{ \pi_{\boldsymbol{\theta}}(z_i^t \mid \boldsymbol{z}^{t-1}, \boldsymbol{c}) }{ \pi_{\boldsymbol{\theta}}(z_i^t \mid \boldsymbol{z}^{t_{m-1}}, \boldsymbol{c}) } \\ &\quad + \sum_{t \in I_m} \log \frac{ p_{\mathrm{unmask}}(\mathcal{U}_t \mid \boldsymbol{z}^{t-1},\boldsymbol{c}) }{ p_{\mathrm{unmask}}(\mathcal{U}_t \mid \boldsymbol{z}^{t_{m-1}},\boldsymbol{c}) } \\ &\leq \sum_{t \in I_m} |\mathcal{U}_t| \left(\epsilon_{\mathrm{tok}}+\epsilon_{\mathrm{unmask}}\right), \end{align}\] where the last inequality follows from the two stability assumptions. Since each generated token is unmasked exactly once, \(\sum_{m=1}^N \sum_{t \in I_m} |\mathcal{U}_t| = n\). Hence \[D_{\mathrm{coarse}} \leq n(\epsilon_{\mathrm{tok}}+\epsilon_{\mathrm{unmask}}).\] Combining this bound with Eq. 15 gives Eq. ?? . ◻
13 Incompleteness of the token-only policy gradient↩︎
In this section, we show that optimizing only the token gradient can miss a direction that improves the reward, in the following setting. Recall that at each denoising step \(t\), the model produces probability \(\pi_{\theta}(\cdot | \boldsymbol{z}^t)\) over the full vocabulary at every position. Let \(\mathcal{M}_t\) denote the set of positions that are still masked in state \(\boldsymbol{z}^t\). For each masked position \(k \in \mathcal{M}_t\), we define an unmasking score \(v_k^{t}\) as the highest token confidence at position \(k\). The generation process then selects the top positions with the highest unmasking scores to unmask, while the rest of the tokens get remasked. We define \(\nabla_{\boldsymbol{\theta}} J_{| \text{pos}}\) and \(\nabla_{\boldsymbol{\theta}} J_{| \text{tok}}\) to be the position part and token part of the policy gradient as follows.
\[\begin{align} \nabla_{\boldsymbol{\theta}} J_{| \text{pos}} &= \mathop{\mathrm{\mathbb{E}}}_{ \hat{\boldsymbol{z}} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c}) }\bigg[ R(\boldsymbol{c}, \boldsymbol{z}^T)\!\sum_{t=1}^{T}\!\Big( \nabla_{\theta} \log p_\text{unmask}(\mathcal{U}_t | \hat{\boldsymbol{z}}^t, \boldsymbol{z}^{t-1}, \boldsymbol{c}) \Big) \bigg] \\ \nabla_{\boldsymbol{\theta}} J_{| \text{tok}} &= \mathop{\mathrm{\mathbb{E}}}_{\hat{\boldsymbol{z}} \sim \pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{c})}\bigg[ R(\boldsymbol{c}, \boldsymbol{z}^T)\!\sum_{t=1}^{T}\!\Big( \nabla_{\theta} \log \pi_\theta(\boldsymbol{z}^t \!\mid\! \boldsymbol{c}, \boldsymbol{z}^{t-1} ) \Big) \bigg] \end{align}\]
Proposition 2. There exists a finite-horizon MDLM generation problem and a linear policy class \(\pi_{\boldsymbol{\theta}}(\cdot \mid \boldsymbol{z}^t )\) such that, when the unmasking policy is defined by selecting positions with highest token confidence, the following holds:
There exists a direction \(\boldsymbol{v}\) such that the token-only policy gradient is orthogonal to this direction, and therefore, does not move parameters along the direction of \(\boldsymbol{v}\): \[\begin{align} \langle \nabla_{\boldsymbol{\theta}} J_{|\mathrm{tok}}(\boldsymbol{\theta}), \boldsymbol{v} \rangle = 0 \quad \text{for all } \boldsymbol{\theta}. \end{align}\] The true objective is not stationary along this direction, i.e., \[\exists \boldsymbol{\theta} \text{ such that } \langle \nabla_{\boldsymbol{\theta}} J(\boldsymbol{\theta}), \boldsymbol{v} \rangle \neq 0.\] In particular, there exists \(\boldsymbol{\theta}\) and \(\boldsymbol{\theta}' = \boldsymbol{\theta} + \epsilon \boldsymbol{v}\) for some \(\epsilon > 0\) such that \[J(\boldsymbol{\theta}') > J(\boldsymbol{\theta}),\] while the token-only gradient provides no update in the \(\boldsymbol{v}\) direction at \(\boldsymbol{\theta}\).
Proof. We construct a two-position example with sequence length \(n=2\) and vocabulary \(\{a,b\}\). At each step, exactly one masked position is selected and filled, producing \(T=3\) states and \(T-1=2\) transitions. The initial state is \(\boldsymbol{z}^1 = ([\texttt{MASK}], [\texttt{MASK}])\), \(\boldsymbol{z}^2\) has one masked token, and \(\boldsymbol{z}^3\) is the completely unmasked sequence.
The reward is defined as \[\begin{align} R(x_1,x_2) = \mathbb{I}\{x_1 = a\}. \end{align}\] We consider the linear policy class on the feature vector \(\phi(\boldsymbol{z}, k)\) for each masked position \(k \in \{1, 2\}\). More specifically, the logits of the policy is given by \[\begin{align} s_{\boldsymbol{\theta}}(\boldsymbol{z}, k) = \boldsymbol{\theta}^\top \phi(\boldsymbol{z}, k). \end{align}\] The probability of each symbol in vocabulary is given by the sigmoid of the logits. \[\begin{align} p(a \; | \; \boldsymbol{z}, k ) = \sigma( s_{\boldsymbol{\theta}}(\boldsymbol{z}, k) ), \quad \text{and} \quad p(b \; | \; \boldsymbol{z}, k ) = 1 - \sigma( s_{\boldsymbol{\theta}}(\boldsymbol{z}, k) ). \end{align}\] We consider the following set of one-hot features in \(\phi(\boldsymbol{z}, k) \in \mathbb{R}^3\) to represent \(\boldsymbol{z}\) and position \(k\): \[\begin{align} \phi(\boldsymbol{z}, k) = \begin{cases} \boldsymbol{e}_1 \quad & \text{if} \quad (\boldsymbol{z}, k) = ((\texttt{[MASK]}, \texttt{[MASK]}), \; 1)\\ \boldsymbol{e}_2 \quad & \text{if} \quad (\boldsymbol{z}, k) = ((\texttt{[MASK]}, \texttt{[MASK]}), \; 2)\\ \boldsymbol{e}_3 \quad & \text{if} \quad (\boldsymbol{z}, k) = ((\texttt{[MASK]}, \; a), \; 1) \quad \text{or} \quad (\boldsymbol{z}, k) = ((\texttt{[MASK]}, \; b), \; 1) \\ \end{cases} \end{align}\] This gives us that \[\begin{align} p(a \; | \; \boldsymbol{z}^1, \; 1) = \sigma( \theta_1 ) \quad & \text{and} \quad p(a \; | \; \boldsymbol{z}^1, \; 2) = \sigma( \theta_2 ) \\ p(a \; | \; \boldsymbol{z}^2, \; 1) = \sigma(\theta_3) \quad & \text{and} \quad p(a \; | \; \boldsymbol{z}^2, \; 2) = \sigma( 0 ) = 0.5 \end{align}\] We use the confidence score as the maximum token probability, that gives \(v_k^{t} = \max_x \pi_{\theta} (x \mid \boldsymbol{z}^t, k) = \sigma(|s_{\boldsymbol{\theta}}(\boldsymbol{z}^t, k)|)\). When \(|\mathcal{U}_t| = 1\), at denoising step \(t\), we choose \(k^{\text{th}}\) position to decode with the following probability: \[\begin{align} p_{\mathrm{unmask}}(u_1 = k \mid \boldsymbol{z}^t) = \frac{\exp\!\big(v_k^{t} / \tau\big)}{\exp\!\big(v_k^{t} / \tau\big) + \sum_{j \in \mathcal{M}_t \setminus \{ k \} } \exp\!\big(v_j^{t} / \tau\big)}. \end{align}\] We now compute the expected reward \(J(\boldsymbol{\theta})\). As \(R(x_1,x_2) = \mathbb{I}\{x_1 = a\}\), we compute the probability of \(x_1 = a\) under policy \(\pi_{\boldsymbol{\theta}}\). The probability of the first position getting selected is \(p_{\mathrm{unmask}}(1 \mid \boldsymbol{z}^1)\) and then the probability of position \(1\) getting filled with \(a\) is \(\sigma(\theta_1)\). The probability of the second position getting selected at initial stage is \(p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\) and irrespective of the value at second position, \(x_1 = a\) with probability \(\sigma(\theta_3)\). Therefore, the total expected reward is \[\begin{align} J(\boldsymbol{\theta}) &= p_{\mathrm{unmask}}(1 \mid \boldsymbol{z}^1) \pi_{\theta} (a \mid \boldsymbol{z}^1, 1) + p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) \pi_{\theta} (a \mid \boldsymbol{z}^2, 1) \\ &= \pi_{\theta}(a \mid \boldsymbol{z}^1, 1) + p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) ( \pi_{\theta}(a \mid \boldsymbol{z}^2, 1) - p_{\mathrm{val}}(a \mid \boldsymbol{z}^1, 1) ) \\ &= \sigma(\theta_1) \; + \; p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) ( \sigma(\theta_3) - \sigma(\theta_1) ) \end{align}\]
13.0.0.1 Zero token-only gradient along direction \(v\).
We define the direction \(\boldsymbol{v} = (0,1,0)\). The token-only gradient along direction \(v\) is given by \[\begin{align} \partial_{\boldsymbol{\theta}_2} J_{|\mathrm{tok}} = \mathop{\mathrm{\mathbb{E}}}_{\xi \sim \pi_{\boldsymbol{\theta}}} \Big[ \sum_t R(\xi)\,\partial_{\theta_2} \log \pi_{\theta} (x_{u_1} \mid \boldsymbol{z}^t, u_1) \Big]. \end{align}\] Observe that \(\pi_{\theta} (x_{u_1} \mid \boldsymbol{z}^t, u_1)\) depends on \(\theta_2\) only for \(t=1\) and \(u_1 = 2\). Therefore, we can write the above as \[\begin{align} \partial_{\boldsymbol{\theta}_2} J_{|\mathrm{tok}} &= \mathop{\mathrm{\mathbb{E}}}_{\xi \sim \pi_{\boldsymbol{\theta}}} \Big[ R(\xi) \mathbf{1} \{ u_1 = 2 \} \,\partial_{\theta_2} \log \pi_{\theta} (x_2 \mid \boldsymbol{z}^1, u_1 = 2) \Big] \\ &= p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) \; \mathop{\mathrm{\mathbb{E}}}_{\xi \sim \pi_{\boldsymbol{\theta}}} \Big[ R(\xi) \,\partial_{\theta_2} \log \pi_{\theta} (x_2 \mid \boldsymbol{z}^1, u_1 = 2) \; \mid \; u_1 = 2 \Big] \\ &= p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) \; \mathop{\mathrm{\mathbb{E}}}_{x_2 \sim \pi_{\theta} (\cdot \mid \boldsymbol{z}^1, 2)} \; \Big[ \; \partial_{\theta_2} \log \pi_{\theta} (x_2 \mid \boldsymbol{z}^1, u_1 = 2) \mathop{\mathrm{\mathbb{E}}}_{u_2, x_1} \Big[ R(\xi) \; \mid \; u_1 = 2, \; x_2 \Big] \Big] \\ &= p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) \; \sigma( \theta_3 ) \; \mathop{\mathrm{\mathbb{E}}}_{x_2 \sim \pi_{\theta}(\cdot \mid \boldsymbol{z}^1, 2)} \; \Big[ \; \partial_{\theta_2} \log \pi_{\theta} (x_2 \mid \boldsymbol{z}^1, u_1 = 2) \Big] \end{align}\] We now show that the expectation in the above equation becomes zero. To show that, we write the expectation as sum and then change the order of expectation and derivative. The change in the order is valid because both are with respect to different variables. \[\begin{align} \mathop{\mathrm{\mathbb{E}}}_{x_2 \sim \pi_{\theta}(\cdot \mid \boldsymbol{z}^1, 2)} \; \Big[ \; \partial_{\theta_2} \log \pi_{\theta}(x_2 \mid \boldsymbol{z}^1, u_1 = 2) \Big] &= \sum_{x_2} \pi_{\theta}(x_2 \mid \boldsymbol{z}^1, 2) \frac{ \partial_{\theta_2} \pi_{\theta}(x_2 \mid \boldsymbol{z}^1, 2) }{ \pi_{\theta}(x_2 \mid \boldsymbol{z}^1, 2) } \\ = \partial_{\theta_2} \sum_{x_2} \pi_{\theta} & (x_2 \mid \boldsymbol{z}^1, 2) = \partial_{\theta_2} 1 = 0 \end{align}\]
13.0.0.2 Increasing the true objective \(J(\boldsymbol{\theta})\) along \(\boldsymbol{v}\) direction.
We show that for any parameter \(\boldsymbol{\theta}\) satisfying \(\sigma( \theta_1 ) > \sigma( \theta_3 )\), the gradient along \(\theta_2\) improves the performance and \(\partial_{\theta_2} J(\boldsymbol{\theta}) \neq 0\).
Observe that for such a parameter vector \(\boldsymbol{\theta}\), \(\sigma(\theta_3)-\sigma(\theta_1) < 0\) and hence \(J(\boldsymbol{\theta})\) is strictly decreasing as \(p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\) increases. Now \(p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\) is a strictly increasing function of \(v_2^{1} = \sigma(|\theta_2|),\) and \(\sigma(|\theta_2|)\) is minimized at \(\theta_2=0\) and strictly larger for \(\theta_2\neq 0\). Therefore, keeping \(\theta_1\) and \(\theta_3\) fixed, the objective \(J(\boldsymbol{\theta})\) is maximized when \(\theta_2=0.\) Therefore, for any \(\boldsymbol{\theta}\) such that \(\sigma( \theta_1 ) > \sigma( \theta_3 )\) and \(\theta_2 \neq 0\), consider \(\boldsymbol{\theta}' = \boldsymbol{\theta}+ \epsilon \boldsymbol{v}\) such that \(|\theta'_{2}| < |\theta_2|\), the true objective improves: \[\begin{align} J(\boldsymbol{\theta}') > J(\boldsymbol{\theta}). \end{align}\]
13.0.0.3 Nonzero gradient along \(\boldsymbol{v}\).
We now show that the true objective exhibits a nonzero directional derivative along \(\boldsymbol{v}=(0,1,0)\). To this end, we explicitly compute the gradient with respect to \(\theta_2\) and verify that it is nonvanishing under the stated conditions.
Differentiating the true objective with respect to \(\theta_2\) yields \[\partial_{\theta_2} J(\boldsymbol{\theta}) = \partial_{\theta_2} p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\, \big(\sigma(\theta_3)-\sigma(\theta_1)\big).\]
Since \(\sigma(\theta_3)-\sigma(\theta_1) < 0\), to conclude that \(\partial_{\theta_2} J(\boldsymbol{\theta})\neq 0\), it suffices to establish that \[\partial_{\theta_2} p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) \neq 0 \qquad \text{for } \theta_2\neq 0.\]
We now analyze the dependence of \(p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\) on \(\theta_2\). Recall that \(v_1^1=\sigma(|\theta_1|), v^1_2=\sigma(|\theta_2|)\), and \[p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) = \frac{\exp(v^1_2/\tau)}{\exp(v^1_1/\tau)+\exp(v^1_2/\tau)}.\] Differentiating \(p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\) gives \[\partial_{\theta_2} p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) = \frac{1}{\tau}\, p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\big(1-p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\big) \,\partial_{\theta_2} v_2^1.\] Observe that the multiplicative factor \(p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\big(1-p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1)\big)\) is strictly positive because \(0 < p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) < 1\). Therefore, the only remaining question is whether \(\partial_{\theta_2} v_2^1\) vanishes. To compute this term, note that \(v_2^1=\sigma(|\theta_2|)\) is a composition of smooth functions away from \(\theta_2=0\). Differentiating yields \[\partial_{\theta_2} v_2^1 = \sigma(|\theta_2|)\big(1-\sigma(|\theta_2|)\big)\,\mathrm{sign}(\theta_2), \qquad \text{for } \theta_2\neq 0.\] Since \(\sigma(|\theta_2|)\in(0,1)\) for all finite \(\theta_2\), the factor \(\sigma(|\theta_2|)(1-\sigma(|\theta_2|))\) is strictly positive, and \(\mathrm{sign}(\theta_2)\neq 0\) whenever \(\theta_2\neq 0\). Hence, \(\partial_{\theta_2} v_2^1 \neq 0\) for all \(\theta_2\neq 0.\)
Combining the above observations, we conclude that \[\partial_{\theta_2} p_{\mathrm{unmask}}(2 \mid \boldsymbol{z}^1) \neq 0 \qquad \text{for all } \theta_2\neq 0,\] and therefore \[\partial_{\theta_2} J(\boldsymbol{\theta})\neq 0 \qquad \text{whenever } \sigma(\theta_1)>\sigma(\theta_3) \text{ and } \theta_2\neq 0.\]
Finally, since \(\boldsymbol{v}=(0,1,0)\), the directional derivative along \(\boldsymbol{v}\) is given by \[\langle \nabla_{\boldsymbol{\theta}} J(\boldsymbol{\theta}), \boldsymbol{v}\rangle = \partial_{\theta_2} J(\boldsymbol{\theta}).\] Thus, \[\langle \nabla_{\boldsymbol{\theta}} J(\boldsymbol{\theta}), \boldsymbol{v}\rangle \neq 0,\] which establishes that the true objective is not stationary along \(\boldsymbol{v}\). ◻
References↩︎
Code available at https://github.com/Haran71/mask-aware-policy-gradients↩︎
Equal contribution.↩︎
We use the term remasking following LLaDA [11]: at each step the model predicts tokens at all masked positions, and those not selected to remain revealed are returned to \(\texttt{[MASK]}\). Such positions were never actually unmasked, so this differs from self-correction methods that revert previously committed tokens.↩︎