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

📄 arXiv: 2602.04344v1 📥 PDF

作者: Kou Misaki, Takuya Akiba

分类: cs.LG, cs.AI

发布日期: 2026-02-04

💡 一句话要点

提出UnMaskFork，通过确定性动作分支实现Masked Diffusion模型测试时性能提升

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Masked Diffusion模型 测试时缩放 蒙特卡洛树搜索 确定性动作分支 语言生成

📋 核心要点

  1. 自回归大语言模型的测试时缩放策略已有效提升推理能力，但Masked Diffusion模型的潜力尚未充分挖掘。
  2. UnMaskFork框架将unmasking过程视为搜索树，利用蒙特卡洛树搜索优化生成路径，实现确定性的部分unmasking。
  3. 实验结果表明，UnMaskFork在编码和数学推理任务上均优于现有方法，展现出良好的性能和可扩展性。

📝 摘要（中文）

本文提出了一种名为UnMaskFork (UMF) 的框架，旨在利用推理时计算资源来增强Masked Diffusion语言模型 (MDLM) 的推理能力。MDLM的迭代和非自回归生成过程使其天然适合高级搜索策略。UMF将unmasking轨迹形式化为搜索树，并采用蒙特卡洛树搜索来优化生成路径。与依赖随机抽样的标准缩放方法不同，UMF通过多个MDLM执行的确定性部分unmasking动作来探索搜索空间。实验评估表明，UMF在复杂的编码基准测试中始终优于现有的测试时缩放基线，并在数学推理任务中表现出强大的可扩展性。

🔬 方法详解

问题定义：论文旨在解决Masked Diffusion语言模型（MDLM）在测试时如何有效利用计算资源来提升性能的问题。现有的测试时缩放方法主要依赖随机抽样，无法充分探索MDLM的生成空间，导致性能提升有限。

核心思路：论文的核心思路是将MDLM的unmasking过程建模为一个搜索树，通过确定性的动作分支来探索不同的生成路径。利用蒙特卡洛树搜索（MCTS）来指导搜索过程，选择最优的unmasking策略，从而提升生成质量。这种确定性搜索避免了随机抽样的盲目性，能够更有效地利用计算资源。

技术框架：UMF框架主要包含以下几个阶段：1) 初始化：将masked的输入作为搜索树的根节点。2) 选择：使用MCTS选择一个待扩展的节点。3) 扩展：对选定的节点执行确定性的部分unmasking动作，生成多个子节点。每个子节点代表一种可能的unmasking结果。4) 评估：使用MDLM对每个子节点进行评估，得到一个奖励值。5) 反向传播：将奖励值反向传播到搜索树中，更新节点的统计信息。重复执行选择、扩展、评估和反向传播，直到达到预定的计算资源限制。最终，选择根节点下访问次数最多的子节点作为最终的生成结果。

关键创新：UMF的关键创新在于将MDLM的unmasking过程形式化为搜索树，并采用确定性的动作分支和蒙特卡洛树搜索来优化生成路径。与传统的随机抽样方法相比，UMF能够更有效地探索生成空间，找到更优的unmasking策略。此外，UMF通过多个MDLM并行执行部分unmasking动作，进一步提升了搜索效率。

关键设计：UMF的关键设计包括：1) 动作空间：定义了确定性的部分unmasking动作，例如，每次unmasking概率最高的k个token。2) 奖励函数：设计了合适的奖励函数来评估每个节点的质量，例如，可以使用MDLM的perplexity或下游任务的性能作为奖励。3) MCTS参数：需要调整MCTS的参数，例如，探索-利用平衡参数，以控制搜索的广度和深度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，UnMaskFork在复杂的编码基准测试中始终优于现有的测试时缩放基线。例如，在HumanEval数据集上，UnMaskFork取得了显著的性能提升。此外，UnMaskFork在数学推理任务中也表现出强大的可扩展性，随着计算资源的增加，性能持续提升。

🎯 应用场景

UnMaskFork可应用于各种需要高质量文本生成的场景，例如代码生成、数学推理、机器翻译等。通过在测试时利用更多的计算资源，UnMaskFork能够显著提升生成质量，从而提高相关应用的性能和用户体验。该方法还可推广到其他类型的生成模型，具有广泛的应用前景。

📄 摘要（原文）

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.