SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

## Abstract

Abstract Diffusion large language models (dLLMs) are emerging as an efficient alternative to autoregressive models due to their ability to decode multiple tokens in parallel. However, aligning dLLMs with human preferences or task-specific rewards via reinforcement learning (RL) is challenging because their intractable log-likelihood precludes the direct application of standard policy gradient methods. While prior work uses surrogates like the evidence lower bound (ELBO), these one-sided approximations can introduce significant policy gradient bias. To address this, we propose the Sandwiched Policy Gradient (SPG) that leverages both an upper and a lower bound of the true log-likelihood. Experiments show that SPG significantly outperforms baselines based on ELBO or one-step estimation. Specifically, SPG improves the accuracy over state-of-the-art RL methods for dLLMs by 3.6% in GSM8K, 2.6% in MATH500, 18.4% in Countdown and 27.0% in Sudoku.

## 1 Introduction

Diffusion models, originally pioneered for high-fidelity image generation (Song et al., 2020; Ho et al., 2020), have recently emerged as a powerful and efficient paradigm for text generation (Austin et al., 2021; Campbell et al., 2022; Sun et al., 2022; Lou et al., 2023; Sahoo et al., 2024; Shi et al., 2024). These models operate in a discrete space but share architectural similarities with their continuous counterparts (Peebles and Xie, 2023). They employ a fixed noising process that progressively corrupts text data, while a neural network is trained to learn the reverse, denoising process. For instance, Masked Diffusion Language Model (MDLM) (Sahoo et al., 2024) uses random masking as its forward noising process and optimizes an Evidence Lower Bound (ELBO) of the log-likelihood. This ELBO-based objective has been widely adopted by subsequent large-scale diffusion language models (dLLMs), including LLaDA (Nie et al., 2025) and DREAM (Gong et al., 2024).

A key advantage of dLLMs over their autoregressive (AR) counterparts is their ability to decode multiple tokens in parallel. This parallelism can significantly reduce inference latency, making it an attractive alternative for scalable language modeling (Wang et al., 2025a; Labs et al., 2025).

Aligning large language models with human preferences (Ouyang et al., 2022) or task-specific rewards (e.g., inducing reasoning behavior) (Shao et al., 2024; Guo et al., 2025) typically requires a post-training stage of reinforcement learning (RL). However, applying RL to dLLMs remains underexplored. A principal challenge is the computationally intractable log-likelihood of dLLMs, which is essential for accurate policy gradient estimation. To circumvent this, recent works (Zhao et al., 2025; Yang et al., 2025; Zhu et al., 2025; Tang et al., 2025) adapt standard RL and preference optimization algorithms, such as GRPO (Shao et al., 2024) and DPO (Rafailov et al., 2023), by using the ELBO or a one-step estimation as a surrogate for the true likelihood. While straightforward, this approximation leads to misaligned policy gradients, and potential suboptimal performance.

Figure: Figure 2: The training process of SPG for MDLM. Left: From a prompt $\bm{c}$, we generate responses $\{\bm{x}^{j}\}_{j=1}^{g}$. We then maximize a lower bound on the likelihood $\pi_{\bm{\theta}}(\bm{x}^{j}\mid\bm{c})$ for high-reward responses while minimizing an upper bound for low-reward ones. Right: The upper/lower bound of likelihood is estimated via Monte Carlo using a block-wise masking strategy, where a random block is selected for masking, with earlier blocks kept clean and later blocks fully masked. The example shows a sequence of length 9 with a block size of 3, where the current generation block is highlighted in yellow.
Refer to caption: iclr2026/images/main.png

To address these limitations, we propose Sandwiched Policy Gradient (SPG), a novel reinforcement learning algorithm for diffusion language models that computes a more robust and less biased policy gradient. As illustrated in [Figure˜2](https://arxiv.org/html/2510.09541v2#S1.F2), our core idea is to “sandwich” the intractable log-likelihood of a generated sequence: we maximize a tractable lower bound for positive-reward sequences while minimizing an upper bound for negative-reward ones. To ensure a stable estimation of these bounds, we also propose a block-wise masking strategy that better aligns data distributions during policy rollout and optimization. SPG achieves state-of-the-art performance on four mathematical and logical reasoning benchmarks, improving accuracy by up to 3.6% on GSM8K, 2.6% on MATH500, 18.4% on Countdown, and 27.0% on Sudoku compared to the state-of-the-art
RL algorithms for diffusion language models.

In summary, our main contributions are:

- •
A new policy gradient algorithm, SPG, which reduces bias by optimizing sandwiched variational bounds based on reward.
- •
A block-wise masking technique that improves the stability of the training objective’s estimation.
- •
State-of-the-art results among RL algorithms for diffusion language models on four reasoning benchmarks, demonstrating the effectiveness of our approach.

## 2 Background

In this section, we provide a brief overview of the masked diffusion language model (MDLM) and reinforcement learning for text diffusion models.

#### Notation.

We denote scalars by lowercase letters ($x$), vectors by bold lowercase ($\bm{x}$), and sequences by $\bm{x}_{1:n}$. $[k]$ represents $\{1,\dots,k\}$. $\text{Cat}(\bm{x}\mid\bm{p})$ is the categorical distribution over $\bm{x}$ with
probabilities $\bm{p}$, and $\mathcal{U}[a,b]$ denotes the uniform distribution in $[a,b]$. Throughout the paper, we use $i\in[n]$ for position of the token, $j\in[g]$ for a sequence in a group of rollouts, and $t$ for the diffusion timestep. For discrete time processes, $t\in[T]$, while for continuous-time Markov chains, $t\in[0,1]$.

### 2.1 Masked Diffusion Language Models

Diffusion models for language learn to generate text by reversing a gradual noising process. Specifically, Masked Diffusion Language Models (MDLMs) (Sahoo et al., 2024) start with clean text $\bm{x}_{1:n}$ and corrupt it into $\bm{z}_{t}\equiv\bm{z}_{t,1:n}$ over a continuous timestep $t\in[0,1]$ by progressively replacing tokens with a special [mask] token. At $t=0$, the data is original ($\bm{z}_{0}=\bm{x}$), while at $t=1$, the sequence is fully masked ($\bm{z}_{1}$ is all [mask] tokens).
Each token is corrupted independently according to the forward transition kernel:

$$ $q_{t\mid 0}(\bm{z}_{t,i}\mid\bm{x}_{i})=\text{Cat}\big(\bm{z}_{t,i}\mid\alpha_{t}\bm{x}_{i}+(1-\alpha_{t})\bm{m}\big),$ (1) $$

where $\bm{m}$ is the one-hot representation of the [mask] token. The noise schedule, $\alpha_{t}\in[0,1]$, is a strictly decreasing function, such as the linear schedule $\alpha_{t}=1-t$, with $\alpha_{0}=1$ and $\alpha_{1}=0$.

In the reverse process, a neural network, which we denote as the policy $\pi_{\bm{\theta}}$, is then trained to perform the reverse process: predicting the original tokens $\bm{x}$ from a corrupted version $\bm{z}_{t}$.
The transition from $\bm{z}_{t}$ to $\bm{z}_{s}$ ($s<t$) is parameterized with $\pi_{\bm{\theta}}$ as follows:

$$ $p_{\theta}(\bm{z}_{s}\mid\bm{z}_{t})=q\left(\bm{z}_{s}\mid\bm{z}_{t},\bm{x}=\pi_{\bm{\theta}}(\cdot\mid\bm{z}_{t})\right)=\begin{cases}\text{Cat}(\mathbf{z}_{s};\mathbf{z}_{t}),&\mathbf{z}_{t}\neq\mathbf{m},\\ \text{Cat}\left(\mathbf{z}_{s};\frac{(1-\alpha_{s})\mathbf{m}+(\alpha_{s}-\alpha_{t})\pi_{\bm{\theta}}(\cdot\mid\bm{z}_{t})}{1-\alpha_{t}}\right),&\mathbf{z}_{t}=\mathbf{m}.\end{cases}$ $$

The policy is achieved by maximizing the Evidence Lower Bound (ELBO) of the log-likelihood of each clean sequence $\bm{x}\sim p_{\text{data}}$, which simplifies to the following objective:

$$ $\mathcal{L}_{\text{ELBO}}(\bm{x};\bm{\theta})=\mathbb{E}_{t,\bm{z}_{t}}\bigg[\sum_{i=1}^{n}w(t)\cdot\mathbbm{1}(\bm{z}_{t,i}=\bm{m})\cdot\log\pi_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t})\bigg],$ (2) $$

where $w(t)=\alpha_{t}^{\prime}/(\alpha_{t}-1)$ is a time-dependent loss weight, and the expectation is over a random timestep $t\sim\mathcal{U}[0,1]$ and the corrupted sequence $\bm{z}_{t}\sim q_{t\mid 0}(\cdot\mid\bm{x})$. In essence, this objective trains the model to “fill in the blanks” by predicting the original tokens at masked positions. For a more comprehensive overview of MDLM, please refer to [Appendix˜A](https://arxiv.org/html/2510.09541v2#A1) and Sahoo et al. (2024).

### 2.2 Reinforcement Learning for Diffusion Language Models

Reinforcement Learning (RL) aligns a language model with desired objectives by treating it as a policy $\pi_{\bm{\theta}}$ that generates a response $\bm{x}$ to a prompt $\bm{c}$. A reward function $R(\bm{c},\bm{x})$ provides a scalar score for the response, and the training goal is to update $\bm{\theta}$ to maximize the expected reward: $\mathcal{J}(\bm{\theta})\vcentcolon=\mathbb{E}_{\bm{x}\sim\pi_{\bm{\theta}}(\cdot\mid\bm{c})}[R(\bm{c},\bm{x})]$. This objective is commonly optimized using policy gradient methods, which rely on the following gradient estimator.

$$ $\nabla_{\bm{\theta}}\mathcal{J}(\bm{\theta})=\mathbb{E}_{\bm{x}\sim\pi_{\bm{\theta}}(\cdot\mid\bm{c})}\bigg[R(\bm{c},\bm{x})\nabla_{\bm{\theta}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{c})\bigg].$ (3) $$

#### The Intractability Challenge.

A central challenge in applying RL to diffusion models is that the policy’s log-likelihood, $\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{c})$, is *intractable* and cannot be computed directly. To overcome this, prior work (Zhu et al., 2025; Yang et al., 2025) approximates this term using its ELBO, effectively replacing $\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{c})$ with a score derived from the pre-training objective in [Equation˜2](https://arxiv.org/html/2510.09541v2#S2.E2).

However, this popular workaround introduces a critical flaw. The ELBO is only a lower bound on the true log-likelihood ($\text{ELBO}\leq\log\pi_{\bm{\theta}}$). Consequently, the RL objective is only a valid lower bound on the true expected reward if all rewards $R(\bm{c},\bm{x})$ are *non-negative*. This constraint prevents the model from effectively learning from negative feedback (i.e., penalizing bad outputs) and is incompatible with advanced RL algorithms that use relative or negative rewards (Shao et al., 2024), biasing the final policy. Our work aims to resolve this limitation.

Figure: Algorithm 1 SPG: Sandwiched Policy Gradient for Masked dLLMs

## 3 Sandwiched Policy Gradient with Evidence Bounds

We introduce SPG, a novel policy gradient algorithm designed for masked diffusion language models ([Algorithm˜1](https://arxiv.org/html/2510.09541v2#alg1)). Our method aims to address a critical issue in applying reinforcement learning to dLLMs by creating a valid optimization objective based on tractable bounds of the model’s evidence.

### 3.1 A Lower Bound Objective for Policy Optimization

Our approach is based on group relative policy optimization (Shao et al., 2024; Liu et al., 2025b). For a given prompt $\bm{c}$, we generate a group of $g$ responses $\{\bm{x}^{j}\}_{j=1}^{g}$ from the policy $\pi_{\bm{\theta}}$. We then compute the advantage $A^{j}(\bm{c},\bm{x}^{j})\vcentcolon=R(\bm{c},\bm{x}^{j})-\frac{1}{g}\sum_{\jmath=1}^{g}R(\bm{c},\bm{x}^{\jmath})$. Moreover, we transform the conventional policy optimization objective as an advantage-weighted log-likelihood objective, for reasons that will be clear later:

$$ $\mathcal{J}^{\text{group}}({\bm{\theta}})=\mathbb{E}_{\bm{c},\{\bm{x}^{j}\}\sim\pi_{{\text{sg}{[{\bm{\theta}}]}}}}\bigg[\penalty 10000\ \frac{1}{g}\sum_{j=1}^{g}A^{j}(\bm{x}^{j},\bm{c})\log\pi_{{\bm{\theta}}}(\bm{x}^{j}\mid\bm{c})\penalty 10000\ \bigg],$ (4) $$

where ${\text{sg}{[{\bm{\theta}}]}}$ indicates that gradients are not computed for the policy that generates the samples. This objective encourages generations with positive advantages ($A^{j}>0$) and discourages those with negative advantages ($A^{j}<0$).

For dLLMs, the log-likelihood $\log\pi_{{\bm{\theta}}}$ is intractable. A common surrogate is the evidence lower bound (ELBO). While maximizing the ELBO is a valid way to increase the true log-likelihood, minimizing the ELBO for negatively-rewarded samples does not guarantee a reduction in the true log-likelihood. To address this, we propose a *sandwiched* objective. For samples with positive advantages, we maximize the ELBO. For samples with negative advantages, we instead minimize a tractable evidence *upper* bound (EUBO), $\mathcal{L}_{\text{EUBO}}$. This creates a true lower bound for the original objective:

$$ $\mathcal{J}_{\text{SPG}}({\bm{\theta}})=\mathbb{E}\bigg[\penalty 10000\ \frac{1}{g}\sum_{j=1}^{g}\left(\mathbbm{1}_{A^{j}\geq 0}\cdot A^{j}\mathcal{L}_{\text{ELBO}}(\bm{x}^{j}\mid\bm{c};{\bm{\theta}})+\mathbbm{1}_{A^{j}<0}\cdot A^{j}\mathcal{L}_{\text{EUBO}}(\bm{x}^{j}\mid\bm{c};{\bm{\theta}})\right)\bigg],$ (5) $$

where the expectation is take with respect to ${\bm{c},\{\bm{x}^{j}\}\sim\pi_{{\text{sg}{[{\bm{\theta}}]}}}}$.
Since $\mathcal{L}_{\text{ELBO}}\leq\log\pi_{\bm{\theta}}\leq\mathcal{L}_{\text{EUBO}}$, it follows that $\mathcal{J}_{\text{SPG}}({\bm{\theta}})\leq\mathcal{J}^{\text{group}}({\bm{\theta}})$. Maximizing this tractable bound therefore serves as a valid proxy for optimizing the true objective.

### 3.2 A Tractable Evidence Upper Bound

To effectively penalize negatively-rewarded samples by minimizing their log-likelihood, we require a tractable EUBO, which we derive in the following theorem based on the Rényi variational bound.

###### Theorem 1 (Evidence Upper Bound for Masked Diffusion) .

Assume the forward denoising process has $T$ steps with a monotonic schedule $\alpha_{t}$. For any $\beta\geq 1$ and a sequence $\bm{x}_{1:n}$, we have:

$$ $\mathcal{L}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})=\frac{1}{\beta}\sum_{i=1}^{n}\log\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}}\Bigg[\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\cdot\mathbbm{1}(\bm{z}_{t+1,i}=\bm{m})\cdot\pi_{{\bm{\theta}}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t+1})\Bigg]+C(T),$ (6) $$

