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

TL;DR

This work tackles decoding in Diffusion Language Models (DLMs), where traditional approaches commit to a single trajectory and fail to explore the rich space of possible generation orders. It introduces Order-Token Search (OTS), a beam-like algorithm that jointly searches over generation order and token values, guided by a stable block-diffusion likelihood estimator for pruning. Across GSM8K, MATH500, Countdown, and HumanEval, OTS yields consistent pass@1 gains, often matching or surpassing post-training methods such as diffu-GRPO, and demonstrates superior scalability with increased compute compared to unstructured sampling. The results establish joint order-token search as a key component for enhancing reasoning in DLM decoding, with broad implications for robust, diverse, and efficient inference in diffusion-based language modeling.

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.

Figures (7)

• Figure 1: Why heuristic remasking fails, and how Order-Token Search helps. (a) Low-confidence remasking greedily fixes high-confidence tokens early and only revises low-confidence positions, which effectively commits to a narrow set of decoding trajectories and can get stuck in an incorrect completion. (b) Random remasking uniformly samples different remasking patterns, producing multiple distinct trajectories (different orders and token assignments), but offers no principled way to select the correct trajectory among competing candidates. (c) Order-Token Search maintains beams of partial trajectories, periodically expands each trajectory by proposing denoising actions that change both which position to update (order) and what token to write (token), and then prunes partial trajectories using a likelihood estimation function to filter incorrect and suboptimal candidates before continuing decoding. (d) MATH500 accuracy plot (example dataset) illustrating Order-Token Search’s advantage over low-confidence and random remasking.
• Figure 2: Decoding strategies exhibit an exploration--exploitation trade-off. Empirical pass@$k$ curves for LLaDA-8B-Instruct and LLaDA-1.5 on reasoning and coding benchmarks. Low-confidence remasking often attains higher pass@$1$, but yields slow pass@$k$ growth as $k$ increases. In contrast, random remasking consistently, and autoregressive (AR) decoding in several settings, starts from lower pass@$1$ yet achieves substantially higher pass@$k$ at large $k$ ($\approx$256), reflecting broader coverage of the solution space through more diverse trajectories. These curves highlight an opportunity for decoding algorithms that translate high-pass@$k$ potential into stronger pass@$1$ by explicitly exploring and selecting among diverse trajectories.
• Figure 3: Illustration of a pruning stage in Order‑Token Search for DLMs. At a search step, we have $2$ fully denoised sequences (on the leftmost), with yellow tokens unmasked in previous steps. We then mask the current block (the middle $3$ tokens) and measure its likelihood through feeding each masked candidate into the DLM to obtain each token's probability. The score function computes the chain-rule product of token probabilities and prunes the lower-likelihood candidate.
• Figure 4: Countdown accuracy versus test-time compute (NFE) for OTS and majority-voting baselines. For each method, we vary beam size (OTS) or the number of samples (AR+MV, Random+MV), and choose the largest configuration so that all right-most points have roughly matched NFE. At this matched-compute point, OTS with beam size 6 attains 29.3% accuracy, compared to 19.9% for AR+MV and 18.4% for Random+MV, indicating more efficient use of additional FLOPs than simply drawing more independent diffusion samples.
• Figure S1: Correlation between generation order and accuracy. The x-axis shows how "chaotic" a generated sample is measured in the Hamming distance of its decoding order from a strict left-to-right (AR) order. The y-axis is the average accuracy of the generated samples with the same chaotic value for a problem. The size of the point represents the number of samples. We find no correlation, indicating that decoding in an AR-like order is not a predictor of success.