Search or Accelerate: Confidence-Switched Position Beam Search for Diffusion Language Models
Abstract
Diffusion Language Models (DLMs) generate text by iteratively denoising a masked sequence, repeatedly deciding which positions to commit at each step. Standard decoding follows a greedy rule, unmasking the most confident positions, yet this local choice can lock the model into a suboptimal unmasking order, especially on reasoning-heavy prompts. We present Search Or AcceleRate (SOAR), a training-free decoding algorithm that adapts its behavior to the model’s uncertainty. When confidence is low, SOAR briefly widens the search over alternative unmasking decisions to avoid premature commitments; when confidence is high, it collapses the search and decodes many positions in parallel to reduce the number of denoising iterations. Across mathematical reasoning and code generation benchmarks (GSM8K, MBPP, HumanEval) on Dream-7B and LLaDA-8B, SOAR improves generation quality while maintaining competitive inference speed, offering a practical way to balance quality and efficiency in DLM decoding.
Lay Summary
Diffusion Language Models start with a fully masked sentence and gradually reveal words, but the standard strategy of always committing to the most confident positions can trap the model in poor directions on hard tasks like math or coding. We introduce SOAR, a training-free decoding algorithm that adapts to the model's uncertainty: when uncertain, it briefly explores alternative choices before committing; when confident, it fills many positions at once to speed things up. On mathematical reasoning and code generation benchmarks, SOAR improves output quality while keeping inference speed competitive, and we show it can be seamlessly combined with other decoding methods for further gains.