Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language Models

📄 arXiv: 2602.12586v1 📥 PDF

作者: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia

分类: cs.AI

发布日期: 2026-02-13

备注: 8 pages, preprint

💡 一句话要点

提出McDiffuSE，利用蒙特卡洛树搜索优化扩散语言模型中的槽填充顺序，提升生成质量。

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 扩散语言模型 蒙特卡洛树搜索 槽填充 序列生成 代码生成

📋 核心要点

  1. 掩码扩散模型在计划-填充解码中对槽填充顺序敏感，导致性能不稳定和输出差异大。
  2. McDiffuSE将槽选择建模为决策过程，利用蒙特卡洛树搜索优化填充顺序，提升生成质量。
  3. 实验表明，McDiffuSE在代码和数学推理任务上显著优于现有方法，尤其在MBPP和MATH500数据集上。

📝 摘要（中文）

本文提出McDiffuSE框架，该框架将槽选择视为决策过程，并通过蒙特卡洛树搜索（MCTS）优化填充顺序，以提升掩码扩散模型（MDM）的性能。MDM在数学和代码推理方面展现出潜力，但其性能对槽填充顺序高度敏感，导致输出差异较大。McDiffuSE通过前瞻性模拟评估部分完成情况，从而系统地探索生成顺序的组合空间。实验表明，McDiffuSE相比自回归基线平均提升3.2%，相比基线计划-填充方法提升8.0%，在MBPP和MATH500上分别显著提升19.5%和4.9%。分析表明，虽然McDiffuSE主要遵循顺序排序，但结合非顺序生成对于最大化性能至关重要。研究发现，更大的探索常数而非增加模拟次数，对于克服模型置信度偏差并发现有效的排序是必要的。这些发现确立了基于MCTS的规划作为增强MDM生成质量的有效方法。

🔬 方法详解

问题定义：论文旨在解决掩码扩散模型（MDM）中槽填充顺序对生成质量的敏感性问题。现有的计划-填充方法性能高度依赖于槽的填充顺序，不同的填充顺序会导致输出结果的巨大差异，这限制了MDM在复杂推理任务中的应用。

核心思路：论文的核心思路是将槽填充顺序的选择视为一个决策过程，并利用蒙特卡洛树搜索（MCTS）来寻找最优的填充顺序。MCTS通过模拟生成过程，评估不同填充顺序的潜在结果，从而选择能够最大化生成质量的顺序。

技术框架：McDiffuSE框架主要包含以下几个阶段：1) 槽识别：识别需要填充的槽位。2) MCTS搜索：使用MCTS算法搜索最佳的填充顺序。MCTS的每个节点代表一个部分完成的序列，边代表一个槽的填充。MCTS通过模拟（rollout）评估每个节点的价值，并根据价值选择下一步要探索的节点。3) 槽填充：根据MCTS选择的顺序，使用MDM填充槽位。

关键创新：论文的关键创新在于将MCTS引入到MDM的槽填充顺序优化中。与传统的固定或启发式填充顺序相比，MCTS能够根据模型的反馈动态地调整填充顺序，从而更好地适应不同的任务和数据。此外，论文还发现，更大的探索常数对于克服模型置信度偏差至关重要。

关键设计：MCTS的奖励函数设计至关重要，论文中奖励函数的设计目标是最大化生成序列的质量。探索常数（exploration constant）的设置影响MCTS的探索-利用平衡，论文发现较大的探索常数能够帮助MCTS发现更好的填充顺序。具体的网络结构和损失函数沿用了MDM的设置，没有进行特别的修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，McDiffuSE在MBPP数据集上提升了19.5%，在MATH500数据集上提升了4.9%，平均超过自回归基线3.2%，超过基线计划-填充方法8.0%。这些结果表明，通过MCTS优化槽填充顺序可以显著提高MDM的生成质量，尤其是在需要复杂推理的任务中。

🎯 应用场景

该研究成果可应用于各种需要序列生成的任务，例如代码生成、数学问题求解、文本摘要等。通过优化生成顺序，可以显著提高生成质量和效率，从而在软件开发、教育、自然语言处理等领域发挥重要作用。未来，该方法有望扩展到更复杂的生成任务中，例如多模态内容生成。

📄 摘要（原文）

While plan-and-infill decoding in Masked Diffusion Models (MDMs) shows promise for mathematical and code reasoning, performance remains highly sensitive to slot infilling order, often yielding substantial output variance. We introduce McDiffuSE, a framework that formulates slot selection as decision making and optimises infilling orders through Monte Carlo Tree Search (MCTS). McDiffuSE uses look-ahead simulations to evaluate partial completions before commitment, systematically exploring the combinatorial space of generation orders. Experiments show an average improvement of 3.2% over autoregressive baselines and 8.0% over baseline plan-and-infill, with notable gains of 19.5% on MBPP and 4.9% on MATH500. Our analysis reveals that while McDiffuSE predominantly follows sequential ordering, incorporating non-sequential generation is essential for maximising performance. We observe that larger exploration constants, rather than increased simulations, are necessary to overcome model confidence biases and discover effective orderings. These findings establish MCTS-based planning as an effective approach for enhancing generation quality in MDMs.