where $C(T)\vcentcolon=\mathbbm{1}(\beta<n)\cdot\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\Big[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\Big]$ is a constant independent of $\bm{{\bm{\theta}}}$.

Here, $\beta\geq 1$ is a hyperparameter that controls the tightness of the bound, with values closer to 1 yielding a tighter bound. The expectation is taken over the timestep $t\sim\mathcal{U}[0,1]$ and the noised latent $\bm{z}_{t}\sim q_{t\mid 0}(\cdot\mid\bm{x})$.

###### Corollary 1 .

Taking the limit of $T\rightarrow\infty$, we have:

$$ $\begin{split}&\nabla_{\bm{{\bm{\theta}}}}\mathcal{L}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})=\nabla_{\bm{{\bm{\theta}}}}\Big(\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})+C(T)\Big)=\nabla_{\bm{{\bm{\theta}}}}\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}}),\penalty 10000\ \penalty 10000\ \penalty 10000\ \penalty 10000\ \text{where}\\ &\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1:n};{\bm{\theta}})=\frac{1}{\beta}\sum_{i=1}^{n}\log\mathbb{E}_{t,\bm{z}_{t}}\left[w(t)\cdot\mathbbm{1}(\bm{z}_{t,i}=\bm{m})\cdot\pi^{\beta}_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t})\right].\end{split}$ (7) $$

