English

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

Computation and Language 2026-02-06 v2 Machine Learning

Abstract

Diffusion Language Models (DLMs) offer order-agnostic generation that can explore many possible decoding trajectories. However, current decoding methods commit to a single trajectory, limiting exploration in trajectory space. We introduce Order-Token Search to explore this space through jointly searching over generation order and token values. Its core is a likelihood estimator that scores denoising actions, enabling stable pruning and efficient exploration of diverse trajectories. Across mathematical reasoning and coding benchmarks, Order-Token Search consistently outperforms baselines on GSM8K, MATH500, Countdown, and HumanEval (3.1%, 3.8%, 7.9%, and 6.8% absolute over backbone), matching or surpassing diffu-GRPO post-trained d1-LLaDA. Our work establishes joint search as a key component for advancing decoding in DLMs.

Keywords

Cite

@article{arxiv.2601.20339,
  title  = {Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space},
  author = {Yangyi Shen and Tianjian Feng and Jiaqi Han and Wen Wang and Tianlang Chen and Chunhua Shen and Jure Leskovec and Stefano Ermon},
  journal= {arXiv preprint arXiv:2601.20339},
  year   = {2026}
}
R2 v1 2026-07-01T09:23:25.634Z