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

Problem Statement

Diffusion Language Models (DLMs), particularly Masked Diffusion Models (MDMs), offer a distinct advantage over autoregressive (AR) models: order-agnostic generation. Unlike AR models that generate tokens strictly left-to-right, MDMs can denoise multiple positions in parallel, allowing for adaptive generation orders. However, current decoding methods fail to fully exploit this flexibility.

Existing decoding strategies typically commit to a single trajectory per run, governed by heuristic remasking rules:

• Low-Confidence Remasking: Greedily fixes high-confidence tokens early. While this often yields strong single-sample accuracy (pass@1), it restricts the search to a narrow set of trajectories, leading to slow growth in multi-sample coverage (pass@k) and potential entrapment in incorrect local optima.
• Random Remasking: Samples remasking patterns uniformly to maximize trajectory diversity. While this improves pass@k by exploring a broader solution space, it suffers from weak pass@1 because it lacks a principled mechanism to select the correct trajectory among competing candidates.

This creates an exploration–exploitation trade-off: current methods either over-exploit local confidence (missing valid solutions) or explore blindly (failing to select the best solution). The paper argues that the potential of DLMs to explore diverse generation orders and token values remains unrealized because standard decoding does not perform explicit search in this joint space.

Methodology: Order-Token Search

To address this, the authors propose Order-Token Search, a decoding algorithm that performs beam search in the joint space of generation orders and token values. The algorithm alternates between expanding candidate trajectories and pruning unpromising ones.

1. Search Process (Expansion)

Order-Token Search maintains a beam of $K$ partially denoised candidates. Unlike standard sampling which follows a fixed sequence of decisions, the algorithm explicitly branches over two sources of uncertainty at user-specified intervals:

• Generation Order: Which positions (tokens) to unmask next.
• Token Values: What tokens to assign to those positions.

To manage computational complexity, the search utilizes block diffusion. Instead of branching at every single denoising step (which would incur prohibitive overhead), the algorithm expands candidates at the boundaries of contiguous token blocks. Within each block, independent denoising occurs. This reduces the decision points from the total sequence length to the number of blocks, allowing for substantial diversity in both order and token assignments while keeping inference costs manageable.

2. Pruning Process (Likelihood Estimation)

The core innovation lies in the likelihood estimator used to score and prune partial trajectories. A standard MDM objective (likelihood of reconstructing the full clean sequence from a noisy state) is unstable for intermediate states with many remaining masks, as it forces the model to solve difficult infilling problems it was not optimized for during training.

The authors introduce an incremental action-level scoring function:

• Instead of scoring the entire state, the algorithm scores the specific denoising actions taken to transition from a more corrupted state $x_t$ to a less corrupted state $x_s$.
• The score $s(x_t; x_s)$ measures how confidently the model can re-predict the newly revealed tokens in $x_s$, conditioned on the model's full-sequence prediction $x_0$ (which serves as context).
• This approach evaluates whether a trajectory's past decisions (order and token choices) are supported by the model, providing a lower-variance and more discriminative signal for pruning.

The cumulative score for a candidate is the sum of these incremental scores over the search intervals. The algorithm retains the top-$K$ candidates after each expansion, effectively filtering out incorrect or suboptimal paths before continuing decoding.

Key Contributions

1. Identification of the Trade-off: The paper empirically demonstrates that standard MDM decoding methods occupy opposite ends of the exploration–exploitation spectrum. Low-confidence remasking over-exploits (high pass@1, low pass@k growth), while random remasking blindly explores (high pass@k, low pass@1).
2. Order-Token Search Algorithm: The proposal of a novel decoding algorithm that searches the joint space of generation orders and token values. It utilizes a stable, incremental likelihood estimator to prune partial trajectories, converting trajectory diversity into improved single-sample accuracy.
3. Empirical Validation: The method is shown to consistently outperform baselines across mathematical reasoning and coding benchmarks, matching or surpassing gains typically achieved through expensive post-training methods like diffu-GRPO.

Experimental Results

The authors evaluated Order-Token Search (OTS) on GSM8K, MATH500, Countdown, and HumanEval using LLaDA-8B-Instruct and LLaDA-1.5 (an RL post-trained variant).

• Performance Gains: OTS consistently achieved the highest average accuracy across benchmarks.
  • On LLaDA-8B, OTS improved pass@1 by +3.1% on GSM8K, +3.8% on MATH500, +7.9% on Countdown, and +6.8% on HumanEval compared to the low-confidence remasking baseline.
  • On LLaDA-1.5, OTS achieved similar relative improvements, reaching 40.4% average accuracy compared to 37.2% for the best baseline (Low-conf + MV).
• Comparison to Post-Training: OTS matched or surpassed the performance of models subjected to reinforcement learning post-training (diffu-GRPO), suggesting that inference-time search can be as effective as model retraining for improving reasoning capabilities.
• Ablation Studies:
  • Likelihood Estimator: The specific incremental scoring function (Eq. 2) was critical. Variants using joint likelihood of all revealed blocks or AR-style likelihoods performed significantly worse, confirming that the tailored estimator is necessary for stable pruning in MDMs.
  • Search Space: OTS outperformed "Order Search" (searching only order) and "Token Search" (searching only tokens) which used AR-style likelihoods, highlighting the necessity of joint search and the mismatch of AR heuristics with MDMs.
• Efficiency: OTS scales more effectively with test-time compute (NFE) than majority voting. At matched compute budgets, OTS achieved higher accuracy (e.g., 29.3% vs. 19.9% on Countdown) by leveraging structured search rather than unstructured sampling.

Significance and Claims

The paper claims that joint order–token search is a key component for advancing DLM decoding. By explicitly exploring the space of generation orders and token values and pruning based on a stable likelihood estimator, the method converts the inherent multi-trajectory potential of diffusion models into reliable single-sample performance.

The authors position this work as a solution to the "exploration–exploitation" bottleneck in current MDM decoding. They assert that their approach allows DLMs to escape the local optima that trap greedy methods and select correct solutions that random sampling misses, all without requiring additional model training. The work establishes that inference-time search strategies, specifically tailored to the parallel, order-agnostic nature of diffusion models, can yield systematic improvements in reasoning and coding tasks.