Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement

arXiv:2507.08390v4 Announce Type: replace 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.

Related

• BADiff: Bandwidth Adaptive Diffusion Model
• Diffusion Processes on Implicit Manifolds
• DiffSketcher: Text Guided Vector Sketch Synthesis through Latent Diffusion Models
• Face2Scene: Using Facial Degradation as an Oracle for Diffusion-Based Scene Restoration
• ODE-free Neural Flow Matching for One-Step Generative Modeling

Source: arXiv cs.LG | 2026-04-10