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

• URL: http://arxiv.org/abs/2601.20339v2
• Date: Thu, 05 Feb 2026 02:28:16 GMT
• Title: Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space
• Authors: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon,
• Abstract summary: Diffusion Language Models (DLMs) offer order-agnostic generation that can explore many possible decoding trajectories.<n>We introduce Order-Token Search to explore this space through jointly searching over generation order and token values.
• Score: 110.80564213032729
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• AlgoVeri: An Aligned Benchmark for Verified Code Generation on Classical Algorithms [54.99368693313797]
  Existing benchmarks test only individual languages/tools, so the performance numbers are not directly comparable.<n>We address this gap with AlgoVeri, a benchmark that evaluates vericoding of $77$ classical algorithms in Dafny, Verus, and Lean.
  arXiv  Detail & Related papers  (2026-02-10T06:58:26Z)
• Increasing LLM Coding Capabilities through Diverse Synthetic Coding Tasks [41.75017840131367]
  Large language models (LLMs) have shown impressive promise in code generation.<n>We present a scalable synthetic data generation pipeline that produces nearly 800k instruction-reasoning-code-test quadruplets.
  arXiv  Detail & Related papers  (2025-10-27T10:54:25Z)
• Zero-Shot Cross-Domain Code Search without Fine-Tuning [12.905068305900356]
  We propose a zero-shot, fine-tuning-free approach for cross-domain code search.<n>CodeBridge combines query-code, query-comment, and code-code matching via PLM-based similarity scoring and sampling-based fusion.<n>Our approach outperforms the state-of-the-art PLM-based code search approaches, i.e., CoCoSoDa and UniXcoder, by an average of 21.4% and 24.9% in MRR, respectively.
  arXiv  Detail & Related papers  (2025-04-10T13:36:37Z)
• ETS: Efficient Tree Search for Inference-Time Scaling [61.553681244572914]
  One promising approach for test-time compute scaling is search against a process reward model.<n> diversity of trajectories in the tree search process affects the accuracy of the search, since increasing diversity promotes more exploration.<n>We propose Efficient Tree Search (ETS), which promotes KV sharing by pruning redundant trajectories while maintaining necessary diverse trajectories.
  arXiv  Detail & Related papers  (2025-02-19T09:30:38Z)
• CodeXEmbed: A Generalist Embedding Model Family for Multiligual and Multi-task Code Retrieval [103.116634967815]
  We introduce CodeXEmbed, a family of large-scale code embedding models ranging from 400M to 7B parameters.<n>Our novel training pipeline unifies multiple programming languages and transforms various code-related tasks into a common retrieval framework.<n>Our 7B model sets a new state-of-the-art (SOTA) in code retrieval, outperforming the previous leading model, Voyage-Code, by over 20% on CoIR benchmark.
  arXiv  Detail & Related papers  (2024-11-19T16:54:45Z)
• Planning In Natural Language Improves LLM Search For Code Generation [5.370466208990696]
  We propose PlanSearch, a novel search algorithm for solving problems in natural language. PlanSearch shows strong results across HumanEval+, MBPP+, and LiveCodeBench. We show that, across all models, search algorithms, and benchmarks analyzed, we can accurately predict performance gains due to search.
  arXiv  Detail & Related papers  (2024-09-05T17:44:49Z)
• D2LLM: Decomposed and Distilled Large Language Models for Semantic Search [18.63768158439252]
  We present D2LLMs-Decomposed and Distilled LLMs for semantic search. We decompose a cross-encoder into an efficient bi-encoder integrated with Pooling by Multihead Attention and an Interaction Emulation Module. Our experiments show that D2LLM surpasses five leading baselines in terms of all metrics across three tasks.
  arXiv  Detail & Related papers  (2024-06-25T04:03:04Z)
• Learning to Rank in Generative Retrieval [62.91492903161522]
  Generative retrieval aims to generate identifier strings of relevant passages as the retrieval target. We propose a learning-to-rank framework for generative retrieval, dubbed LTRGR. This framework only requires an additional learning-to-rank training phase to enhance current generative retrieval systems.
  arXiv  Detail & Related papers  (2023-06-27T05:48:14Z)
• Improving Code Search with Hard Negative Sampling Based on Fine-tuning [15.341959871682981]
  We introduce a cross-encoder architecture for code search that jointly encodes the concatenation of query and code. We also introduce a Retriever-Ranker (RR) framework that cascades the dual-encoder and cross-encoder to promote the efficiency of evaluation and online serving.
  arXiv  Detail & Related papers  (2023-05-08T07:04:28Z)
• Revisiting Code Search in a Two-Stage Paradigm [67.02322603435628]
  TOSS is a two-stage fusion code search framework. It first uses IR-based and bi-encoder models to efficiently recall a small number of top-k code candidates. It then uses fine-grained cross-encoders for finer ranking.
  arXiv  Detail & Related papers  (2022-08-24T02:34:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.