In practice, we estimate $\tilde{\mathcal{L}}_{\text{EUBO}}$ using Monte Carlo sampling and plug it in Equation [5](https://arxiv.org/html/2510.09541v2#S3.E5) in place of $\mathcal{L}_{\text{EUBO}}$. The proof and theoretical analysis are provided in Appendix [B](https://arxiv.org/html/2510.09541v2#A2).

###### Theorem 1 (Evidence Upper Bound for Masked Diffusion) .

###### Corollary 1 .

#### Remark.

A key structural difference from $\mathcal{L}_{\text{ELBO}}$ is that the logarithm in $\mathcal{L}_{\text{EUBO}}$ ([Equation˜6](https://arxiv.org/html/2510.09541v2#S3.E6)) appears outside the expectation.
Therefore, in practice, due to Jensen’s inequality, applying the concave logarithm to a Monte Carlo estimate of the expectation’s argument yields a biased estimate of the true EUBO. While it is possible to derive a looser but unbiased bound using inequalities like $\log(x)\leq x-1$, we found this approach empirically worse by widening the gap to the true log-likelihood, as shown in [Table˜10](https://arxiv.org/html/2510.09541v2#A5.T10). We therefore retain the tighter, albeit slightly biased, formulation.

### 3.3 Practical Considerations

#### Block-Wise Masking Strategy for Monte Carlo Estimation.

In practice, we approximate $\mathcal{L}_{\text{ELBO}}$ and $\tilde{\mathcal{L}}_{\text{EUBO}}$ in [Equation˜5](https://arxiv.org/html/2510.09541v2#S3.E5) via Monte Carlo sampling: for each $\bm{x}^{j}$, we randomly sample $m$ timesteps $\{t_{\tau}\}_{\tau=1}^{m}$ and generate the corresponding partially masked samples $\{\bm{z}_{t_{\tau}}^{j}\}_{\tau=1}^{m}\sim q(\cdot\mid\bm{x}^{j})$. One straightforward approach as used in Yang et al. (2025) would be to apply random masking to clean sequences. However, recent dLLMs like LLaDA (Nie et al., 2025) employ a block-wise semi-autoregressive unmasking strategy during generation and achieve state-of-the-art performance over random unmasking. As a result, the policy rollout process actually encounters a much narrower and more structured set of partially masked sequences than with fully random masking.

To better align data distributions during policy rollout and optimization, we adopt a block-wise masking strategy rather than random masking. As depicted in [Figure˜2](https://arxiv.org/html/2510.09541v2#S1.F2), the sequence is divided into several blocks, and a random block is selected, with all preceding blocks left clean and all following blocks fully masked. Within the chosen block, tokens are randomly masked. Additionally, following D1 (Zhao et al., 2025), we lightly perturb the prompt and clean blocks by randomly masking tokens with a small probability $p_{\text{mask}}=0.15$ to enhance stability and generalization.

Altogether, our block-wise masking strategy improves the stability of the objective’s estimation and the efficiency of policy optimization. While similar block-wise masking approaches have been explored in concurrent work for supervised fine-tuning or block diffusion models (Sun et al., 2025; Wang et al., 2025b), our focus is on RL for full-attention masked dLLMs. As shown in [Figure˜6](https://arxiv.org/html/2510.09541v2#S4.F6), our models trained with block-wise masking generalize well to various inference strategies.

#### Mixture of Upper and Lower Bound for Negative Advantage Traces.

Monte Carlo estimation of [Equation˜6](https://arxiv.org/html/2510.09541v2#S3.E6) leads to a biased estimation to $\tilde{\mathcal{L}}_{\text{EUBO}}$ and potentially requires a substantial number of samples to get reliable approximations, resulting in high computational costs and instability during training. To address these challenges, we use a mixture of $\tilde{\mathcal{L}}_{\text{EUBO}}$ and $\mathcal{L}_{\text{ELBO}}$ as a more practical log-likelihood approximation for negative advantage traces:

$$ $\tilde{\mathcal{L}}_{\text{Mix}}(\bm{x}\mid\bm{c};\bm{{\bm{\theta}}}):=\omega\cdot\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}\mid\bm{c};\bm{{\bm{\theta}}})+(1-\omega)\cdot\mathcal{L}_{\text{ELBO}}(\bm{x}\mid\bm{c};\bm{{\bm{\theta}}})$ (8) $$

where $0\leq\omega\leq 1$ is a blend coefficient. Intuitively, the upper bound $\tilde{\mathcal{L}}_{\text{EUBO}}$ sharpens the model decisions by applying a $\beta$-power adjustment to the original model output, acting as a strong correction signal for negative advantage traces. In contrast, the lower bound $\mathcal{L}_{\text{ELBO}}$ is easier and more stable to estimate with a small number of Monte Carlo samples, but it tends to introduce larger, systematic bias relative to the true log-likelihood. In particular, as a conservative approximation, $\mathcal{L}_{\text{ELBO}}$ alone is insufficient for effectively penalizing negative advantage traces, thus limiting its efficacy. Therefore, combining them allows us to harness the strengths of each, resulting in a more effective log-likelihood estimation in practice. In the following proposition, we formalize the advantages of using the mixture by deriving the gradient of the mixture loss and analyzing the variance of the gradient.

###### Proposition 1 (Optimal Mixture Strictly Reduces Variance) .

Fix a coordinate $k$ and let

$$ $\rho_{\beta}\coloneqq{w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})}/{\mathbb{E}\left[w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]},$ $$

where $w(t,\bm{z}_{t})\coloneqq w(t)\mathbbm{1}(z_{t}=\bm{m})$. Then, the gradient of mixture objective ([8](https://arxiv.org/html/2510.09541v2#S3.E8)) is given by

$$ $\displaystyle g_{\omega,k}=((1-\omega)w(t,\bm{z}_{t})+\omega\rho_{\beta})\,\partial_{\bm{\theta}_{k}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t},\bm{c}).$ (9) $$

If $\operatorname{Var}((\rho_{\beta}-w(t,\bm{z}_{t}))\partial_{\bm{\theta}_{k}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t},\bm{c}))>0$, then $\operatorname{Var}[g_{\omega,k}]$ is a strictly convex quadratic in $\omega$ and thus admits a unique minimizer $\omega_{k}^{\star}$. Moreover,

$$ $\operatorname{Var}[g_{\omega_{k}^{\star},k}]\;<\;\min\!\big\{\operatorname{Var}[g_{0,k}],\ \operatorname{Var}[g_{1,k}]\big\},$ $$

A proof for the above proposition is provided in [Section˜C.1](https://arxiv.org/html/2510.09541v2#A3.SS1). A few remarks are in order:

- •
Confidence-aware weighting: The mixture gradient in [Equation˜9](https://arxiv.org/html/2510.09541v2#S3.E9) realizes a confidence-aware weighting: uncertain tokens with small $\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$, indicating a low recovery chance, have a smaller weight, while confident tokens with large $\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$ are upweighted. The sharpness is controlled by parameter $\beta$ and the blend by $\omega$. Furthermore, the convex interpolation of the confidence-aware coefficient of the upper bound with the lower bound ensures clipping tiny gradients to a minimum value and thus prevents vanishing gradients.
- •
Lower variance and more stable training: According to [Proposition˜1](https://arxiv.org/html/2510.09541v2#S3.Ex4), the gradient of the optimal mixture, i.e., $g_{\omega_{k}^{\star},k}$, has strictly smaller coordinate-wise variance than the gradient of either the lower bound ($g_{0,k}$) or the upper bound ($g_{1,k}$)(^1^11[Proposition 1](https://arxiv.org/html/2510.09541v2#S3.Ex4) extends directly to a single, coordinate-independent optimizer $\omega^{\star}$ obtained by minimizing the sum of coordinate-wise variances.). In our experiments, we fix $\beta$ and $\omega$ as hyperparameters for simplicity. These values can also be adaptively adjusted during training to better match the evolving training dynamics and data distribution.

Thus, the mixture approach offers theoretical advantages over using either the upper or lower bound alone, as supported by our experimental results in [Section˜4](https://arxiv.org/html/2510.09541v2#S4). Further discussions of the mixture approach and empirical evidence of reduced gradient variance are provided in Appendix [C.2](https://arxiv.org/html/2510.09541v2#A3.SS2) and [Figure˜7](https://arxiv.org/html/2510.09541v2#A3.F7), and Appendix [C.3](https://arxiv.org/html/2510.09541v2#A3.SS3) presents a toy example illustrating the distinct behaviors of the lower and upper bounds.

###### Proposition 1 (Optimal Mixture Strictly Reduces Variance) .

## 4 Experiments

In this section, we present experimental results highlighting the superior performance of SPG across various benchmarks. Further, we provide detailed analysis and ablations of SPG to assess the contribution of each component, examine the influence of key hyperparameters, and evaluate the robustness of our approach under different inference strategies.

### 4.1 Experimental Setup and Main Results

#### Experimental Setup.

We conduct RL fine-tuning with SPG following the experimental settings in D1 (Zhao et al., 2025) and WD1 (Tang et al., 2025). We employ LLaDA-8B-Instruct (Nie et al., 2025), a state-of-the-art open-sourced dLLM without post-training, as the base model, and experiment on four benchmarks: two for mathematical reasoning (GSM8K (Cobbe et al., 2021) and MATH500 (Lightman et al., 2023)) and two for logical reasoning (Countdown (Pan et al., 2025) and Sudoku (Arel, 2025)). We follow the same train-test splitting, reward functions, and evaluation protocol as D1 and WD1, except for Sudoku. For Sudoku, to avoid train-test leakage, we take the training set from D1 and split the data by Sudoku answers, ensuring that the test set contains entirely new puzzle solutions. This guarantees that the model cannot solve test puzzles merely by memorizing possible answers.
All experiments are conducted in the zero-shot setting, except for Sudoku, where 3-shot generation is used for both training and evaluation(^2^22We use 3-shot generation for Sudoku because zero-shot is too difficult for this task, resulting in very few meaningful RL rollouts. Few-shot examples used in our experiments are provided in [Section D.3](https://arxiv.org/html/2510.09541v2#A4.SS3).).
For all models, we employ Low-Rank Adaptation (LoRA) with a rank of $r=128$ and scaling factor $\alpha=64$. For SPG, we report results using both $\tilde{\mathcal{L}}_{\text{EUBO}}$ (i.e., SPG w/ EUBO) and $\tilde{\mathcal{L}}_{\text{Mix}}$ (i.e., SPG w/ Mixture) for negative advantage traces. We select the value of $\beta$ in the EUBO from $\{1.0,1.5,2.0\}$ based on the best average test accuracy across all generation lengths, and fix the mixture coefficient $\omega$ at 0.5.
Further experimental details are in [Section˜D.1](https://arxiv.org/html/2510.09541v2#A4.SS1) and [Section˜D.2](https://arxiv.org/html/2510.09541v2#A4.SS2).

#### Baselines.

We compare our method with several recent RL algorithms for dLLMs, including D1 (Zhao et al., 2025), WD1 (Tang et al., 2025), and UniGRPO (Yang et al., 2025). For D1 and WD1, we reproduce results using the official codebases and instructions, and for fair comparison, we omit the additional SFT stage in D1 across all models. For UniGRPO, since the code is not publicly available and the original work focuses on vision-language multimodal models, we reimplement the algorithm within our setup. For consistency, we set the number of inner gradient updates $\mu$ to 4 for all models, following GRPO (Shao et al., 2024). We also evaluate LLaDA-1.5 (Zhu et al., 2025) under our settings, which fine-tune LLaDA-8B-Instruct using VRPO, a preference optimization approach on 350K preference pairs.

Table: Table 1: Model performance on four reasoning benchmarks. The best results are bolded and the second best are underlined. SPG consistently outperforms all other methods. We denote the absolute gain of test accuracy to the previous state-of-the-art in green.

Figure: Figure 3: Reward dynamics of SPG w/ Mixture during RL training, compared with D1, WD1, and UniGRPO. SPG consistently leads to faster convergence and higher reward level. We report mean and standard deviation over a rolling window of 50 steps.
Refer to caption: x2.png

#### Generation and Evaluation Setup.

For both RL rollouts and evaluation, we use the semi-autoregressive confidence-based decoding strategy, following LLaDA, D1 and WD1. We apply the same generation setup as D1, with the denoising timestep set to half the total sequence length. The sequence is divided into blocks of 32 tokens, and in diffusion step, we unmask the 2 tokens with the highest confidence (measured by the probability of the sampled token) within the current incomplete block. During RL rollout, to encourage diverse outputs, we use a generation length of 256 and a sampling temperature of 0.9 across all benchmarks, except for sudoku, where the temperature is set to 0.3 as in D1. During evaluation, the sampling temperature is set to 0.0. We evaluate the models every 100 steps, reporting results from the checkpoint that achieves the highest average test accuracy across generation lengths of 128, 256, and 512.

#### Results.

We provide the performance of SPG on each benchmark in comparison to the base model and other baselines in [Table˜1](https://arxiv.org/html/2510.09541v2#S4.T1). Both SPG w/ EUBO and SPG w/ Mixture consistently achieve significant improvements over the baselines across all tasks and generation lengths, with the Mixture approach that combines ELBO and EUBO for negative advantage traces yielding the best performance. In particular, at a generation length of 256, SPG w/ Mixture improves the test accuracy over the previous state-of-the-art by 3.6% on GSM8K, 2.6% on MATH500, 18% on Countdown, and 27% on Sudoku, showcasing the effectiveness of SPG to conduct RL for dLLMs.
Reward dynamics throughout training are illustrated in [Figure˜3](https://arxiv.org/html/2510.09541v2#S4.F3), where SPG shows a rapid and steady increase in reward over the optimization steps, further demonstrating its efficiency and robustness.
We provide additional results and comparisons to the baselines in [Table˜4](https://arxiv.org/html/2510.09541v2#A5.T4) and [Section˜E.1](https://arxiv.org/html/2510.09541v2#A5.SS1).

### 4.2 Ablations and Further Analysis

We conduct a series of ablation studies to gain deeper insights from the following aspects:

- •
The contribution of each individual component, including log-likelihood estimation methods for negative advantage traces ([Section˜4.2](https://arxiv.org/html/2510.09541v2#S4.SS2.SSS0.Px1)) and the masking strategy in Monte Carlo estimation ([Table˜3](https://arxiv.org/html/2510.09541v2#S4.T3)).
- •
The effect of key hyperparameters, including $\beta$ that controls the tightness of the upper bound and the mixture coefficient $\omega$ ([Figure˜5](https://arxiv.org/html/2510.09541v2#S4.F5)).
- •
The robustness of our approach under various inference strategies ([Figure˜6](https://arxiv.org/html/2510.09541v2#S4.F6)).

Due to computational constraints, some ablation experiments are conducted on a representative mathematical reasoning benchmark (MATH500) and a logical reasoning benchmark (Countdown). Unless otherwise noted, we report average test accuracy across generation lengths 128, 256, and 512 for the ablation studies, with detailed results for each generation length provided in [Section˜E.2](https://arxiv.org/html/2510.09541v2#A5.SS2).
In [Section˜E.2](https://arxiv.org/html/2510.09541v2#A5.SS2), we also investigate alternative log-likelihood estimation methods for positive advantage traces in place of ELBO, as detailed in [Table˜11](https://arxiv.org/html/2510.09541v2#A5.T11), and study the diversity of model generations by evaluating the pass@K performance of each model in [Table˜12](https://arxiv.org/html/2510.09541v2#A5.T12).

Figure: Figure 4: Reward dynamics of different log-likelihood estimation methods for negative advantage traces on Sudoku. SPG w/ Mixture leads to both fast convergence and high rewards.
Refer to caption: x3.png

#### Ablations on Algorithm Components.

We first study the impact of different log-likelihood estimation methods for negative advantage traces in [Section˜4.2](https://arxiv.org/html/2510.09541v2#S4.SS2.SSS0.Px1). Specifically, we compare our approach using $\tilde{\mathcal{L}}_{\text{EUBO}}$ or $\tilde{\mathcal{L}}_{\text{Mix}}$ with those using $\mathcal{L}_{\text{ELBO}}$ (SPG w/ ELBO) or omitting the negative advantage loss entirely (SPG wo/ neg).
Removing the negative advantage loss results in a substantial performance drop, highlighting the importance of negative advantage penalties to RL. Additionally, both Mixture and EUBO methods outperform ELBO (except for EUBO in MATH500), showcasing the benefits of evidence upper bound regularization for negative rewards. We provide complete results for each generation length in [Table˜6](https://arxiv.org/html/2510.09541v2#A5.T6).

The effect of log-likelihood estimation methods is further illustrated by the reward dynamics of each model in [Figure˜4](https://arxiv.org/html/2510.09541v2#S4.F4), taking Sudoku as an example. SPG w/ ELBO converges rapidly during training but plateaus early, as minimizing the lower bound does not necessarily minimize the true log-likelihood for negative advantage traces. In contrast, SPG w/ EUBO achieves higher final rewards but converges more slowly and less stably. Combining both, SPG w/ Mixture attains fast, stable convergence and high rewards, leading to an effective balance. This aligns with our discussions in [Section˜3.3](https://arxiv.org/html/2510.09541v2#S3.SS3.SSS0.Px2).

We also conduct ablations on the masking strategies in Monte Carlo estimation of $\mathcal{L}_{\text{ELBO}}$, $\tilde{\mathcal{L}}_{\text{EUBO}}$, and $\tilde{\mathcal{L}}_{\text{Mix}}$. As shown in [Table˜3](https://arxiv.org/html/2510.09541v2#S4.T3), the block-wise masking strategy outperforms random masking, demonstrating the importance of aligning input distributions between policy rollout and optimization. We provide complete results for each generation length in [Table˜7](https://arxiv.org/html/2510.09541v2#A5.T7).

Table: Table 2: Ablations on log-likelihood estimation methods for negative advantage traces. The best results are bolded and the second best underlined. We denote the absolute gain of test accuracy to SPG w/ ELBO in green. SPG w/ Mixture consistently outperforms other likelihood estimation methods.

Figure: Figure 5: *(a)-(d)*: ablations on the effect of $\beta$ in the upper bound; *(e)-(f)*: ablations on the mixture coefficient $\omega$. The best performed $\beta\geq 1$ and $\omega\in[0,1]$ are marked by triangle in each setting.
Refer to caption: x4.png

#### Ablations on Key Hyperparameters β \beta and ω \omega .

We first examine the effect of $\beta$, a crucial hyperparameter in evidence upper bound estimation, in panels (a)-(d) of [Figure˜5](https://arxiv.org/html/2510.09541v2#S4.F5). In general, a relatively small value of $\beta$ (i.e., close to 1.0) leads to a tighter bound and thus better performance. Nevertheless, SPG consistently performs well across a range of $\beta$ values on most tasks, indicating its robustness. For our main results in [Table˜1](https://arxiv.org/html/2510.09541v2#S4.T1), we fix $\omega=0.5$ and select the optimal $\beta\geq 1$, resulting in $\beta=1.0$ for Sudoku and $\beta=1.5$ for the other three benchmarks, except for Countdown with SPG w/ EUBO where $\beta=2.0$.
Besides, since the ELBO corresponds to the case of $\beta=0$ theoretically and EUBO corresponds to $\beta\geq 1$, we also investigate intermediate values $0<\beta<1$, which may serve as an implicit mixture of lower and upper bounds. However, it is unstable in Sudoku and underperform SPG w/ Mixture on most benchmarks.

We also experiment on the effect of the mixture coefficient $\omega$, keeping $\beta$ fixed at its optimal value determined for $\omega=0.5$ as mentioned before. As illustrated in panels (e)-(f) of [Figure˜5](https://arxiv.org/html/2510.09541v2#S4.F5), combining lower and upper bounds with $\omega\in(0,1)$ leads to better performance than leveraging either bound solely, resulting in an inverted U-shaped curve. This observation is consistent with our analysis in [Proposition˜1](https://arxiv.org/html/2510.09541v2#S3.Ex4) and [Section˜3.3](https://arxiv.org/html/2510.09541v2#S3.SS3.SSS0.Px2). We provide complete ablation results of $\beta$ and $\omega$ for each generation length in [Table˜8](https://arxiv.org/html/2510.09541v2#A5.T8) and [Table˜9](https://arxiv.org/html/2510.09541v2#A5.T9).

#### Ablations on Inference Strategies.

In the above experiments, we adopt a consistent state-of-the-art inference setup during both RL rollout and evaluation, i.e., confidence-based, block-wise semi-autoregressive generation with a block size of 32. The same configuration and block size are also used in our block-wise masking strategy. This raises the question of whether our approach generalizes well to alternative inference strategies. To assess this, we evaluate the base model and all RL fine-tuned models using various inference strategies, as shown in [Figure˜6](https://arxiv.org/html/2510.09541v2#S4.F6). Despite being trained under confidence-based semi-AR decoding, SPG consistently outperforms all baselines by a substantial margin across all inference strategies, demonstrating its robustness and strong generalizability. Complete results for each benchmark individually are provided in [Table˜13](https://arxiv.org/html/2510.09541v2#A5.T13).

Figure: Figure 6: Ablations on inference strategies, including different combinations of decoding orders (i.e., semi-autoregressive (semi-AR) decoding with varying block sizes and full sequence decoding) and unmasking approaches (i.e., confidence-based and random unmasking). We set generation length to 256 and report the average accuracy across four benchmarks. SPG consistently outperforms all baselines by a large margin across different inference strategies.
Refer to caption: x5.png

## 5 Related Work

#### Diffusion Language Models.

Building on the remarkable success of diffusion models for image generation in continuous domains (Song et al., 2020; Ho et al., 2020), researchers have explored their extension to discrete data such as text. Initial attempts focused on training continuous diffusion models in the text embedding space (Li et al., 2022; Gong et al., 2022; Han et al., 2022; Sahoo et al., 2025a), while they face challenges in optimization and generalization due to the discrete nature of text data. Masked diffusion models (Lou et al., 2023; Zheng et al., 2023; Campbell et al., 2024; Sahoo et al., 2024; Shi et al., 2024) address this by defining the diffusion process directly in the discrete token space, using random masking as the forward process, and have achieved strong empirical results. Block Diffusion (Arriola et al., 2025) further advances this direction by combining the strengths of autoregressive models, such as the capability to generate variable-length outputs and using KV cache to accelerate inference, with the benefits of diffusion language models like parallel decoding and flexible, any-order generation within blocks. Recently, large-scale diffusion language models trained with masked diffusion objectives have demonstrated performance competitive with similarly sized autoregressive models (Nie et al., 2025; Gong et al., 2024). More recent works (Wu et al., 2025; Ma et al., 2025; Liu et al., 2025a; Sahoo et al., 2025a, b) have introduced caching and parallel decoding algorithms that greatly enhance the inference efficiency of dLLMs.

#### Reinforcement Learning for LLMs and Reasoning.

The seminal works apply reinforcement learning to large language models (LLMs) to align them with human preferences via reinforcement learning from human feedback (RLHF) (Christiano et al., 2017; Ouyang et al., 2022). More recently, reinforcement learning has proven highly effective at enhancing the reasoning abilities of LLMs during the post-training stage, where rewards can be provided by a process reward model (Lightman et al., 2023) or verifiable reward signals. Algorithms such as Proximal Policy Optimization (PPO) and Trust Region Policy Optimization (TRPO) constrain policy updates to a trust region, reducing variance and promoting stable learning by preventing excessive shifts from the reference policy (Schulman et al., 2015, 2017). Group Relative Policy Optimization (GRPO) (Shao et al., 2024) introduces group-relative rewards, enabling efficient training without the need for an additional value (critic) model. GRPO and its variants have demonstrated strong empirical performance in state-of-the-art models such as DeepSeek-R1 (Guo et al., 2025), particularly on mathematical reasoning tasks, where incorporating long reasoning traces with self-reflection and verification steps yields significant improvements. Recent works (Liu et al., 2025b; Zheng et al., 2025; Team et al., 2025; Cohen et al., 2025) further improve RL algorithms for LLMs by reducing the bias introduced by the GRPO objective, enhancing sample efficiency, and introducing additional regularization.

#### Reinforcement Learning for Diffusion Language Models.

Numerous studies have explored RL-based fine-tuning algorithms for diffusion models with continuous objectives (Fan et al., 2023; Black et al., 2023; Clark et al., 2023). While RL algorithms have achieved notable success to LLMs and continuous diffusion models, their applications to diffusion language models in the discrete space remain underexplored.
DRAKES (Wang et al., 2024) leverages reward backpropagation along the denoising trajectory, but is computationally intensive for large scale models as the gradients are propagated through each denoising step. Alternatively, methods like D1 (Zhao et al., 2025) and UniGRPO Yang et al. (2025) utilize the GRPO framework, approximating the log-likelihood through either a one-step unmasking (as in D1) or Monte Carlo estimation using the ELBO (as in UniGRPO). VRPO (Zhu et al., 2025) adapts DPO (Rafailov et al., 2023) to fine-tune dLLMs by applying MC estimation of the ELBO. WD1 (Tang et al., 2025) starts from the GRPO formulation and the same log-likelihood estimation as in D1, while avoiding direct estimation of the old and reference policy log-likelihoods by integrating them into a weighted policy optimization objective. Despite these advances, a principled analysis of RL algorithms for dLLMs, especially the challenging log-likelihood estimation, is missing. This results in substantial bias in the optimization objective and suboptimal performance.

## 6 Conclusion

We propose SPG, a novel reinforcement learning algorithm for diffusion large language models. SPG addresses the intractable log-likelihood in dLLMs by maximizing a tractable lower bound on positive reward sequences and minimizing an upper bound on negative ones, resulting in a more robust and less biased policy gradient. Additionally, we propose a block-wise masking strategy for Monte Carlo estimation to enhance optimization stability and efficiency.
Extensive experiments on four mathematical and logical reasoning benchmarks demonstrate the superior performance of SPG, achieving significant improvement over baselines and the state-of-the-art performance.

## References

- Arel (2025)
Arel.
Arel’s sudoku generator.
https://www.ocf.berkeley.edu/ arel/sudoku/main.html, 2025.
- Arriola et al. (2025)
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov.
Block diffusion: Interpolating between autoregressive and diffusion language models.
*arXiv preprint arXiv:2503.09573*, 2025.
- Austin et al. (2021)
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg.
Structured denoising diffusion models in discrete state-spaces.
*Advances in neural information processing systems*, 34:17981–17993, 2021.
- Black et al. (2023)
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine.
Training diffusion models with reinforcement learning.
*arXiv preprint arXiv:2305.13301*, 2023.
- Campbell et al. (2022)
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet.
A continuous time framework for discrete denoising models.
*Advances in Neural Information Processing Systems*, 35:28266–28279, 2022.
- Campbell et al. (2024)
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola.
Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design.
*arXiv preprint arXiv:2402.04997*, 2024.
- Christiano et al. (2017)
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei.
Deep reinforcement learning from human preferences.
*Advances in neural information processing systems*, 30, 2017.
- Clark et al. (2023)
Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet.
Directly fine-tuning diffusion models on differentiable rewards.
*arXiv preprint arXiv:2309.17400*, 2023.
- Cobbe et al. (2021)
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al.
Training verifiers to solve math word problems.
*arXiv preprint arXiv:2110.14168*, 2021.
- Cohen et al. (2025)
Taco Cohen, David W Zhang, Kunhao Zheng, Yunhao Tang, Remi Munos, and Gabriel Synnaeve.
Soft policy optimization: Online off-policy rl for sequence models.
*arXiv preprint arXiv:2503.05453*, 2025.
- Dao (2023)
Tri Dao.
Flashattention-2: Faster attention with better parallelism and work partitioning.
*arXiv preprint arXiv:2307.08691*, 2023.
- Fan et al. (2023)
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee.
Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models.
*Advances in Neural Information Processing Systems*, 36:79858–79885, 2023.
- Gong et al. (2022)
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong.
Diffuseq: Sequence to sequence text generation with diffusion models.
*arXiv preprint arXiv:2210.08933*, 2022.
- Gong et al. (2024)
Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, et al.
Scaling diffusion language models via adaptation from autoregressive models.
*arXiv preprint arXiv:2410.17891*, 2024.
- Guo et al. (2025)
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al.
Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.
*arXiv preprint arXiv:2501.12948*, 2025.
- Han et al. (2022)
Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov.
Ssd-lm: Semi-autoregressive simplex-based diffusion language model for text generation and modular control.
*arXiv preprint arXiv:2210.17432*, 2022.
- Ho et al. (2020)
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Labs et al. (2025)
Inception Labs, Samar Khanna, Siddhant Kharbanda, Shufan Li, Harshit Varma, Eric Wang, Sawyer Birnbaum, Ziyang Luo, Yanis Miraoui, Akash Palrecha, et al.
Mercury: Ultra-fast language models based on diffusion.
*arXiv preprint arXiv:2506.17298*, 2025.
- Li et al. (2022)
Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto.
Diffusion-lm improves controllable text generation.
*Advances in neural information processing systems*, 35:4328–4343, 2022.
- Lightman et al. (2023)
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe.
Let’s verify step by step.
In *The Twelfth International Conference on Learning Representations*, 2023.
- Liu et al. (2025a)
Zhiyuan Liu, Yicun Yang, Yaojie Zhang, Junjie Chen, Chang Zou, Qingyan Wei, Shaobo Wang, and Linfeng Zhang.
dllm-cache: Accelerating diffusion large language models with adaptive caching.
*github*, 2025a.
- Liu et al. (2025b)
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin.
Understanding r1-zero-like training: A critical perspective.
*arXiv preprint arXiv:2503.20783*, 2025b.
- Loshchilov and Hutter (2017)
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
*arXiv preprint arXiv:1711.05101*, 2017.
- Lou et al. (2023)
Aaron Lou, Chenlin Meng, and Stefano Ermon.
Discrete diffusion modeling by estimating the ratios of the data distribution.
*arXiv preprint arXiv:2310.16834*, 2023.
- Ma et al. (2025)
Xinyin Ma, Runpeng Yu, Gongfan Fang, and Xinchao Wang.
dkv-cache: The cache for diffusion language models.
*arXiv preprint arXiv:2505.15781*, 2025.
- Nie et al. (2025)
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li.
Large language diffusion models.
*arXiv preprint arXiv:2502.09992*, 2025.
- Ouyang et al. (2022)
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.
Training language models to follow instructions with human feedback.
*Advances in neural information processing systems*, 35:27730–27744, 2022.
- Pan et al. (2025)
Jiayi Pan, Junjie Zhang, Xingyao Wang, Lifan Yuan, Hao Peng, and Alane Suhr.
Tinyzero.
https://github.com/Jiayi-Pan/TinyZero, 2025.
Accessed: 2025-01-24.
- Peebles and Xie (2023)
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In *Proceedings of the IEEE/CVF international conference on computer vision*, pages 4195–4205, 2023.
- Rafailov et al. (2023)
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn.
Direct preference optimization: Your language model is secretly a reward model.
*Advances in neural information processing systems*, 36:53728–53741, 2023.
- Rényi (1961)
Alfréd Rényi.
On measures of entropy and information.
In *Proceedings of the fourth Berkeley symposium on mathematical statistics and probability, volume 1: contributions to the theory of statistics*, volume 4, pages 547–562. University of California Press, 1961.
- Sahoo et al. (2024)
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov.
Simple and effective masked diffusion language models.
*Advances in Neural Information Processing Systems*, 37:130136–130184, 2024.
- Sahoo et al. (2025a)
Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin Chiu, and Volodymyr Kuleshov.
The diffusion duality.
*arXiv preprint arXiv:2506.10892*, 2025a.
- Sahoo et al. (2025b)
Subham Sekhar Sahoo, Zhihan Yang, Yash Akhauri, Johnna Liu, Deepansha Singh, Zhoujun Cheng, Zhengzhong Liu, Eric Xing, John Thickstun, and Arash Vahdat.
Esoteric language models.
*arXiv preprint arXiv:2506.01928*, 2025b.
- Schulman et al. (2015)
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz.
Trust region policy optimization.
In *International conference on machine learning*, pages 1889–1897. PMLR, 2015.
- Schulman et al. (2017)
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
Proximal policy optimization algorithms.
*arXiv preprint arXiv:1707.06347*, 2017.
- Shao et al. (2024)
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al.
Deepseekmath: Pushing the limits of mathematical reasoning in open language models.
*arXiv preprint arXiv:2402.03300*, 2024.
- Shi et al. (2024)
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias.
Simplified and generalized masked diffusion for discrete data.
*Advances in neural information processing systems*, 37:103131–103167, 2024.
- Song et al. (2020)
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
*arXiv preprint arXiv:2011.13456*, 2020.
- Sun et al. (2025)
Bowen Sun, Yujun Cai, Ming-Hsuan Yang, and Yiwei Wang.
Blockwise sft for diffusion language models: Reconciling bidirectional attention and autoregressive decoding.
*arXiv preprint arXiv:2508.19529*, 2025.
- Sun et al. (2022)
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai.
Score-based continuous-time discrete diffusion models.
*arXiv preprint arXiv:2211.16750*, 2022.
- Tang et al. (2025)
Xiaohang Tang, Rares Dolga, Sangwoong Yoon, and Ilija Bogunovic.
wd1: Weighted policy optimization for reasoning in diffusion language models.
*arXiv preprint arXiv:2507.08838*, 2025.
- Team et al. (2025)
Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al.
Kimi k1. 5: Scaling reinforcement learning with llms.
*arXiv preprint arXiv:2501.12599*, 2025.
- Van Erven and Harremos (2014)
Tim Van Erven and Peter Harremos.
Rényi divergence and kullback-leibler divergence.
*IEEE Transactions on Information Theory*, 60(7):3797–3820, 2014.
- Wang et al. (2024)
Chenyu Wang, Masatoshi Uehara, Yichun He, Amy Wang, Tommaso Biancalani, Avantika Lal, Tommi Jaakkola, Sergey Levine, Hanchen Wang, and Aviv Regev.
Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design.
*arXiv preprint arXiv:2410.13643*, 2024.
- Wang et al. (2025a)
Xu Wang, Chenkai Xu, Yijie Jin, Jiachun Jin, Hao Zhang, and Zhijie Deng.
Diffusion llms can do faster-than-ar inference via discrete diffusion forcing.
*arXiv preprint arXiv:2508.09192*, 2025a.
- Wang et al. (2025b)
Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, and Mengdi Wang.
Revolutionizing reinforcement learning framework for diffusion large language models.
*arXiv preprint arXiv:2509.06949*, 2025b.
- Wu et al. (2025)
Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie.
Fast-dllm: Training-free acceleration of diffusion llm by enabling kv cache and parallel decoding.
*arXiv preprint arXiv:2505.22618*, 2025.
- Yang et al. (2025)
Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang.
Mmada: Multimodal large diffusion language models.
*arXiv preprint arXiv:2505.15809*, 2025.
- Zhao et al. (2025)
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover.
d1: Scaling reasoning in diffusion large language models via reinforcement learning.
*arXiv preprint arXiv:2504.12216*, 2025.
- Zheng et al. (2025)
Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, et al.
Group sequence policy optimization.
*arXiv preprint arXiv:2507.18071*, 2025.
- Zheng et al. (2023)
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong.
A reparameterized discrete diffusion model for text generation.
*arXiv preprint arXiv:2302.05737*, 2023.
- Zhu et al. (2025)
Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, et al.
Llada 1.5: Variance-reduced preference optimization for large language diffusion models.
*arXiv preprint arXiv:2505.19223*, 2025.

## Appendix A Basics of dLLMs

In this section, we provide a more self-contained overview of masked dLLMs. Please also refer to Sahoo et al. (2024) for more details.

#### Notation.

We denote scalars by lowercase letters ($x$), vectors by bold lowercase ($\bm{x}$), and sequences by $\bm{x}_{1:n}$. A superscript (e.g., $\bm{x}^{j}$) denotes an item’s index within a group. We define the set of the first $k$ integers as $[k]\vcentcolon=\{1,\dots,k\}$ and the $k$-dimensional probability simplex as $\Delta^{k-1}$. Distributions include the categorical $\text{Cat}(\cdot\mid\bm{p})$ and the uniform $\mathcal{U}[a,b]$. Throughout the paper, we use the following primary indices: $i\in[n]$ for position, $j\in[g]$ for a sequence in a group, and $t\in[0,1]$ for the continuous diffusion timestep.

We start from a discrete time version of the diffusion models with finite $t\in[T]$. Assume a one-hot categorical variable $\bm{x}\in\{e_{1},\dots,e_{k}\}\subset\Delta^{k-1}$.
Further assume we gradually corrupt $\bm{x}$ into an absorbing state $\bm{m}$ (i.e., $e_{[\texttt{mask}]\penalty 10000\ }$) with transition matrix $\bm{Q}_{t}$ at time $t$. Then:

$$ $q(\bm{z}_{t}\mid\bm{x})=\text{Cat}(\bm{z}_{t}\mid\overline{\bm{Q}_{t}}\bm{x})=\text{Cat}(\bm{z}_{t}\mid\prod_{\tau=1}^{t}\bm{Q}_{\tau}\bm{x}).$ $$

Here, $\bm{z}_{t}$ is also a one-hot categorical random variable in $\Delta^{k-1}$. In practice, one could choose $\bm{Q}_{t}$ such that:

$$ $q(\bm{z}_{t}\mid\bm{x})=\text{Cat}(\bm{z}_{t}\mid\alpha_{t}\bm{x}+(1-\alpha_{t})\bm{m}).$ $$

Here, $\alpha_{1}=1,\alpha_{T}=0,\alpha^{\prime}_{t}<0$.

Normally, the goal is to construct the lower bound of the evidence (ELBO) and maximize it. For this particular case, consider the discretized Markov chain with $T$ latent variables $\bm{z}_{1},\bm{z}_{2},\dots,\bm{z}_{T}$, where $\bm{z}_{T}=\bm{m}$ and $\bm{z}_{1}=\bm{x}$. We use the shorthand $\bm{z}=\bm{z}_{1:T}$ and write

$$ $\begin{split}\mathcal{L}_{\text{ELBO}}(\bm{x};\bm{{\bm{\theta}}})&=\mathbb{E}_{\bm{z}\sim q(\cdot\mid\bm{x})}\bigg[\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{x},\bm{z})}{q(\bm{z}\mid\bm{x})}\bigg]\\ =&\mathbb{E}_{\bm{z}\sim q(\cdot\mid\bm{x})}\bigg[\underbrace{\log p_{\bm{{\bm{\theta}}}}(\bm{x}\mid\bm{z}_{1})}_{=0}+\sum_{t=1}^{T-1}\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}\mid\bm{z}_{t+1})}{q(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x})}+\underbrace{\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{T})}{q(\bm{z}_{T}\mid\bm{x})}}_{=0}\bigg]\\ =&\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t},\bm{z}_{t+1}\sim q}\bigg[\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}\mid\bm{z}_{t+1})}{q(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x})}\bigg]\\ =&\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\mathbb{E}_{\bm{z}_{t}\sim q(\cdot\mid\bm{z}_{t+1},\bm{x})}\bigg[\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}\mid\bm{z}_{t+1})}{q(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x})}\bigg].\end{split}$ (10) $$

Here, $\log p_{\bm{{\bm{\theta}}}}(\bm{x},\bm{z}_{1})=0$ because we assume $\bm{z}_{1}=\bm{x}$, and $p_{\bm{{\bm{\theta}}}}(\bm{z}_{T})=q(\bm{z}_{T}\mid\bm{x})$ because we assume $\bm{z}_{T}=\bm{m}$. A common method to parameterize $p_{\bm{{\bm{\theta}}}}$ is via predicting $\bm{x}$ with model $\pi_{\bm{\theta}}$ in $q$:

$$ $p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}\mid\bm{z}_{t+1})=q\left(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x}=\pi_{\bm{\theta}}(\cdot\mid\bm{z}_{t+1})\right).$ $$

Now, given that $\bm{z}_{t+1}$ is either $\bm{m}$ or $\bm{x}$ (assuming $\bm{m}\neq\bm{x}$). Then the KL term in equation [10](https://arxiv.org/html/2510.09541v2#A1.E10) decomposes into the following.

$$ $\begin{split}\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}\mid\bm{z}_{t+1})}{q(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x})}=\begin{cases}0&\bm{z}_{t}=\bm{z}_{t+1}=\bm{x},\\ 0&\bm{z}_{t}=\bm{m},\bm{z}_{t+1}=\bm{x},\qquad{\small{\color[rgb]{1,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,1}\pgfsys@color@cmyk@stroke{0}{1}{0}{0}\pgfsys@color@cmyk@fill{0}{1}{0}{0}\text{(Impossible)}}}\\ \log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t+1})&\bm{z}_{t}=\bm{x},\bm{z}_{t+1}=\bm{m},\\ 0&\bm{z}_{t}=\bm{z}_{t+1}=\bm{m}.\\ \end{cases}\end{split}$ (11) $$

Moreover, $q(\bm{z}_{t}=\bm{x}\mid\bm{z}_{t+1}=\bm{m},\bm{x})=\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}$, and note that $\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t})=1$
when $\bm{z}_{t}=\bm{x}$, so we have:

