arXiv:2602.04344cs.LGcs.AI2026-02中稿 · ICML被引 1

用确定性搜索提升掩码扩散模型生成质量,不依赖随机采样

UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching

  • 将未掩码路径建模为搜索树,用确定性动作分支探索生成路径
  • 在复杂编码任务上超越现有测试时扩展方法,数学推理也表现优异
  • 适合需要高质量文本生成的开发者和研究者

测试时扩展策略已有效利用推理阶段算力提升自回归大语言模型的推理能力。本文表明,掩码扩散语言模型(MDLMs)因其迭代且非自回归的生成过程,天然适合高级搜索策略。为此,我们提出UnMaskFork(UMF),将未掩码轨迹建模为搜索树,并采用蒙特卡洛树搜索优化生成路径。与依赖随机采样的标准扩展方法不同,UMF通过多个MDLM执行确定性部分解掩动作来探索搜索空间。实证评估显示,UMF在复杂编码基准上持续优于现有测试时扩展基线,同时在数学推理任务中也表现出强可扩展性。

原文摘要 · Abstract (English)

Test-time scaling strategies have effectively leveraged inference-time compute to enhance the reasoning abilities of Autoregressive Large Language Models. In this work, we demonstrate that Masked Diffusion Language Models (MDLMs) are inherently amenable to advanced search strategies, owing to their iterative and non-autoregressive generation process. To leverage this, we propose UnMaskFork (UMF), a framework that formulates the unmasking trajectory as a search tree and employs Monte Carlo Tree Search to optimize the generation path. In contrast to standard scaling methods relying on stochastic sampling, UMF explores the search space through deterministic partial unmasking actions performed by multiple MDLMs. Our empirical evaluation demonstrates that UMF consistently outperforms existing test-time scaling baselines on complex coding benchmarks, while also exhibiting strong scalability on mathematical reasoning tasks.

扩散模型生成优化测试时扩展

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。