Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement

TL;DR

This paper introduces PG-DLM, an inference-time algorithm for diffusion language models that refines entire generation trajectories, improving reward-guided generation efficiency and accuracy without retraining.

Contribution

The paper proposes particle Gibbs sampling for diffusion language models, enabling trajectory-level refinement and adaptive compute allocation during inference.

Findings

PG-DLM outperforms prior methods on reward-guided tasks.

Achieves 90.07% accuracy on GSM8K with 2.9 particles.

Achieves 94.47% accuracy on GSM8K with 16 particles.

Abstract

Discrete diffusion models have recently emerged as strong alternatives to autoregressive language models, matching their performance through large-scale training. However, inference-time control remains relatively underexplored. In this work, we study how to steer generation toward desired rewards without retraining the models. Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement. We introduce particle Gibbs sampling for diffusion language models (PG-DLM), an inference-time algorithm enabling trajectory-level refinement. PG-DLM constructs a Markov chain over full denoising trajectories and applies a conditional sequential Monte Carlo kernel to resample them. By doing so, PG-DLM introduces a new scaling axis, the number of refinement iterations, which is unavailable to prior methods. Increasing iterations remains effective even as gains from adding more parallel samples saturate. Furthermore, PG-DLM enables adaptive compute allocation by performing additional iterations only when needed, leading to further efficiency gains. We derive theoretical guarantees for convergence and variance bounds, and analyze trade-offs across different scaling axes. Empirically, PG-DLM outperforms prior methods across compute budgets on reward-guided generation tasks. On GSM8K, it achieves 90.07% accuracy with 2.9 particles on average and 94.47% accuracy with 16 particles.