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

TL;DR

This paper introduces Order-Token Search, a novel decoding method for Diffusion Language Models that jointly explores generation order and token space, significantly improving performance on reasoning and coding benchmarks.

Contribution

It proposes a joint search approach for DLM decoding, combining order and token exploration with a likelihood estimator for stable pruning, advancing decoding strategies.

Findings

Outperforms baselines on GSM8K, MATH500, Countdown, and HumanEval.

Achieves 3.1%, 3.8%, 7.9%, and 6.8% absolute improvements respectively.

Matches or surpasses diffu-GRPO post-trained models.

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.