Title: Improving Diffusion Language Model Reasoning through Joint Search in Generation Order and Token Space
Primary area: foundation or frontier models, including LLMs
Keywords: Diffusion Language Model, Masked Diffusion Model, Test-Time Search Algorithm

Abstract: The order-agnostic generation of Diffusion Language Models (DLMs) presents a promising alternative to autoregressive models for complex reasoning. However, standard decoding strategies fail to effectively explore the vast search space. We analyze common heuristics like low-confidence remasking and identify a critical trade-off: while they improve single-trial accuracy, their greedy nature prematurely narrows the reasoning path. To address this, we introduce **Order-Token Search**, an algorithm that jointly searches over token content and generation order. Its core is a likelihood estimation function that scores block-level denoising actions, enabling stable path pruning. This allows for efficient exploration of diverse reasoning trajectories. Extensive experiments on mathematical reasoning and planning benchmarks show that our method consistently outperforms baselines, achieving a 13.7\% absolute gain on the challenging Countdown task—an improvement comparable to costly post-training. Our work establishes structured search as a key missing component for advancing reasoning in DLMs.

TL;DR: Order-Token Search enhances Diffusion LM reasoning by jointly searching generation order and token space, overcoming greedy methods' limitations.