$$ $\begin{split}\mathcal{L}_{\text{ELBO}}(\bm{x};\bm{{\bm{\theta}}})&=\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\bigg[\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t+1})\mathbbm{1}(\bm{z}_{t+1}=\bm{m})\bigg]\\ &=\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\bigg[\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t+1})\bigg].\penalty 10000\ \penalty 10000\ \penalty 10000\ {\small{\color[rgb]{1,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,1}\pgfsys@color@cmyk@stroke{0}{1}{0}{0}\pgfsys@color@cmyk@fill{0}{1}{0}{0}\text{(If $\bm{z}_{t+1}=\bm{x}$, then $\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t+1})=0$)}}}\end{split}$ (12) $$

Taking the above limit as $T\rightarrow\infty$, we have:

$$ $\mathcal{L}_{\text{ELBO}}(\bm{x};\bm{{\bm{\theta}}})=\int_{t=0}^{1}\mathbb{E}_{\bm{z}_{t}\sim q(\cdot\mid\bm{x})}\bigg[\frac{\alpha^{\prime}_{t}}{\alpha_{t}-1}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t})\bigg].$ (13) $$

#### Generalization to Sequence

The above is for a single categorical variable $\bm{x}$. In practice as in language modeling, it becomes a sequence of categorical variables $\bm{x}_{1:n}$. Then we write

