Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models

Ziwei Luo, Ziqi Jin, Lei Wang, Lidong Bing, Thomas B. Schön
  Uppsala University  
  MiroMind AI  
  Nanyang Technological University  
arXiv Code

Self-Rewarding SMC for Masked Diffusion Language Models


TL;DR. Self-Rewarding SMC is an inference-time scaling method that leverages trajectory-level confidence from diffusion models as importance weights to steer generation toward globally confident, high-quality samples.

Summary

This work presents self-rewarding sequential Monte Carlo (SMC), an inference-time scaling algorithm enabling effective sampling of masked diffusion language models (MDLMs). Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy, where only tokens with the highest prediction confidence are preserved at each step. This restricts the generation to a noise-sensitive, greedy decoding paradigm, resulting in an inevitable collapse in the diversity of possible paths. We address this problem by launching multiple interacting diffusion processes in parallel, referred to as \emph{particles}, for trajectory exploration. Importantly, we introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights. During sampling, particles are iteratively weighted and resampled to systematically steer generation towards globally confident, high-quality samples. Our self-rewarding SMC is verified on various masked diffusion language models and benchmarks, achieving significant improvement without extra training or reward guidance, while effectively converting parallel inference capacity into improved sampling quality.


Reformulate the Sampling of Masked Diffusion Models

Given pretrained model $p_\theta(\mathbf{x}_t)$ and a mask set $\mathcal{M}_t \triangleq \{\ell: \mathbf{x}_t(\ell)=\mathtt{[MASK]}\}$ at time $t$. For each token $\mathbf{x}_{t-1}(j)$, we directly define its confidence as the model probability on $j$, as


$\mathbf{c}_t (j) := p_\theta\big(\hat{\mathbf{x}}_0(j) \mid \mathbf{x}_t \big), \quad j \in \mathcal{M}_t.$


Low-confidence remasking: 1) Top-K $\mathcal{S}_t^{\text{top-k}} = \{\, j\in \mathcal{M}_t:\; \text{Top-}\rho_t \{ \mathbf{c}_t{(j)} \} \,\}$ and 2) Threshold $\mathcal{S}_t^{\text{thr}} = \{\, j\in \mathcal{M}_t:\; \mathbf{c}_t{(j)} \ge \rho_t \,\}.$

In summary, the reverse transition distribution of each token $\mathbf{x}_t (j)$ can be formulated by

$p_\theta(\mathbf{x}_{t-1} (j) \mid \mathbf{x}_t) = \begin{cases} p_\theta(\mathbf{x}_{t-1} (j) \mid \mathbf{x}_t, \hat{\mathbf{x}}_0), & j\in \mathcal{S}_t,\\ \mathrm{Cat}(\mathbf{x}_{t-1} (j); \mathbf{m}), & j\in \mathcal{M}_t\setminus \mathcal{S}_t,\\ \mathrm{Cat}(\mathbf{x}_{t-1} (j); \mathbf{x}_t), & j\notin \mathcal{M}_t, \end{cases}$


Confidence-based Sequential Monte Carlo

Proposition 3.1. Given a pretrained diffusion model $p_\theta$, let $\{\tilde\pi_t(\mathbf{x}_{t:T})\}_{t=0}^T$ denote the unnormalized path measures defined by a Feynman--Kac recursion. If the sequential proposal in SMC is chosen to be the diffusion transition kernel, i.e., $q_{t-1}(\mathbf{x}_{t-1} \mid \mathbf{x}_t) = K_t(\mathbf{x}_t,\mathbf{x}_{t-1})$, then the incremental importance weights at step $t-1$ is given by


$\tilde{w}_{t-1}(\mathbf{x}_{t-1:T}) = \prod_{j\in \mathcal{S}_t} \mathbf{c}_t(j),$


where $\mathbf{c}_t (j) := p_\theta\big(\hat{\mathbf{x}}_0(j) \mid \mathbf{x}_t \big)$ is the token confidence and $\mathcal{S}_t$ denotes the selected mask subset to be updated at step $t$.

SMC maintains multiple diffusion processes, called particles, to explore the sampling trajectories in parallel. At each iteration, we take three steps: resample, propagate, and re-weight, to perform as an interactive optimization process. Importantly, traditional diffusion sampling only considers token-level confidence, while our algorithm uses the trajectory-level confidence as importance weights to select globally confident outputs.

Experiments on Diffusion Large language Models (dLLMs)


Self-rewarding SMC improves dLLMs in math and coding


Overall performance trends as #particle increases


Effect of Gumbel noise temperature on model performance

Key Innovations

  • Self-rewarding SMC at inference time. We introduce a general sequential Monte Carlo framework for masked diffusion language models that improves sampling quality without additional training or external reward models.
  • Trajectory-level confidence as an intrinsic reward. From a probabilistic perspective, we reinterpret MDLM sampling and remasking, showing that trajectory confidence naturally provides a self-rewarding signal for particle weighting.
  • Consistent gains across models and tasks. Experiments on multiple pretrained MDLMs and diffusion LLMs demonstrate improved sample quality on a range of benchmarks.

BibTeX

If our code helps your research or work, please consider citing our paper. The following are BibTeX references:

@article{luo2026self,
  title={Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models},
  author={Luo, Ziwei and Jin, Ziqi and Lei, Wang and Bing, Lidong and Sch{\"o}n, Thomas B},
  journal={arXiv preprint arXiv:2602.01849},
  year={2026}
}