PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models

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

arXiv:2602.01849v1cs.LG

TL;DR

Masked diffusion language models can become trapped by greedy, step-wise confidence sampling, while external reward guidance is task-specific and requires hand-crafted tuning. This work introduces self-rewarding sequential Monte Carlo with interacting particles and trajectory-level confidence, improving pretrained MDLM quality and diversity without extra training or external rewards.

  • Problem

    Existing MDLMs use greedy confidence-based sampling that promotes locally optimal tokens and myopic trajectory exploration, while external reward-guided methods restrict general applicability through task-specific signals and hand-crafted tuning.

  • Method

    Self-rewarding SMC runs multiple interacting diffusion particles in parallel and uses trajectory-level confidence as an implicit reward for importance weighting and resampling.

  • Results

    Self-rewarding SMC significantly improves pretrained MDLM sample quality and diversity across multiple benchmarks and model families; with N = 4, average scores rise from 49.3 to 52.1 for LLaDA-1.5 and from 51.9 to 56.4 for Dream-7B.

  • Takeaways & Limitations

    The method provides an inference-time scaling dimension that converts parallel particle computation into more effective masked-diffusion generation without additional training or external reward guidance.

  • Takeaways & Limitations

    The method increases inference-time computation and relies solely on model likelihood, so it does not explicitly optimize downstream objectives such as reasoning correctness or human preference.

Abstract

from arXiv · show

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.

1. Introduction

Masked diffusion language models commonly use greedy confidence-based sampling, which limits trajectory exploration. The paper proposes self-rewarding SMC, using multiple particles and trajectory-level confidence to improve sampling without extra training or external reward guidance.

  • Existing MDLMs preserve only tokens with the highest prediction probability at each step, causing myopic trajectory exploration and suboptimal generation.
  • Inference-time scaling methods can improve diversity and quality without modifying pretrained models but often depend on task-specific external rewards and hand-crafted tuning.
  • Self-rewarding SMC runs interacting diffusion particles in parallel and resamples them using accumulated trajectory-level confidence as an implicit reward.
  • The method improves sampling without extra training or reward signals and is evaluated across multiple pretrained masked diffusion models and benchmarks.
  • The method unifies MDLM sampling and remasking from a probabilistic perspective, showing that trajectory-level confidence naturally supplies a self-rewarding signal for SMC.

2. Background

Masked diffusion models reverse a forward masking process to generate text, while importance sampling and SMC provide trajectory-based inference procedures. SMC reduces practical variance through sequential weighting and resampling, producing weighted particles that approximate a target trajectory distribution.

  • Masked Diffusion Models: Masked diffusion models reverse a forward masking process over sequences, with masked tokens represented by a special [MASK] category.
  • Masked Diffusion Models: During inference, the reverse process samples predicted clean tokens and applies the learned posterior to masked positions while leaving unmasked tokens unchanged.
  • Importance Sampling: Importance sampling rewrites expectations under a target trajectory distribution using a proposal distribution, but can suffer from high variance.
  • Sequential Monte Carlo: SMC introduces intermediate path measures and combines sequential weighting with resampling to reduce variance across the trajectory.
  • Sequential Monte Carlo: At each reverse-time iteration, SMC resamples ancestors, propagates particles, and re-weights them using incremental importance weights.
  • Sequential Monte Carlo: The resulting weighted particle collection provides an asymptotically consistent approximation of the trajectory target distribution.
  • Sequential Monte Carlo: The paper applies this framework to diffusion sampling by maintaining multiple particles and using trajectory-level confidence for importance weighting.

3. Self-Rewarding Sequential Monte Carlo

