Search or Accelerate: Confidence-Switched Position Beam Search for Diffusion Language Models

TL;DR

SOAR is a confidence-adaptive decoding algorithm for diffusion language models that improves reasoning and code generation quality by dynamically balancing search breadth and speed based on model confidence.

Contribution

It introduces a training-free, confidence-switched decoding method that enhances diffusion language model performance without additional training.

Findings

Improves reasoning and code generation benchmarks.

Balances quality and efficiency effectively.

Maintains competitive inference speed.

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: unmask the most confident positions, yet this local choice can lock the model into a suboptimal unmasking order, especially on reasoning-heavy prompts. We present 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. Our Code is available at https://github.com/duterscmy/SOAR