Diffusion Language Model Inference with Monte Carlo Tree Search

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Huang, Zheng, Ramnath, Kiran, Chen, Yueyan, Feng, Aosong, Woo, Sangmin, Srinivasan, Balasubramaniam, Xu, Zhichao, Zhou, Kang, Wang, Shuai, Ding, Haibo, Cheong, Lin Lee
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866914296039997440
author Huang, Zheng
Ramnath, Kiran
Chen, Yueyan
Feng, Aosong
Woo, Sangmin
Srinivasan, Balasubramaniam
Xu, Zhichao
Zhou, Kang
Wang, Shuai
Ding, Haibo
Cheong, Lin Lee
author_facet Huang, Zheng
Ramnath, Kiran
Chen, Yueyan
Feng, Aosong
Woo, Sangmin
Srinivasan, Balasubramaniam
Xu, Zhichao
Zhou, Kang
Wang, Shuai
Ding, Haibo
Cheong, Lin Lee
contents Diffusion language models (DLMs) have recently emerged as a compelling alternative to autoregressive generation, offering parallel generation and improved global coherence. During inference, DLMs generate text by iteratively denoising masked sequences in parallel; however, determining which positions to unmask and which tokens to commit forms a large combinatorial search problem. Existing inference methods approximate this search using heuristics, which often yield suboptimal decoding paths; other approaches instead rely on additional training to guide token selection. To introduce a principled search mechanism for DLMs inference, we introduce MEDAL, an inference-time scaling framework that integrates Monte Carlo Tree SEarch initialization for Diffusion LAnguage Model inference. We employ Monte Carlo Tree Search at the initialization stage to explore promising unmasking trajectories, providing a robust starting point for subsequent refinement. This design enables efficient inference-time scaling, allowing generation quality to improve as the search budget increases, without additional training. Across multiple benchmarks, MEDAL achieves up to 22.0% improvement over existing inference strategies, establishing a new paradigm for search-based inference in DLMs.
format Preprint
id arxiv_https___arxiv_org_abs_2512_12168
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Diffusion Language Model Inference with Monte Carlo Tree Search
Huang, Zheng
Ramnath, Kiran
Chen, Yueyan
Feng, Aosong
Woo, Sangmin
Srinivasan, Balasubramaniam
Xu, Zhichao
Zhou, Kang
Wang, Shuai
Ding, Haibo
Cheong, Lin Lee
Computation and Language
Artificial Intelligence
Diffusion language models (DLMs) have recently emerged as a compelling alternative to autoregressive generation, offering parallel generation and improved global coherence. During inference, DLMs generate text by iteratively denoising masked sequences in parallel; however, determining which positions to unmask and which tokens to commit forms a large combinatorial search problem. Existing inference methods approximate this search using heuristics, which often yield suboptimal decoding paths; other approaches instead rely on additional training to guide token selection. To introduce a principled search mechanism for DLMs inference, we introduce MEDAL, an inference-time scaling framework that integrates Monte Carlo Tree SEarch initialization for Diffusion LAnguage Model inference. We employ Monte Carlo Tree Search at the initialization stage to explore promising unmasking trajectories, providing a robust starting point for subsequent refinement. This design enables efficient inference-time scaling, allowing generation quality to improve as the search budget increases, without additional training. Across multiple benchmarks, MEDAL achieves up to 22.0% improvement over existing inference strategies, establishing a new paradigm for search-based inference in DLMs.
title Diffusion Language Model Inference with Monte Carlo Tree Search
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2512.12168