Self-rewarding SMC extends masked diffusion sampling with interacting particles whose trajectory-level confidence weights guide resampling toward globally confident sequences. It combines local confidence-based diffusion transitions with sequential weighting, adaptive resampling, and maximum-weight final selection.

  • Sampling reformulation: Masked diffusion sampling preserves high-probability tokens, remasks low-confidence tokens, and deterministically retains already unmasked tokens during reverse transitions.The update policy selects a subset of masked tokens for iterative unmasking while newly accepted tokens are sampled from model predictions.
  • Sampling reformulation: Step-wise confidence can bias generation toward locally optimal tokens, producing noise-sensitive and myopic trajectory exploration.This motivates replacing purely local selection with a trajectory-level confidence signal.
  • Confidence-based SMC: SR-SMC runs N interacting diffusion particles in parallel and defines a potential from the joint probability of accepted tokens.The potential serves as an implicit self-rewarding signal for SMC updates, tilting sampling toward globally confident sequence generation.
  • Confidence-based SMC: Choosing the diffusion transition kernel as the SMC proposal yields incremental importance weights that accumulate token-confidence scores across sampling steps.This proposal corresponds to a bootstrap SMC scheme, so trajectory-level reweighting follows from the diffusion formulation rather than being a heuristic.
  • Practical sampling: At each iteration, SR-SMC resamples, propagates, and reweights particles, then selects the final sequence with maximum weight.The procedure initializes fully masked particles with uniform weights and uses adaptive resampling when weight degeneracy is significant; it also supports Gumbel-Max token sampling.

4. Experiment

Experiments evaluate self-rewarding SMC across pretrained masked diffusion models and diffusion large language models, showing improved sample quality and benchmark performance under inference-time scaling.

  • Experimental Setup: Self-rewarding SMC is evaluated across multiple pretrained diffusion language models and benchmarks.The experiments include MDLM, BD3-LMs, LLaDA-1.5, and Dream-7B.
  • Sample Quality Evaluation: Generative perplexity results on sequences of lengths L = 1024, 2048 show improved sample quality for diffusion and block diffusion baselines.The evaluation uses GPT2-Large generative perplexity on 300 generated samples.
  • Diffusion Large Language Models: The LLaDA-1.5 and Dream-7B experiments evaluate GSM8K, MATH, HumanEval, and MBPP using generation lengths L ∈{256, 512} with block size 32.These tasks cover mathematical reasoning and code generation.
  • Diffusion Large Language Models: SR-SMC consistently improves results across benchmarks, model architectures, and generation lengths.Average gains are 2.8+ for LLaDA-1.5 and 4.5+ for Dream-7B.
  • Particle Scaling: Increasing the number of particles produces consistent performance improvements across both evaluated models and all four benchmarks.The strongest gains typically occur at N = 3 or N = 4, while N = 1 corresponds to standard parallel decoding.

5. Discussion and Analysis

Scaling particle count improves performance across benchmarks, while trajectory-level resampling improves exploration, robustness, and recovery from locally poor decoding choices.

  • Scaling the Number of Particles: LLaDA-1.5 improves from 49.3 to 52.1 and Dream-7B from 51.9 to 56.4 when scaling to N = 4 particles.The strongest gains typically occur at N = 3 or N = 4.
  • Analysis of Particle Overtake: 24% to 31% of analyzed blocks contain particle overtake events, showing that initially non-dominant trajectories can later surpass local greedy choices.The analysis covers 640 blocks per benchmark using LLaDA-1.5.
  • Qualitative Comparison: Mathematical Reasoning (MATH): SR-SMC recovers the correct multi-step conversion that greedy decoding misses through a persistent hallucinated identity.The comparison contrasts step-wise confidence with trajectory-level confidence and parallel trajectory exploration.
  • Analysis of Gumbel Noise: SR-SMC with N = 4 particles is more robust across sampling temperatures, steering generation away from repetitive or low-confidence regions.The figure compares standard parallel decoding with SR-SMC across LLaDA-1.5 and Dream-7B on MBPP and MATH.
  • Analysis of Gumbel Noise: SR-SMC consistently outperforms the baseline across most temperature settings for LLaDA-1.5, with the largest gains at moderate to high temperatures.These settings provide more diverse candidates from the stochastic diffusion process.
  • Zero-Shot Evaluation: SR-SMC enhances both LLaDA-1.5 and Dream-7B in zero-shot GSM8K and MATH evaluations without prompting demonstrations.Trajectory-level confidence helps navigate a wider exploration space and steer away from low-quality outputs.

6. Related Work

Related work has advanced masked diffusion language models through scaling, architectural changes, and training interventions, while inference-time scaling remains less explored.

  • Masked Diffusion Language Models: Masked diffusion language models and diffusion large language models have emerged as alternatives to autoregressive models for discrete text generation.The cited model families include MDLMs, LLaDAs, Dream, and DiffuLLaMA.
  • Inference-Time Scaling: Existing MDLMs primarily improve performance through model scaling, architectural modifications, or training-time interventions.The role of inference-time scaling is described as largely unexplored.