$$ $\begin{split}\mathcal{L}_{\text{ELBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})&=\mathbb{E}_{\bm{z}_{1:n}\sim q(\cdot\mid\bm{x}_{1:n})}\bigg[\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{x}_{1:n},\bm{z}_{1:n})}{q(\bm{z}_{1:n}\mid\bm{x}_{1:n})}\bigg]\\ &=\mathbb{E}_{\{\bm{z}_{i}\sim q(\cdot\mid\bm{x}_{i})\}_{i=1}^{n}}\bigg[\sum_{i=1}^{n}\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{x}_{i},\bm{z}_{1:n})}{q(\bm{z}_{i}\mid\bm{x}_{i})}\bigg]\qquad{\small{\color[rgb]{1,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,1}\pgfsys@color@cmyk@stroke{0}{1}{0}{0}\pgfsys@color@cmyk@fill{0}{1}{0}{0}\text{(Independence of $q(\cdot\mid\bm{x}_{i})$)}}}\\ &=\sum_{i=1}^{n}\mathbb{E}_{\{\bm{z}_{i^{\prime}}\sim q(\cdot\mid\bm{x}_{i^{\prime}})\}_{i^{\prime}=1}^{n}}\bigg[\log\frac{p_{\bm{{\bm{\theta}}}}(\bm{x}_{i},\bm{z}_{1:n})}{q(\bm{z}_{i}\mid\bm{x}_{i})}\bigg]\\ &=\sum_{i=1}^{n}\mathcal{L}_{\text{ELBO}}(\bm{x}_{i};\bm{{\bm{\theta}}}).\end{split}$ (14) $$

The key distinction from the single-token formulation (mentioned beforehand) is that the reverse process $p_{\bm{{\bm{\theta}}}}$ is conditioned on all $\bm{z}_{1:n}$ instead of a single token’s $\bm{z}_{i}$.

## Appendix B Evidence Upper Bound for dLLMs

In this section, we provide the derivation of the evidence upper bound. Following the above section, we start from the discrete time version of the diffusion models.

###### Lemma 1 (Rényi Variational Bound; Rényi ( 1961 ); Van Erven and Harremos ( 2014 ) ) .

Fix an observation $x$. Let $q(\cdot\mid x)$ be any distribution on $\mathcal{Z}$ such that $p(\cdot\mid x)\ll q(\cdot\mid x)$, denoting that $p(\cdot\mid x)$ is absolutely continuous with respect to $q(\cdot\mid z)$. Then, the following holds for any $\beta\geq 1$:

$$ $\mathbb{E}_{z\sim q(\cdot\mid x)}\left[\log\frac{p(x,z)}{q(z\mid x)}\right]\leq\log p(x)\leq\frac{1}{\beta}\log\mathbb{E}_{z\sim q(\cdot\mid x)}\left[\bigg(\frac{p(x,z)}{q(z\mid x)}\bigg)^{\beta}\right].$ (15) $$

In view of the above lemma, we derive an evidence upper bound for masked diffusion models in the following theorem.

###### Theorem 1 (Evidence Upper Bound for Masked Diffusion) .

Assume the forward denoising process has $T$ steps with a monotonic schedule $\alpha_{t}$.
For any $\beta\geq 1$ and a sequence of categorical variables $\bm{x}_{1:n}$, we have:

$$ $\log\pi_{{\bm{\theta}}}(\bm{x}_{1:n})\leq\mathcal{L}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}}),$ (16) $$

where

$$ $\mathcal{L}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})\vcentcolon=\frac{1}{\beta}\sum_{i=1}^{n}\log\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}}\Bigg[\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\cdot\mathbbm{1}(\bm{z}_{t+1,i}=\bm{m})\cdot\pi_{{\bm{\theta}}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t+1})\Bigg]+C(T),$ (17) $$

and $C(T)=\begin{cases}\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\Big[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\Big],&1\leq\beta<n\\
0,&\beta\geq n\end{cases}$ is a constant independent of $\bm{{\bm{\theta}}}$.

###### Proof.

We first consider the case with a single categorical variable $\bm{x}$. On the account of Lemma [15](https://arxiv.org/html/2510.09541v2#A2.E15) and following a similar argument as in equation [10](https://arxiv.org/html/2510.09541v2#A1.E10), for any $\beta\geq 1$, we can write

$$ $\displaystyle\log\pi_{\bm{\theta}}(\bm{x})$ $\displaystyle\leq\frac{1}{\beta}\log\mathbb{E}_{\bm{z}\sim q(\cdot\mid\bm{x})}\left[\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{x},\bm{z})}{q(\bm{z}\mid\bm{x})}\bigg)^{\beta}\right]$ $\displaystyle\quad=\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{z}_{t}\mid\bm{z}_{t+1})}{q(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]$ (18) $$

Note that the sequence $\bm{z}_{1:T}$ has a form $\{\bm{x},\dots,\bm{x},\bm{m},\dots,\bm{m}\}$. Define the transition event:

$$ $\displaystyle\mathcal{A}_{t}\coloneqq\{\bm{z}_{t}=\bm{x},\bm{z}_{t+1}=\bm{m}\}$ (19) $$

Then, by the law of total expectations, equation [18](https://arxiv.org/html/2510.09541v2#A2.E18) can be expressed as:

$$ $\displaystyle\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}\mid\bm{z}_{t+1})}{q(\bm{z}_{t}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\log\sum_{t=1}^{T-1}\mathbb{P}(\mathcal{A}_{t})\mathbb{E}_{\bm{z}\sim q(\cdot\mid\bm{x})}\left[\prod_{s=1}^{T-1}\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{s}\mid\bm{z}_{s+1})}{q(\bm{z}_{s}\mid\bm{z}_{s+1},\bm{x})}\bigg)^{\beta}\middle|\mathcal{A}_{t}\right]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\log\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\bigg[\mathbbm{1}(\bm{z}_{t+1}=\bm{m})q(\bm{z}_{t}=\bm{x}\mid\bm{z}_{t+1}=\bm{m},\bm{x})\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t}=\bm{x}\mid\bm{z}_{t+1}=\bm{m})}{q(\bm{z}_{t}=\bm{x}\mid\bm{z}_{t+1}=\bm{m},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\log\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\bigg[\mathbbm{1}(\bm{z}_{t+1}=\bm{m})\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\pi_{{\bm{\theta}}}^{\beta}(\bm{x}\mid\bm{z}_{t+1})\bigg]$ (20) $$

The penultimate line is due to the fact that conditioned on the event $\mathcal{A}_{t}$, the ratio $\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{s}\mid\bm{z}_{s+1})}{q(\bm{z}_{s}\mid\bm{z}_{s+1},\bm{x})}$ is equal to one for any $s\neq t$. The last line uses the formula for $q$. The indicator $\mathbbm{1}(\bm{z}_{t}=\bm{m})$ appears in the final expression because the terms in the bound are only non-trivial when the model must make a prediction from a corrupted state.

Now we generalize the above to a sequence of categorical variables $\bm{x}=\bm{x}_{1:n}$. Similar as [Equation˜18](https://arxiv.org/html/2510.09541v2#A2.E18), we have

$$ $\displaystyle\log\pi_{\bm{{\bm{\theta}}}}(\bm{x}_{1:n})\leq\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\prod_{i=1}^{n}\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]$ $$

Denote $\tilde{C}(T)\vcentcolon=\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\Big[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\Big]$. The upper bound in the RHS can be further derived as

$$ $\displaystyle\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\prod_{i=1}^{n}\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\prod_{i=1}^{n}\sum_{\bm{y}_{1:T}^{i}}q(\bm{y}_{1:T}^{i}\mid\bm{x})\mathbbm{1}(\bm{y}_{1:T}^{i}=\bm{z}_{1:T})\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{y}_{t,i}^{i}\mid\bm{y}_{t+1}^{i}}{q(\bm{y}_{t,i}^{i}\mid\bm{y}_{t+1}^{i},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle\leq$ $\displaystyle\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\prod_{i=1}^{n}\sum_{\bm{y}_{1:T}^{i}}q(\bm{y}_{1:T}^{i}\mid\bm{x})\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{y}_{t,i}^{i}\mid\bm{y}_{t+1}^{i}}{q(\bm{y}_{t,i}^{i}\mid\bm{y}_{t+1}^{i},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\log\Bigg(\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\Big[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\Big]\cdot\bigg(\prod_{i=1}^{n}\sum_{\bm{y}_{1:T}^{i}}q(\bm{y}_{1:T}^{i}\mid\bm{x})\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{y}_{t,i}^{i}\mid\bm{y}_{t+1}^{i}}{q(\bm{y}_{t,i}^{i}\mid\bm{y}_{t+1}^{i},\bm{x})}\bigg)^{\beta}\bigg)\Bigg)$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\log\prod_{i=1}^{n}\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]+\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\Big[q(\bm{z}_{1:T}\mid\bm{x})^{-n}\Big]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\sum_{i=1}^{n}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]+\tilde{C}(T)$ (21) $$

Here, $\bm{y}_{1:T}^{i}$ are copies of $\bm{z}_{1:T}$ enforced to agree with $\bm{z}_{1:T}$ using the indicator $\mathbbm{1}(\bm{y}_{1:T}^{i}=\bm{z}_{1:T})$. $C(T)$ is a constant independent of $\bm{{\bm{\theta}}}$, and the first term in [Equation˜21](https://arxiv.org/html/2510.09541v2#A2.E21) can be derived similar to the single variable case in [Equation˜20](https://arxiv.org/html/2510.09541v2#A2.E20):

$$ $\displaystyle\frac{1}{\beta}\sum_{i=1}^{n}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{\theta}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\sum_{i=1}^{n}\log\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\Bigg[\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\cdot\mathbbm{1}(\bm{z}_{t+1,i}=\bm{m})\cdot\pi_{{\bm{\theta}}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t+1})\Bigg]$ $$

Additionally, for $\beta\geq n$, we can obtain a tighter bound using Hölder’s inequality:

$$ $\displaystyle\frac{1}{\beta}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\prod_{i=1}^{n}\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{n\beta}\sum_{i=1}^{n}\log\mathbb{E}_{\bm{z}_{1:T}\sim q(\cdot\mid\bm{x})}\bigg[\prod_{t=1}^{T-1}\bigg(\frac{p_{\bm{{\bm{\theta}}}}(\bm{z}_{t,i}\mid\bm{z}_{t+1})}{q(\bm{z}_{t,i}\mid\bm{z}_{t+1},\bm{x})}\bigg)^{n\beta}\bigg]$ $\displaystyle=$ $\displaystyle\frac{1}{n\beta}\sum_{i=1}^{n}\log\sum_{t=1}^{T-1}\mathbb{E}_{\bm{z}_{t+1}\sim q(\cdot\mid\bm{x})}\Bigg[\frac{\alpha_{t}-\alpha_{t+1}}{1-\alpha_{t+1}}\cdot\mathbbm{1}(\bm{z}_{t+1,i}=\bm{m})\cdot\pi_{{\bm{\theta}}}^{n\beta}(\bm{x}_{i}\mid\bm{z}_{t+1})\Bigg]$ (22) $$

Denote $\tilde{\beta}\vcentcolon=n\beta$. Then the upper bound $\log\pi_{{\bm{\theta}}}(\bm{x}_{1:n})\leq\mathcal{L}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})$ holds with the constant term $C(T)$ being zero and the hyperparameter $\tilde{\beta}\geq n$. Combining these two cases, we have $C(T)=\begin{cases}\tilde{C}(T),&1\leq\beta<n\\
0,&\beta\geq n\end{cases}$.

∎

Furthermore, we can derive the continuous time version by omitting the constant term that does not affect the gradient with respect to $\bm{{\bm{\theta}}}$, and taking the limit of $T\rightarrow\infty$ similar as the derivations for $\mathcal{L}_{\text{ELBO}}$, as shown in [Equation˜7](https://arxiv.org/html/2510.09541v2#S3.E7):

###### Corollary 1 .

Taking the limit of $T\rightarrow\infty$, we have:

$$ $\begin{split}&\nabla_{\bm{{\bm{\theta}}}}\mathcal{L}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})=\nabla_{\bm{{\bm{\theta}}}}\Big(\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}})+C(T)\Big)=\nabla_{\bm{{\bm{\theta}}}}\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1:n};\bm{{\bm{\theta}}}),\penalty 10000\ \penalty 10000\ \penalty 10000\ \penalty 10000\ \text{where}\\ &\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1:n};{\bm{\theta}})=\frac{1}{\beta}\sum_{i=1}^{n}\log\mathbb{E}_{t,\bm{z}_{t}}\left[w(t)\cdot\mathbbm{1}(\bm{z}_{t,i}=\bm{m})\cdot\pi^{\beta}_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t})\right].\end{split}$ (23) $$

