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

Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Stefano Ermon

Keywords: Diffusion Language Model, Masked Diffusion Model, Test-Time Search Algorithm

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

Abstract: The order-agnostic generation of Diffusion Language Models (DLMs) presents a promising alternative to autoregressive models for complex reasoning. We model reasoning as traversals of a problem-specific graph of logical dependencies, and view DLM decoding as sampling trajectories from a joint space over generation orders and token values. We show that standard decoding heuristics such as low-confidence remasking collapse this reasoning space. 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, matching or surpassing the gains of fully post-trained d1-LLaDA with diffu-GRPO on Countdown, GSM8K, and MATH500 (e.g. achieving a 13.7% absolute gain on Countdown). Our work establishes structured search as a key missing component for advancing reasoning in DLMs.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 1903