Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models

Ziwei Luo, Ziqi Jin, Lei Wang, Lidong Bing, Thomas B. Schön

TL;DR

This work presents self-rewarding sequential Monte Carlo (SMC), an inference-time scaling algorithm enabling effective sampling of masked diffusion language models (MDLMs), and introduces the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.

Abstract

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 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. Our code is available at https://github.com/Algolzw/self-rewarding-smc.

Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models

TL;DR

This work presents self-rewarding sequential Monte Carlo (SMC), an inference-time scaling algorithm enabling effective sampling of masked diffusion language models (MDLMs), and introduces the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.

Abstract

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 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. Our code is available at https://github.com/Algolzw/self-rewarding-smc.
Paper Structure (32 sections, 1 theorem, 40 equations, 7 figures, 7 tables, 1 algorithm)

This paper contains 32 sections, 1 theorem, 40 equations, 7 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Notation
  4. Masked Diffusion Models
  5. Importance Sampling and Sequential Monte Carlo
  6. Self-Rewarding Sequential Monte Carlo
  7. Reformulate the Sampling of MDMs
  8. Confidence-based Sequential Monte Carlo
  9. Practical Sampling with Self-Rewarding SMC
  10. Adaptive resampling.
  11. Gumbel-Max sampling.
  12. Experiment
  13. Experimental Setup
  14. Pretrained models
  15. Implementation
  16. ...and 17 more sections

Key Result

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 the recursion in Eq. (eq:fk_recursion). If the sequential proposal in SMC is chosen to be the diffusion transition kernel, i.e., $q_{t-1}({\mathbf{x}}_{ where ${\mathbf{c}}_t (j) \coloneqq p_\theta(\hat{{\mathbf{x}}}_0(j) \mid {\mathbf{x}}_t )$ is the

Figures (7)

  • Figure 1: Illustrative example of text generation using (a) masked diffusion models and (b) our self-rewarding SMC framework. Here, 'M' represents $\mathtt{[MASK]}$ tokens and 'Resa.' denotes resampling. 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, calculated using Eq. (\ref{['eq:conf_iiweights']}), to select globally confident outputs.
  • Figure 2: Generative perplexity ($\downarrow$) comparison of our self-rewarding SMC and the corresponding baselines.
  • Figure 3: Comparison results of LLaDA 1.5 and Dream-7B using different numbers of particles on four tasks. Each marker denotes the empirical result while the dashed curves indicate first-order polynomial fits used solely to illustrating overall trends as $N$ increases.
  • Figure 4: A qualitative comparison of reasoning trajectories. Greedy decoding focus on step-wise confidence sahoo2024simple leads to a hallucinated identity ($3$ Blinkets $= 7$ Blinkets) that persists through the chain. SR-SMC utilizes trajectory-level confidence to explore multiple trajectories in parallel and successfully recovers the correct multi-step conversion.
  • Figure 5: Effect of sampling temperature $\tau$ on model performance across different benchmarks. We report the accuracy of LLaDA-1.5 and Dream-7B on MBPP and MATH datasets as the temperature varies uniformly from 0.0 to 1.0. The blue circles represent the baseline with standard parallel decoding, while the red stars denote the results using our SR-SMC with $N=4$ particles. SR-SMC consistently demonstrates better robustness across the entire temperature range. Notably, while the baseline performance of Dream-7B collapses at low temperatures (startig from 0.1) due to repetition, SR-SMC maintains stable and high accuracy by effectively exploring the generative space through particle re-weighting and resampling.
  • ...and 2 more figures

Theorems & Definitions (2)

  • Proposition 3.1
  • proof