One caveat of the above $\tilde{\mathcal{L}}_{\text{EUBO}}$ is that the $\log$ is outside of the expectation, which in general makes Monte Carlo sample estimates biased. One could certainly further loosen the bound using the inequality $\log x\leq x-1$:

$$ $\displaystyle\mathcal{L}_{\text{EUBO}}(\bm{x})\leq\frac{1}{\beta}\sum_{i=1}^{n}\mathbb{E}_{t\sim\mathcal{U}[0,1],\bm{z}_{t}\sim q}\Bigg[w(t)\cdot\mathbbm{1}(\bm{z}_{t,i}=\bm{m})\cdot\pi_{\bm{{\bm{\theta}}}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t})\Bigg]-\frac{n}{\beta}$ (24) $$

But in practice we found this results in much worse performance, as demonstrated in [Table˜10](https://arxiv.org/html/2510.09541v2#A5.T10), potentially due to the much larger gap between EUBO and likelihood.

###### Lemma 1 (Rényi Variational Bound; Rényi ( 1961 ); Van Erven and Harremos ( 2014 ) ) .

###### Theorem 1 (Evidence Upper Bound for Masked Diffusion) .

###### Proof.

###### Corollary 1 .

## Appendix C Additional Analysis on Upper and Lower Bounds

### C.1 Proof of Proposition ˜ 1

###### Proposition 1 (Optimal Mixture Strictly Reduces Variance) .

Fix a coordinate $k$ and let

$$ $\rho_{\beta}\coloneqq{w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})}/{\mathbb{E}\left[w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]},$ $$

where $w(t,\bm{z}_{t})\coloneqq w(t)\mathbbm{1}(z_{t}=\bm{m})$. Then, the gradient of mixture objective ([8](https://arxiv.org/html/2510.09541v2#S3.E8)) is given by

$$ $\displaystyle g_{\omega,k}=((1-\omega)w(t,\bm{z}_{t})+\omega\rho_{\beta})\,\partial_{\bm{\theta}_{k}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t},\bm{c}).$ (25) $$

If $\operatorname{Var}((\rho_{\beta}-w(t,\bm{z}_{t}))\partial_{\bm{\theta}_{k}}\log\pi_{\bm{\theta}}(\bm{x}\mid\bm{z}_{t},\bm{c}))>0$, then $\operatorname{Var}[g_{\omega,k}]$ is a strictly convex quadratic in $\omega$ and thus admits a unique minimizer $\omega_{k}^{\star}$. Moreover,

$$ $\operatorname{Var}[g_{\omega_{k}^{\star},k}]\;<\;\min\!\big\{\operatorname{Var}[g_{0,k}],\ \operatorname{Var}[g_{1,k}]\big\},$ $$

###### Proof.

We first derive the formulas for the gradient of each objective. Consider a specific example $x_{i}$.
The gradient of the $\mathcal{L}_{\text{ELBO}}$ and $\tilde{\mathcal{L}}_{\text{ELBO}}$ are given by:

$$ $\displaystyle\nabla_{\bm{\theta}}\mathcal{L}_{\text{ELBO}}$ $\displaystyle=\mathbb{E}\left[w(t,\bm{z}_{t})\nabla\log\pi_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]$ (26) $\displaystyle\nabla_{\bm{\theta}}\tilde{\mathcal{L}}_{\text{EUBO}}$ $\displaystyle=\frac{\mathbb{E}\left[w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\nabla\log\pi_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]}{\mathbb{E}\left[w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]}$ (27) $$

Then the gradient of the mixture objective $\tilde{\mathcal{L}}_{\text{Mix}}$ is given by:

$$ $\displaystyle\nabla_{\bm{\theta}}\tilde{\mathcal{L}}_{\text{Mix}}=\mathbb{E}\left[\Big((1-\omega)w(t,\bm{z}_{t})+\omega\rho_{\beta}\Big)\nabla_{\bm{\theta}}\log\pi_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]$ (28) $$

We further compute the per-parameter (per-dimension) variance of the gradient of $\tilde{\mathcal{L}}_{\text{Mix}}$ and consider the optimal mixture coefficient $\omega$ to minimize the variance. For simplicity, we use the following short-hand notation:

$$ $\displaystyle s_{k}$ $\displaystyle\coloneqq\partial_{\bm{\theta}_{k}}\log\pi_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$ $$

We denote the $k$-th coordinate of the gradient $\nabla_{\bm{\theta}}\tilde{\mathcal{L}}_{\text{Mix}}$ by $g_{\omega,k}$. Then, the coordinate-wise variance of the gradient is given by

$$ $\displaystyle\operatorname{Var}\!\big[g_{\omega,k}\big]$ $\displaystyle=\mathbb{E}\!\Big[\big((1-\omega)\,w+\omega\,\rho_{\beta}\big)^{2}\,s_{k}^{2}\Big]-\Big(\mathbb{E}\!\big[\big((1-\omega)\,w+\omega\,\rho_{\beta}\big)\,s_{k}\big]\Big)^{2}$ $\displaystyle=\operatorname{Var}(ws_{k})+2\omega\operatorname{Cov}(ws_{k},(\rho_{\beta}-w)s_{k})+\omega^{2}\operatorname{Var}((\rho_{\beta}-w)s_{k})$ $$

where we used the shorthand $w\equiv w(t,z_{t})$. The above expression is quadratic in $\omega$ and we find the optimal $\omega$ by setting the derivative of variance to zero:

$$ $\displaystyle\frac{\partial}{\partial\omega}\operatorname{Var}\!\big[g_{\omega,k}\big]=2\,\operatorname{Cov}\!\big(w\,s_{k},\;(\rho_{\beta}-w)\,s_{k}\big)+2\omega\,\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)=0$ $\displaystyle\Rightarrow$ $\displaystyle\;\omega_{k}^{\star}=-\,\frac{\operatorname{Cov}\!\big(w\,s_{k},\;(\rho_{\beta}-w)\,s_{k}\big)}{\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)}.$ $$

The above yields a per-coordinate optimal $\omega^{\star}_{k}$. Equivalently, we can write $\omega^{\star}_{k}$ as follows:

$$ $\displaystyle\omega_{k}^{\star}=\frac{\operatorname{Var}\!\big(w\,s_{k}\big)-\operatorname{Cov}\!\big(w\,s_{k},\rho_{\beta}\,s_{k}\big)}{\operatorname{Var}\!\big(w\,s_{k}\big)+\operatorname{Var}\!\big(\rho_{\beta}\,s_{k}\big)-2\,\operatorname{Cov}\!\big(w\,s_{k},\rho_{\beta}\,s_{k}\big)}$ $$

Furthermore, $\omega^{\star}_{k}$ is a minimizer of coordinate-wise variance in the non-degenerative case with $\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)>0$, as the variance is strongly convex in $\omega$.

The coordinate-wise variance of gradients in $\mathcal{L}_{\text{ELBO}}$ ($\omega=0$) and $\tilde{\mathcal{L}}_{\text{ELBO}}$ ($\omega=1$), and the optimal mixture coefficient $\omega^{\star}$ are then given by

$$ $\displaystyle\mathcal{L}_{\text{ELBO}}:\quad\operatorname{Var}\!\big[g_{0,k}\big]$ $\displaystyle=\operatorname{Var}\!\big[w\,s_{k}\big],$ $\displaystyle\tilde{\mathcal{L}}_{\text{ELBO}}:\quad\operatorname{Var}\!\big[g_{1,k}\big]$ $\displaystyle=\operatorname{Var}\!\big[w\,s_{k}\big]+2\,\operatorname{Cov}\!\big(w\,s_{k},\;(\rho_{\beta}-w)\,s_{k}\big)+\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big),$ $\displaystyle\text{Optimal:}\;\;\operatorname{Var}\!\big[g_{\omega_{k}^{\star},k}\big]$ $\displaystyle=\operatorname{Var}\!\big[w\,s_{k}\big]-\frac{\Big(\operatorname{Cov}\!\big(w\,s_{k},\;(\rho_{\beta}-w)\,s_{k}\big)\Big)^{2}}{\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)}\,,$ $$

The difference between the variance of $\mathcal{L}_{\text{ELBO}}$ and $\tilde{\mathcal{L}}_{\text{ELBO}}$ with the optimal mixture coefficient can then be derived as follows:

$$ $\displaystyle\operatorname{Var}\!\big[w\,s_{k}\big]-\operatorname{Var}\!\big[g_{\omega_{k}^{\star},k}\big]$ $\displaystyle=\frac{\Big(\operatorname{Cov}\!\big(w\,s_{k},\;(\rho_{\beta}-w)\,s_{k}\big)\Big)^{2}}{\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)}\geq 0$ $\displaystyle\operatorname{Var}\!\big[\rho_{\beta}\,s_{k}\big]-\operatorname{Var}\!\big[g_{\omega_{k}^{\star},k}\big]$ $\displaystyle=\frac{\Big(\operatorname{Cov}\!\big(w\,s_{k},\;(\rho_{\beta}-w)\,s_{k}\big)+\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)\Big)^{2}}{\operatorname{Var}\!\big((\rho_{\beta}-w)\,s_{k}\big)}\geq 0$ $$

∎

###### Proposition 1 (Optimal Mixture Strictly Reduces Variance) .

###### Proof.

### C.2 Additional Comparison Between the Mixture Loss and the Lower and Upper Bounds

#### Comparing Mixture with the Lower Bound.

Consider the ratio of the coefficient of score function $\nabla_{\bm{\theta}}\log\pi_{\bm{\theta}}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$ in the gradient in the case of the mixture objective (i.e., $\nabla_{\bm{{\bm{\theta}}}}\tilde{\mathcal{L}}_{\text{Mix}}$ in [Equation˜28](https://arxiv.org/html/2510.09541v2#A3.E28)) over using only the lower bound (i.e., $\nabla_{\bm{{\bm{\theta}}}}\mathcal{L}_{\text{ELBO}}$ in [Equation˜26](https://arxiv.org/html/2510.09541v2#A3.E26)):

$$ $\displaystyle\frac{w_{\text{Mix}}}{w_{\text{ELBO}}}=\frac{(1-\omega)w(t,\bm{z}_{t})+\omega\rho_{\beta}}{w(t,\bm{z}_{t})}=(1-\omega)+\omega\frac{\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})}{\mathbb{E}\left[w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]}$ $$

Treating the expectation over all samples $\mathbb{E}\left[w(t,\bm{z}_{t})\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})\right]$ as a constant (since it is averaged), the second term in the above ratio is strictly increasing in $\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$. This realizes a confidence-aware weighting: uncertain tokens with small $\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$, i.e., those with a low recovery chance, have a smaller weight, while confident tokens with large $\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$ are upweighted, with sharpness being controlled by parameter $\beta$ and the blend by $\omega$.

#### Comparing Mixture with the Upper Bound.

We compute the ratio of coefficient of score function in the gradient of upper bound (i.e., $\nabla_{\bm{{\bm{\theta}}}}\tilde{\mathcal{L}}_{\text{EUBO}}$ in [Equation˜27](https://arxiv.org/html/2510.09541v2#A3.E27)) over the mixture gradient:

$$ $\displaystyle\frac{w_{\text{EUBO}}}{w_{\text{Mix}}}=\frac{\omega\rho_{\beta}}{(1-\omega)w(t,\bm{z}_{t})+\omega\rho_{\beta}}$ $$

Considering the above ratio, when $\pi_{\bm{\theta}}^{\beta}(\bm{x}_{i}\mid\bm{z}_{t},\bm{c})$ is very small, the coefficient of score function in $\nabla_{\bm{{\bm{\theta}}}}\tilde{\mathcal{L}}_{\text{EUBO}}$, $w_{\text{EUBO}}$, becomes very small, preventing updates to the parameters. However, the mixing approach maintains per-sample weights by preventing that from collapsing to (near) zero. In other words, for each sample, the mixture coefficient computes a convex interpolation that simultaneously floors very small EUBO weights to a minimum value and applies an uncertainty-aware capping to large EUBO weights.

Figure: Figure 7: Dynamics of the gradient norm of models trained with different log-likelihood estimation methods. SPG w/ Mixture achieves lower gradient norm and more stable optimization. We report mean and standard deviation over a rolling window of 50 steps.
Refer to caption: x6.png

#### Empirical Evidence of Reduced Gradient Variance.

As a practical indicator of gradient variance, we plot the gradient norm of each model trained with different log-likelihood estimation methods for negative advantage traces in [Figure˜7](https://arxiv.org/html/2510.09541v2#A3.F7). When using the mixture objective, the model has consistently smaller and more stable gradient norm throughout training, aligning well with our theoretical analysis.

### C.3 Toy Example for Upper and Lower Bounds.

In this section, we provide a toy example highlighting the contrasting behaviors and landscapes of the upper and lower bounds, further demonstrating the necessity to select the appropriate bound for optimization based on the optimization direction.

Consider a simple case where the sequence length is 2 and the vocabulary size is 2, i.e., $\bm{x}=[\bm{x}_{1},\bm{x}_{2}]$ and $\mathcal{V}=\{A,B\}$. Then, We can calculate $\mathcal{L}_{\text{ELBO}}$ and $\tilde{\mathcal{L}}_{\text{EUBO}}$ in closed form:

$$ $\displaystyle\mathcal{L}_{\text{ELBO}}(\bm{x}=\text{AA})$ $\displaystyle=\frac{1}{2}\Big[\log\pi_{\bm{\theta}}(\bm{x}_{1}=\text{A}\mid\text{MA})+\log\pi_{\bm{\theta}}(\bm{x}_{1}=\text{A}\mid\text{MM})$ (29) $\displaystyle+\log\pi_{\bm{\theta}}(\bm{x}_{2}=\text{A}\mid\text{AM})+\log\pi_{\bm{\theta}}(\bm{x}_{2}=\text{A}\mid\text{MM})\Big]$ (30) $\displaystyle\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}=\text{AA})$ $\displaystyle=\frac{1}{\beta}\log\Big(\frac{\pi_{\bm{\theta}}^{\beta}(\bm{x}_{1}=\text{A}\mid\text{MA})+\pi_{\bm{\theta}}^{\beta}(\bm{x}_{1}=\text{A}\mid\text{MM})}{2}\Big)$ (31) $\displaystyle+\frac{1}{\beta}\log\Big(\frac{\pi_{\bm{\theta}}^{\beta}(\bm{x}_{2}=\text{A}\mid\text{AM})+\pi_{\bm{\theta}}^{\beta}(\bm{x}_{2}=\text{A}\mid\text{MM})}{2}\Big)$ (32) $$

For simplicity, denote $a:=\pi_{\bm{\theta}}(\bm{x}_{1}=\text{A}\mid\text{MA})$ and $b:=\pi_{\bm{\theta}}(\bm{x}_{1}=\text{A}\mid\text{MM})$, and consider the of the likelihood of the first token $\bm{x}_{1}$. We have

$$ $\displaystyle\mathcal{L}_{\text{ELBO}}(\bm{x}_{1})$ $\displaystyle=\frac{1}{2}(\log a+\log b)$ (33) $\displaystyle\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1})$ $\displaystyle=\frac{1}{\beta}\log\Big(\frac{a^{\beta}+b^{\beta}}{2}\Big)$ (34) $$

