摘要
测试时扩展策略已有效利用推理计算提升了自回归大语言模型的推理能力。本文证明,得益于其迭代和非自回归的生成过程,掩码扩散语言模型(MDLMs)天然适用于高级搜索策略。为此,我们提出 UnMaskFork (UMF) 框架,将去掩码轨迹构建为搜索树,并采用蒙特卡洛树搜索优化生成路径。与依赖随机采样的标准扩展方法不同,UMF 通过多个 MDLM 执行的确定性部分去掩码动作探索搜索空间。实证评估表明,UMF 在复杂编码基准上持续优于现有测试时扩展基线,同时在数学推理任务中展现出强大的可扩展性。
AI 推荐理由
论文核心提出利用测试时扩展和搜索策略增强掩码扩散模型的推理能力,并在数学和编码任务中验证。
研究机构
Sakura AI, Tokyo, Japan
论文信息