7. Conclusion

The paper formulates self-rewarding SMC for masked diffusion language models by assigning trajectory-confidence importance weights and recursively updating particles during sampling.

  • 7. Conclusion: The method uses trajectory confidence as importance weights to promote globally confident generation trajectories without additional training or external reward guidance.Experiments and ablations span multiple benchmarks and model families.
  • 7. Conclusion: The resulting framework is presented as an inference-time scaling dimension that improves pretrained MDLM sample quality and diversity.The conclusion attributes this result to experiments and ablation studies across multiple benchmarks and model families.
  • 7. Conclusion: SMC maintains intermediate path measures and recursively updates particle weights as trajectories are extended through the reverse diffusion process.The incremental weight updates combine with normalized-weight resampling.

A.2. Proof for Confidence-based Sequential Monte Carlo

The proof derives the SMC incremental weight for masked diffusion when the proposal equals the diffusion transition kernel, linking the resulting potential to token confidence.

  • A.2. Proof for Confidence-based Sequential Monte Carlo: Choosing the sequential proposal as the diffusion transition kernel simplifies the incremental importance weight expression.The proof explicitly sets qt−1(xt−1 | xt) = Kt(xt, xt−1).
  • A.2. Proof for Confidence-based Sequential Monte Carlo: The incremental importance weight is defined as the ratio between successive unnormalized path measures and the sequential proposal.This is the standard SMC weight used in the proof.
  • A.2. Proof for Confidence-based Sequential Monte Carlo: The potential is the joint probability of accepted tokens within S_t and acts as the self-rewarding confidence signal for the SMC update.The proof connects the potential to the token confidence used by the method.

B. Limitation and Future Work

The method has higher inference-time computation because it runs multiple diffusion processes in parallel, and its likelihood-based trajectory confidence does not directly optimize downstream objectives.

  • Running multiple diffusion processes in parallel increases inference-time computation, although particle count can control this trade-off.This cost is described as inherent to inference-time scaling methods.
  • Likelihood-only trajectory confidence is task-agnostic but does not explicitly optimize reasoning correctness or human preference.Future work proposes look-ahead or twisted diffusion transitions to improve sampling efficiency and quality.

C.1. Entropy Results of Text Generation

The text-generation evaluation reports generative perplexity and entropy for samples of two lengths, comparing them with the original data to assess quality and diversity.

  • The original data has generative perplexity 14.8 and entropy 5.44, providing reference values for the generated text.The paper states these references demonstrate improved sample quality while preserving text diversity.
  • Generative perplexity and entropy are reported for 300 samples with lengths L = 1024 and 2048.The table also reports the number of function evaluations and identifies OWT as the training dataset.

C.2. Detailed Results of Inference with Gumbel Noise

Across Gumbel noise temperatures from 0.0 to 1.0, SR-SMC outperforms the LLaDA-1.5 and Dream-7B baselines on MBPP and MATH, while qualitative comparisons connect particle resampling with global coherence.

  • C.2. Detailed Results of Inference with Gumbel Noise: SR-SMC consistently outperforms LLaDA-1.5 and Dream-7B across Gumbel noise temperatures from 0.0 to 1.0 on MBPP and MATH.The reported robustness spans both benchmarks and the full tested temperature range.
  • C.2. Detailed Results of Inference with Gumbel Noise: Dream-7B performance degrades severely when its noise temperature increases slightly to 0.1 and 0.2, whereas SR-SMC improves robustness across temperatures.The comparison highlights sensitivity in the baseline and greater stability for the proposed method.
  • C.2. Detailed Results of Inference with Gumbel Noise: Qualitative comparisons use arithmetic and physical reasoning examples to contrast greedy decoding paths with self-rewarding SMC paths.Figures 6 and 7 provide additional path comparisons, including distance calculation examples.
  • C.2. Detailed Results of Inference with Gumbel Noise: Greedy decoding is prone to local consistency errors and calculation hallucinations, whereas SR-SMC maintains global coherence through particle resampling.This qualitative outcome is stated for the compared reasoning paths.
Loading 2602.01849v1…