Take the partial gradient with respect to $a$ and $b$ respectively,

$$ $\displaystyle\frac{\partial\mathcal{L}_{\text{ELBO}}(\bm{x}_{1})}{\partial a}$ $\displaystyle=\frac{1}{2a};\ \frac{\partial\mathcal{L}_{\text{ELBO}}(\bm{x}_{1})}{\partial b}=\frac{1}{2b}$ (35) $\displaystyle\frac{\partial\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1})}{\partial a}$ $\displaystyle=\frac{a^{\beta-1}}{a^{\beta}+b^{\beta}};\ \frac{\partial\tilde{\mathcal{L}}_{\text{EUBO}}(\bm{x}_{1})}{\partial b}=\frac{b^{\beta-1}}{a^{\beta}+b^{\beta}}$ (36) $$

Therefore, for $\tilde{\mathcal{L}}_{\text{EUBO}}$, the gradient direction is dominated by the larger one between $a$ and $b$, while for $\mathcal{L}_{\text{ELBO}}$, the gradient direction is dominated by the smaller one. Such property is illustrated in the landscapes of $-\mathcal{L}_{\text{ELBO}}$ and $-\tilde{\mathcal{L}}_{\text{EUBO}}$ for $a,b\in(0,1)$ in [Figure˜8](https://arxiv.org/html/2510.09541v2#A3.F8).

Figure: Figure 8: Landscapes of $-\mathcal{L}_{\text{ELBO}}$ and $-\tilde{\mathcal{L}}_{\text{EUBO}}$ for $0<a,b<1$. $-\tilde{\mathcal{L}}_{\text{EUBO}}$ is flatter among low value regions while sharper among high value regions, making it more suitable for log-likelihood minimization; vice versa for $-\mathcal{L}_{\text{ELBO}}$.
Refer to caption: iclr2026/images/elbo_3d.png

When $\bm{x}=\text{AA}$ has negative advantage, the corresponding $\mathcal{L}_{\text{ELBO}}$ and $\tilde{\mathcal{L}}_{\text{EUBO}}$ are minimized. For $\mathcal{L}_{\text{ELBO}}$, the model benefits more from further decreasing the smaller one between probabilities $a$ and $b$. In the extreme case, $\mathcal{L}_{\text{ELBO}}=-\infty$ when either $a$ or $b$ equals to zero, leaving the other term not sufficiently decreased. Instead, when using $\tilde{\mathcal{L}}_{\text{EUBO}}$ for negative advantage traces, the larger one between $a$ and $b$ is preferentially minimized, leading to a more balanced optimization that stably decreases the log-likelihood.

Similarly, when $\bm{x}=\text{AA}$ has positive advantage, the corresponding $\mathcal{L}_{\text{ELBO}}$ and $\tilde{\mathcal{L}}_{\text{EUBO}}$ are maximized. Using $\mathcal{L}_{\text{ELBO}}$ enables effectively increasing the smaller likelihood, while $\tilde{\mathcal{L}}_{\text{EUBO}}$ focuses on the larger one, leading to a less efficient optimization.

## Appendix D Additional Experimental Details

### D.1 Datasets and Reward Functions

We follow the setting in D1 (Zhao et al., 2025) and WD1 (Tang et al., 2025), using the same reward functions and train-test splitting, except for Sudoku. The rewards are designed to encourage both correctness and proper formatting, with varying levels of granularity tailored for each task. For completeness, we provide details as follows.

#### GSM8K.

We utilize the train split of the GSM8K dataset(^3^33[https://huggingface.co/datasets/openai/gsm8k](https://huggingface.co/datasets/openai/gsm8k)) for RL training, and evaluate model performance on the test split. We follow the Unsloth reward setup(^4^44[https://unsloth.ai/blog/r1-reasoning](https://unsloth.ai/blog/r1-reasoning)), utilizing five equally-weighted additive components:

- •
XML Structure Reward: +0.125 per correct formatting tag; small penalties for extra contents after the closing tag.
- •
Soft Format Reward: +0.5 for outputs matching 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 answer is a valid integer.
- •
Correctness Reward: +2.0 of the answer matches the ground truth.

#### MATH500.

We utilize the train split of the MATH dataset(^5^55[https://huggingface.co/datasets/ankner/math-500](https://huggingface.co/datasets/ankner/math-500)) for RL training, and evaluate model performance on the test split. We use a format reward and a correctness reward:

- •
Format Reward: We award 1.00 if <answer></answer> tags are present with \boxed inside them; 0.75 if answer tags are present without \boxed; 0.50 if answer tags are not present but \boxed is present; 0.25 if neither the answer tags nor \boxed is present.
- •
Correctness Reward: We award 2.00 if the answer in \boxed{} matches the ground truth.

#### Countdown.

We utilize the train split of the Countdown dataset(^6^66[https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4](https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4)) for RL training, restricting to instances that use only three numbers. We evaluate on the same set of 256 synthetically generated countdown questions with 3 numbers as in D1 (Zhao et al., 2025). The reward covers three cases: +1.0 if the expression reaches the target using the exact numbers; +0.1 if the numbers are correct but does not reach the target; +0.0 otherwise.

#### Sudoku.

We experiment on the 4$\times$4 Sudoku dataset(^7^77[https://github.com/Black-Phoenix/4x4-Sudoku-Dataset](https://github.com/Black-Phoenix/4x4-Sudoku-Dataset)) generated by Arel (2025). The original training split contains 1M unique Sudoku puzzles covering all 288 4$\times$4 Soduku solutions. To avoid train-test leakage and potential cheating by memorizing all the solutions, we randomly select 200 solutions and include all puzzles corresponding to these solutions into the new training set, resulting in 694,006 training puzzles. We then randomly select 2 or 3 puzzles corresponding to the left 88 solutions to construct the test set, which has 256 Soduku puzzles in total.

We observe that the zero-shot setting is too difficult for the base LLaDA-8B-Instruct model, which has test accuracy below 7% with a generation length of 256 and struggles to correctly interpret the questions, leading to very few meaningful RL rollouts. Therefore, we instead use 3-shot for all the Sudoku experiments. We ensure that the solutions presented in the 3-shot samples do not appear in test set solutions, and the puzzles do not appear in both train and test set. The detailed few-shot samples are provided in [Section˜D.3](https://arxiv.org/html/2510.09541v2#A4.SS3).

### D.2 Hyperparameter Settings and Implementation Details

We follow D1 (Zhao et al., 2025) for most hyperparameter settings. We employ Low-Rank Adaptation (LoRA) with a rank of $r=128$ and scaling factor $\alpha=64$. The training was conducted on 8 NVIDIA A100-80G or NVIDIA H100-80G GPU, with the following hyperparameters: batch size of 6 per GPU, and gradient accumulation steps of 2. We set the number of inner gradient update $\mu$ as 4 for all models. We use the AdamW optimizer (Loshchilov and Hutter, 2017), with $\beta_{1}=0.9$, $\beta_{2}=0.99$, weight decay of 0.1, learning rate of $3\times 10^{-6}$, and gradient clipping at 0.2. We utilize Flash Attention 2 (Dao, 2023)and 4-bit quantization.

For RL rollout, we use sequence length of 256 tokens, and 128 diffusion steps. We employ confidence-based semi-autoregressive generation with block size 32, and set the temperature as 0.9 (except for Sudoku where temperature is set as 0.3 following D1). We set number of completions per prompt $g$ as 6, and number of Monte Carlo estimation samples $m$ as 2 due to computational constraint. Since the rollout stage dominates the training time, the average time per gradient update step for SPG is similar to that of the other baselines.

We train 6000 steps (i.e., number of gradient updates) for GSM8K and Countdown, 4000 steps for MATH500, and 2500 steps for Sudoku. For all RL models, we run evaluation every 100 steps with generation sequence length 128, 256, 512, and report the result of the checkpoint with the highest average accuracy over the three generation lengths, except for the ablations on inference strategies in [Table˜13](https://arxiv.org/html/2510.09541v2#A5.T13), where we only evaluate on generation length 256.

### D.3 Few-Shot Samples for Sudoku

#### Sample 1:

#### Sample 2:

#### Sample 3:

## Appendix E Additional Results

### E.1 Additional Evaluations to the Main Results

#### Complete evaluation results.

We provide the complete evaluation results, along with those reported in D1 (Zhao et al., 2025) and WD1 (Tang et al., 2025), in [Table˜4](https://arxiv.org/html/2510.09541v2#A5.T4). Our reproduced numbers closely match the reported results. d1-LLaDA (Zhao et al., 2025) denotes the model that conducts first SFT and then RL (using D1). All other models are trained solely with RL.
In D1 and d1-LLaDA, the best result for each generation length is reported separately, whereas we select a single checkpoint with the highest average accuracy across all three generation lengths, leading to slightly worse results than the reported numbers. The reported results in WD1 are based on evaluations on fewer checkpoints, so they are generally a bit lower than our reproduced values.

Table: Table 4: Complete model performance on four reasoning benchmarks compared with baselines. We provide both the reported and the reproduced results for D1 and WD1. The best results are bolded and the second best are underlined. SPG consistently outperforms all other models.

#### Dynamics of Completion Length.

We provide the dynamics of the effective sequence length of SPG during RL training in [Figure˜9](https://arxiv.org/html/2510.09541v2#A5.F9). We also report the effective length of the best checkpoint in [Table˜5](https://arxiv.org/html/2510.09541v2#A5.T5). SPG leads to effective usage of the total given length and good adaptation to task difficulties.

Figure: Figure 9: Dynamics of the effective generation length of SPG during RL training, compared with D1, WD1, and UniGRPO. SPG leads to concise solutions with better token efficiency. We report mean and standard deviation over a rolling window of 50 steps.
Refer to caption: x7.png

Table: Table 5: Effective sequence length of each model at the best checkpoint corresponding to [Table˜1](https://arxiv.org/html/2510.09541v2#S4.T1) on four reasoning benchmarks.

### E.2 Additional Ablation Results

In this section, we provide the complete results for each generation length and task in supplement to [Section˜4.2](https://arxiv.org/html/2510.09541v2#S4.SS2). We also include additional ablation studies on the looser upper bound, different log-likelihood estimation methods for positive advantage traces, and Pass@K performance.

#### Ablations on Algorithm Components.

We provide the complete results for ablations on log-likelihood estimation methods in [Table˜6](https://arxiv.org/html/2510.09541v2#A5.T6) and for ablations on masking strategies in [Table˜7](https://arxiv.org/html/2510.09541v2#A5.T7).

Table: Table 6: Ablations on log-likelihood estimation methods for negative advantage traces. The best results are bolded and the second best are underlined. SPG w/ Mixture consistently outperforms other likelihood estimation methods.

Table: Table 7: Ablations on the masking strategies in Monte Carlo estimation. Our block-wise masking strategy leads to consistent improvement to random masking on both benchmarks.

#### Ablations on Key Hyperparameters β \beta and ω \omega .

We provide the complete results for ablations on $\beta$ in [Table˜8](https://arxiv.org/html/2510.09541v2#A5.T8) and for ablations on $\omega$ in [Table˜9](https://arxiv.org/html/2510.09541v2#A5.T9).

Table: Table 8: Ablations on the value of $\beta$ in the upper bound.

Table: Table 9: Ablations on the mixture coefficient $\omega$ on MATH500 and Countdown.

#### Ablations on Inference Strategies.

We provide complete results for ablations on different inference strategies in [Table˜13](https://arxiv.org/html/2510.09541v2#A5.T13). Note that the reported numbers of each method for “Semi-AR, Block=32, Confidence” is in general slightly higher than the results in [Table˜1](https://arxiv.org/html/2510.09541v2#S4.T1) under the same inference setting. This is because in [Table˜13](https://arxiv.org/html/2510.09541v2#A5.T13), we select best checkpoint specifically for generation length 256 to maintain consistency with other inference settings, while in [Table˜1](https://arxiv.org/html/2510.09541v2#S4.T1), we choose the checkpoint with the highest average accuracy across generation lengths 128, 256, and 512.

#### Ablations on the Looser Upper Bound.

As mentioned in [Section˜3.2](https://arxiv.org/html/2510.09541v2#S3.SS2) and [Appendix˜B](https://arxiv.org/html/2510.09541v2#A2), a looser but unbiased bound can be derived using inequalities like $\log(x)\leq x-1$, i.e., $\tilde{\mathcal{L}}_{\text{Loose}}$ ([Equation˜24](https://arxiv.org/html/2510.09541v2#A2.E24)). However, as shown in [Table˜10](https://arxiv.org/html/2510.09541v2#A5.T10), this looser bound performs worse empirically than the tighter upper bound $\tilde{\mathcal{L}}_{\text{EUBO}}$ we used, possibly due to a larger discrepancy from the true log-likelihood.

Table: Table 10: Ablations on the looser upper bound. The loose bound performs worse than the tighter upper bound we used, indicating inferior performance due to a larger discrepancy from the true log-likelihood.

Table: Table 11: Ablations on log-likelihood estimation for positive advantage traces. Using the upper bound for log-likelihood estimation of positive advantage traces perform worse than using the lower bound.

#### Ablations on Log-Likelihood Estimations for Positive Advantage Traces.

Instead of always using $\mathcal{L}_{\text{ELBO}}$ for positive advantage traces, we experiment on MATH500 and Countdown benchmarks using both $\tilde{\mathcal{L}}_{\text{EUBO}}$ and $\tilde{\mathcal{L}}_{\text{Mix}}$ for positive advantage traces.
Correspondingly, we use $\omega=0.5$ and the best performed $\beta$ as previously discussed for negative advantage traces.
For the positive advantage traces, we always use the tightest $\beta=1.0$ for both $\tilde{\mathcal{L}}_{\text{EUBO}}$ and $\tilde{\mathcal{L}}_{\text{Mix}}$. The results are shown in [Table˜11](https://arxiv.org/html/2510.09541v2#A5.T11), indicating that using the upper bound for likelihood estimation of positive advantage traces performs worse than using $\mathcal{L}_{\text{ELBO}}$. This aligns well with our theoretical insights that the lower bound is a better objective for log-likelihood maximization.

#### Ablations on Pass@K Performance.

In all previous experiments, we apply greedy sampling by setting temperature as 0.0 following D1 and LLaDA. However, beyond accuracy, it is essential for models to generate a diverse set of outputs that can cover the correct solution and allow for explorations. In this section, we investigate the models’ ability to generate diverse outputs using a higher temperature, and evaluate their Pass@K performance on MATH500 and Countdown, as shown in [Table˜12](https://arxiv.org/html/2510.09541v2#A5.T12). Specifically, we set temperature to 0.9 and generation length to 256, conduct evaluations every 100 steps, and report results from the checkpoint with the highest accuracy. For comparison, we also include results from greedy sampling, denoted as Pass@1Greedy. As expected, increasing the temperature leads to a decrease in Pass@1 performance across all models, aligning with observations from previous work. For K>1, the Pass@K scores improve for all models as K increases from 1 to 4. SPG achieves the best performance across all settings, with SPG w/ Mixture reaching 55.6% Pass@4 accuracy on MATH500 and 76.6% on Countdown, demonstrating the ability of SPG to generate diverse outputs that can recover the correct solution.

Table: Table 12: Pass@K performance of each model on MATH500 and Countdown. We set temperature as 0.9 and report results of the best checkpoint of each case at a generation length of 256. For comparison, we also include the greedy sampling performance, i.e., Pass@1Greedy. The best results are bolded and the second best are underlined.

Table: Table 13: Ablations on the inference strategy. SPG leads to consistently superior performance to baselines with different inference strategies. The best results are bolded and the second best are underlined for each setting. We report results for generation length 256.

### E.3 Qualitative Examples on GSM8K

We present qualitative examples for all models on GSM8K, using a generation length of 256 and 128 diffusion steps. Incorrect reasoning steps or answers are highlighted in red and the correct answers are highlighted in green. SPG leads to more concise and accurate reasoning compared to other models. Note that the models may sometimes achieve the correct answer through flawed reasoning, while such artifacts are expected to diminish with a stronger base model.

#### Model: LLaDA-8B-Instruct

#### Model: D1

#### Model: WD1

#### Model: UniGRPO

#### Model: SPG w/ EUBO

#### Model: SPG w/ Mixture

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. In all main experiments, the paper set the number of Monte Carlo samples to $m=2$. Ignoring computational constraints, how does increasing $m$ affect optimization variance and end-task performance?

2. Figure 9 depicts the training dynamics of the effective generation length. What do these trajectories imply about how SPG allocates its reasoning budget over training? For tasks such as Countdown and Sudoku, the effective length shows large-amplitude fluctuations; what factors drive these oscillations?
soundness: 3
strengths: 1. The paper articulates the RL bottleneck for dLLMs and motivates a natural sandwiched approach.

2. Experimental results show consistent gains across benchmarks, sequence lengths, and inference strategies; ablations on some choices further support robustness.

3. The paper is easy to follow and uses a compact structure that makes the method clear.
summary: The paper introduces SPG for dLLMs. Because dLLMs have intractable log-likelihoods, standard policy gradients cannot be applied directly. SPG sandwiches the likelihood by maximizing an ELBO for positive advantage samples and minimizing an EUBO for negative advantage samples, paired with blockwise masking for Monte Carlo estimation. SPG achieves SOTA accuracy on GSM8K, MATH500, Countdown, and Sudoku among RL for dLLMs.
weaknesses: 1. LoRA-only evidence without full fine-tuning controls: all experiments rely on LoRA, so findings may hinge on LoRA’s inductive biases; it remains unclear whether the gains and stability persist under full fine-tuning.

2. The authors mention that masking 15% of the prompt can improve performance. Although they claim this choice follows d1, they do not seem to explain here the motivation or why it leads to better results.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: What’s the wall-clock and token-throughput overhead of block-wise MC vs random masking and vs GRPO/D1/W1, at equal budgets?
Could the authors add (even small-scale) comparisons to (a) DPO-style for diffusion, (b) score-matching/pathwise surrogates with a learned decoder, (c) GRPO + stronger variance reduction?
Can the authors add more datasets? (non-reasoning, multi-turn, tool-use benchmarks)
soundness: 3
strengths: 1. Very clear problem: ELBO-only training biases gradients when rewards can be negative; upper bounds make the algorithm penalize low-reward traces without relying on true $\log \pi$.
2. Strong experimental results: 4 reasoning tasks achieve very high improvement/performance gains.
3. The theory is reasonable. EUBO derived from a Renyi variational inequality gives a tractable surrogate. Mixing lower/upper bounds reduces variance both intuitively and theoretically (Prop. 1) and avoids gradient clipping/vanishing.
summary: This paper addresses policy gradients for Diffusion language models where $\log \pi_{\theta} (x |c)$ is intractable by optimizing a lower bound for positive-advantage traces and an upper bound for negative advantage traces with a practical block-wise masking estimator. Because dLLMs do not expose a tractable log-likelihood, prior works optimize ELBO or one-step proxies, which bias gradients (ELBO is only a lower bound). Sandwiched Policy Gradient proposes to "sandwich" the objective: maximize a tractable lower bound $L_{ELBO}$ on positive-advantage samples while minimizing a tractable upper bound $U_{EUBO}$ on negative advantage samples (For instance, GRPO).  The problem is well-defined with clear mathematical analysis and extensive experiments on RL baselines.
weaknesses: 1. Baselines. Comparisons focus on GRPO-like/diffusion-RL. Missing: preference optimization tailored to diffusion (DPO-style for dLLMs), score-function or pathwise estimators based on learned surrogates or likelihood-ratio via learned decoder controls.
2. The datasets are limited to reasoning tasks. Though the performance gains are significant, the datasets are limited to math/logic. No tool-use or multi-turn agents settings. No natural-language preferences.
3. Training cost vs baselines and wall-clock time are not clearly examined. Blockwise MC might add computational overhead. Can the authors report that the computational overhead for each method?

---

confidence: 5
contribution: 2
initial_rating: 4
presentation: 3
questions: - Regarding Weakness 1: Would the authors consider adopting a more standard evaluation setup (e.g., validation-based checkpoint selection or reporting final-step results) to reduce overfitting concerns?

- Regarding Weakness 2: Could the authors clarify why SPG does not adopt clipping or KL regularization? Is this omission due to the intractability of computing importance ratios caused by the constant $C(T)$ term in EUBO, or was it omitted for simplicity?

- Regarding the tightness of the EUBO: As discussed in [3] and [4], the ELBO of dLLMs equals the AO-ARM loss, and becomes tight when the joint distribution $p(x|\sigma)$ is consistent across different orders of $\sigma$ (i.e., the inequality in Eq. (2) of [3] holds with equality). However, as shown in Appendix C.3 of this paper, the EUBO does not appear to be tight even under this condition. Could the authors explain why this phenomenon occurs, and whether it implies a fundamental looseness in the Rényi-based upper bound?

[3] AUTOREGRESSIVE DIFFUSION MODELS, ICLR 2022

[4] YOUR ABSORBING DISCRETE DIFFUSION SECRETLY  MODELS THE CONDITIONAL DISTRIBUTIONS OF  CLEAN DATA, ICLR 2025
soundness: 2
strengths: - Originality:
A novel Sandwiched Policy Gradient (SPG) framework is proposed to leverage both lower and upper bounds of log-likelihood for diffusion LLM, which is a clear conceptual advance over prior ELBO-only RL approaches.

- Quality:
Theoretical development is coherent and well-motivated, with a solid integration of Rényi-based upper bounds and a mixture formulation that balances bias and variance.

- Clarity:
The paper is well-written and logically structured. Figures and equations clearly illustrate the SPG process.
summary: This paper proposes a new reinforcement learning (RL) algorithm for diffusion large language models (dLLMs), which have intractable likelihoods that make standard policy gradient methods infeasible. Prior approaches approximate the log-likelihood using the ELBO, but the lower-bound approximation introduces gradient bias and limits learning from negative rewards. To overcome this, the authors introduce the Sandwiched Policy Gradient (SPG) method, which optimizes a “sandwiched” objective combining both a lower bound (ELBO) for positive-reward samples and an upper bound (EUBO) for negative-reward samples, thereby reducing bias in the policy gradient. SPG further employs a block-wise masking strategy to stabilize Monte Carlo estimation and a mixture formulation that adaptively blends upper and lower bounds to reduce gradient variance. Experiments on four reasoning benchmarks—GSM8K, MATH500, Countdown, and Sudoku—show that SPG achieves consistent improvements.
weaknesses: 1. Non-standard evaluation protocol.
The paper selects checkpoints every 100 steps based on the highest test accuracy, which risks test set overfitting. While this follows d1 for consistency, it is methodologically questionable. The model should instead be selected via a validation set or by reporting the final checkpoint performance, as adopted in [1]. Revising the evaluation protocol would improve the experimental rigor.

2. Absence of standard RL stabilization techniques.
SPG appears to use a naive policy gradient update without employing the clipping or KL-regularization mechanisms used in PPO or GRPO. Without importance sampling ratio corrections, it is unclear whether SPG remains stable for off-policy updates $\mu > 1$. This raises concerns about potential instability or divergence during training.

 3. Lack of comparison with key related work.
The paper omits comparison with [2], which is the first to successfully apply trajectory-level RL to diffusion LLMs and reports state-of-the-art results on similar reasoning benchmarks. Including this baseline is essential for contextualizing SPG’s improvements and validating its claimed advantage.


[1] DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation, arXiv:2506.20639

[2] Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models, NeurIPS 2025.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: ### Clarity on the reinforcement learning aspect

The reinforcement learning part is presented in a quite limited way, raising a number of questions.
* The underlying Markov decision process is not even defined, it could
* The fundamental argument of the paper is that the ELBO doesn’t hold for a negative reward. This is true. However, in RL, the optimal policy is invariant to a reward shifting. So in principle one can assume the reward to be positive without loss of generality. This may not be the case empirically speaking, for various reasons (one being the variance of the policy gradient, for example), but this would call for at least some discussion, ideally some baseline/ablation. The more close ablation is SPG w/ neg, which is something fundamentally different (it’s closer to SFT on positive examples, ignoring negative ones, a smooth version of best-of-n). Such a baseline could be Reinforce, but one could also imagine a GRPO variation with a normalization leading to positive advantages (using the fact that a state-dependent baseline does not bias the gradient). Please discuss this aspect as thoroughly as possible.
* Objective in Eq (4) is biased, because GRPO one is biased (because the baseline does depend on the sampled action). The correct way to do it would be to have a leave-one-out empirical expectation of the return, known as RLOO [A, B]. It’s not a big deal, probably doesn’t change much empirically speaking, but given that the whole point is about better sandwiching the policy objective, it is worth starting from an unbiased one.
* Obj in Eq (4), and then the proxy in Eq (5), do not consider importance sampling at all. However, it seems that most of the baselines are. Is the proposed approach purely on-policy, like Reinforce, or does it insures some off-policyness, like GRPO, but not taken into account in the loss, implying some additional bias in what is sandwiched ? (As a side remark, RLOO without importance sampling is not biased in the off-policy case, as a corollary of [C], but the overall discussion point remains)
* Many approaches in LLMs, but also in dLLMs, consider regularized RL (typically KL-regularization towards the initial model). It is not discussed here. Is it that no regularization is considered (which is perfectly fine, but worth discussing), or that it is skipped to lighten notations (but in this case it should not, could have implications)

### Too favorable aspect of experiments

This may be a wrong feeling, happy to be corrected, but it relies on the following points.

* As an initial side note, it would make sense to consider some baselines/ablations suggested above (related to the optimal policy invariance to shifting rewards).
* The main point is that the considered baselines (D1, WD1, UniGRPO) are not described enough (even considering the related works section in the appendix), and so we do not know what are the key differences with the proposed approach. The current ablation shows well that the different components are important to proposed approach, but what is missing is that the considered baselines are probably a form of ablation somehow. For example, SPG w/ ELBO seems pretty close to some of the baselines (like uniGRPO, but without IS/clipping?), how does it compare? Another example is that, if we look at the results of Table 1 for the baselines, and compare it to the results of SPG with the random masking strategy, the results are much closer (eg uniGRPO vs SPG w/ EUBO and random masking). So one may wonder if just applying the block-wise masking to uniGRPO would not lead to very good results too (weakening the EUBO contribution). Maybe it is not the case thanks to the mixing, but also maybe that the results are more nuanced then how they are showed (or not, but this would then call for a better discussion about baselines, their key differences, notably wrt SPG). 
* Another point that presents too favorably the proposed approach is that the SPG specific hyper parameters ($\beta$ and $\omega$) are basically tuned on the test set, which biases too favorably SPG.


[A] Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLM, Ahmadian et al., 2024   
[B] Buy 4 REINFORCE Samples, Get a Baseline for Free! Kool et al, 2019.  
[C] Contrastive Policy Gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion, Flet-Berliac et al, 2024.
soundness: 2
strengths: * The paper is clearly written and structured, easy to follow globally and well argued 
* The point of the ELBO being no valid for negative rewards is a very good one, and the proposed approach makes a lot of sense, and provides good experimental results
* The empirical study is quite thorough with interesting ablations and relevant baselines
summary: This paper introduces a policy gradient approach for RL fine-tuning diffusion LLMs (dLLMs). The authors not that the policy gradient involves the gradient of the log-likelihood of the policy, this log-likelihood being intractable in the case of dLLLMs. Existing approaches rely on an evidence lower bound (ELBO), which holds only when the return is positive. The authors build an evidence upper bound (EUBO) as a corollary of the Rényi variational bound, and use it for negative returns (ELBO being used for positive rewards). The authors also provide additional improvements. One is the mixing of ELBO and EUBO for negative returns, with a theoretical argument about decreasing variance (not discussing the bias though), the other one being a block-wise masking strategy for MC estimation. The authors then present a quite thorough experimental study, comparing to a number of baselines and ablating the different components of the proposed approach.
weaknesses: * From a reinforcement learning perspective, a lot of aspects regarding the proposed approach lack clarity. See questions for more details.
* The experiments may be presented in a too favorable aspect for the proposed approach. This may also be related to not enough discussed baselines (especially what are the key differences). See questions for more details.


Decision: Accept (Poster)