用蒙特卡洛树搜索优化扩散语言模型的生成路径,提升文本质量。
Diffusion Language Model Inference with Monte Carlo Tree Search
- 在初始化阶段引入蒙特卡洛树搜索,探索更优的去掩码策略。
- 相比现有方法,生成质量最高提升22.0%,且无需额外训练。
- 适合追求高质量生成、可扩展推理的开发者与研究者。
扩散语言模型(DLMs)作为自回归生成的替代方案,具备并行生成和更强全局连贯性的优势。推理时,DLM通过迭代去噪掩码序列实现文本生成;然而,确定哪些位置解码以及选择哪些词构成一个巨大的组合搜索问题。现有方法依赖启发式近似,常导致次优路径;另一些方法则需额外训练引导选词。为引入一种原则性搜索机制,我们提出MEDAL——一个将蒙特卡洛树搜索用于扩散语言模型推理初始化的推理时扩展框架。在初始化阶段使用蒙特卡洛树搜索探索有前景的去掩码轨迹,为后续精炼提供稳健起点。该设计实现高效的推理时扩展,生成质量随搜索预算增加而提升,无需额外训练。在多个基准测试中,MEDAL相比现有推理策略最高提升22.0%,确立了基于搜索的DLM推理新范式。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。