我能要你的顺序吗？蒙特卡洛树搜索用于扩散语言模型中的槽填充顺序优化 Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language Models
用MCTS优化掩码扩散语言模型的槽填充顺序，在推理任务上平均提升3.2%
前置知识
掩码扩散语言模型 (Masked Diffusion Models, MDMs)
掩码扩散语言模型是一种非自回归的文本生成范式，它通过迭代去噪过程生成文本，而非像传统自回归模型那样从左到右逐个预测token。MDMs假设目标token之间条件独立，允许任意顺序生成，这为发现超越严格从左到右轨迹的生成顺序提供了可能。典型的MDMs包括LLaDA和Dream等模型，它们通过掩码和逐步揭示token的方式工作。
本文的核心问题正是MDMs中的槽填充顺序优化，理解MDMs的工作原理是理解本文方法的基础。
蒙特卡洛树搜索 (Monte Carlo Tree Search, MCTS)
MCTS是一种启发式搜索算法，通过结合树搜索和蒙特卡洛模拟来平衡探索与利用。算法包含四个阶段：选择（使用UCT或PUCT准则选择子节点）、扩展（为未探索的动作创建子节点）、模拟（从当前状态执行随机或策略引导的rollout）、反向传播（将模拟结果沿路径向上传播）。MCTS在AlphaGo/AlphaZero中取得巨大成功，现在被广泛应用于游戏、规划和推理任务。
本文将MCTS应用于MDMs的槽顺序决策，理解MCTS的机制是理解MCDIFFUSE如何工作的关键。
Plan-and-infill框架
Plan-and-infill是一种为MDMs设计的解码框架，每次去噪迭代包含两个步骤：规划阶段选择一个子序列（槽），然后自回归地在该槽内填充token。这种方法通过将token生成分解为槽内依赖和槽间依赖来降低生成难度。ReFusion是目前唯一的plan-and-infill扩散语言模型实现。
MCDIFFUSE正是针对plan-and-infill框架中的规划阶段（即槽选择）进行优化，用MCTS替代启发式选择策略。
PUCT算法
PUCT (Predictor-Upper Confidence Tree) 是MCTS中用于动作选择的准则，它平衡利用项（经验平均Q值）和探索项（基于先验概率的奖励）。公式为 $PUCT(s, a) = Q(s, a) + c \cdot P(a|s) \cdot \frac{\sqrt{N_s}}{1+N_{s,a}}$，其中 $P(a|s)$ 是先验概率，$c$ 是探索常数。PUCT会优先探索先验概率高的动作，同时确保所有动作最终都被访问。
MCDIFFUSE使用PUCT进行树搜索中的动作选择，理解PUCT如何平衡探索与利用是理解其核心机制的关键。
槽 (Slot) 在MDMs中的概念
在plan-and-infill框架中，输出序列被划分为K个连续且不重叠的槽（slot），每个槽包含L个token。槽内token通过自回归方式生成，而槽本身可以按任意顺序生成。例如，一个Python函数可以被划分为语法声明槽、函数定义槽、逻辑体槽等。槽的顺序选择对最终生成质量有显著影响。
MCDIFFUSE优化的正是这些槽的生成顺序，理解槽的概念是理解问题建模的基础。
研究动机
掩码扩散语言模型（MDMs）虽然通过非自回归解码提供了推理效率，但在复杂推理任务上往往不如自回归模型（ARMs）。一个关键原因是MDMs的生成顺序敏感性：当相互依赖的token被同时生成而没有相互条件约束时，输出质量对生成顺序高度敏感。在plan-and-infill框架中，虽然槽内自回归生成降低了局部依赖难度，但槽间生成顺序仍然至关重要——错误的顺序会导致跨迭代传播的未建模依赖，破坏全局连贯性。当前MDMs使用的启发式规划器（如基于置信度的选择）往往无法考虑长期依赖，导致性能下降。例如，在代码生成任务中，错误地先生成import语句而不是函数定义可能导致后续生成陷入语法破碎状态。
本文的目标是本文的目标是开发一个无需训练的框架MCDIFFUSE，通过将槽选择形式化为决策问题并使用蒙特卡洛树搜索（MCTS）来优化MDMs中的槽填充顺序，从而缩小扩散模型与自回归模型在推理任务上的性能差距。具体目标包括：在六个推理基准测试上超越现有MDMs和plan-and-infill方法；通过前瞻模拟评估部分完成情况，系统探索槽顺序的组合空间；理解非顺序生成何时以及为何对性能至关重要。
与已有工作不同的是，本文的独特切入角度在于将槽选择问题视为序列决策问题而非简单的启发式排序。与现有方法直接使用模型置信度贪婪选择不同，MCDIFFUSE通过MCTS的前瞻能力探索多个可能的生成路径，评估长期轨迹连贯性而非仅仅局部置信度。这种视角的转变抓住了一个被忽视的关键点：槽选择不是静态优化问题，而是动态规划问题，当前的槽选择会影响后续槽的生成质量。此外，本文发现大探索常数比增加模拟次数更能有效克服模型置信度偏差，这是对MCTS在文本生成中应用的深入洞察。
核心方法
MCDIFFUSE的核心思想可以用一个比喻来理解：想象你在建造一座乐高城堡，每块乐高代表一个代码槽。贪婪方法会先拿起看起来最容易放置的那块，但这可能导致后续拼接困难。MCDIFFUSE则像一个有远见的建筑师，会模拟多种拼接顺序，评估每种顺序下整座城堡的构建难度，然后选择最有希望的起始块。在技术路线上，MCDIFFUSE将槽选择形式化为确定性马尔可夫决策过程（MDP），状态编码已生成的槽和部分填充序列，动作为选择下一个要填充的槽，奖励基于槽级置信度。然后使用PUCT算法引导的MCTS进行树搜索，通过随机置信度rollout估计长期价值。
MCDIFFUSE最核心的创新点在于它使用MCTS的前瞻能力来克服模型内在的置信度偏差。传统方法直接使用模型置信度选择槽，但高置信度的槽可能不是最佳起始点——它可能将后续生成推向低概率区域。MCDIFFUSE通过模拟未来生成步骤来评估长期轨迹质量，而非仅仅局部置信度。具体来说，它使用混合价值函数 $V(s_t, a) = \lambda \cdot R(s_t, a) + (1-\lambda) \cdot E_{\pi_{roll}}[G]$，其中 $\lambda$ 平衡即时置信度与长期rollout回报。实验表明 $\lambda=0.3$ 效果最好，这意味着长期rollout反馈比即时置信度更重要。此外，本文发现大探索常数（c=50）比增加模拟次数更能有效探索，表明槽规划的主要挑战是足够的探索广度而非搜索深度。
方法步骤详情
MCDIFFUSE的执行流程包含四个MCTS阶段的迭代：(1) 选择阶段：从根节点 $s_0$ 开始，使用PUCT准则递归遍历树，选择子节点直到达到未扩展的叶节点。PUCT公式为 $PUCT(s,a) = Q(s,a) + c \cdot P(a|s) \cdot \frac{\sqrt{N_s}}{1+N_{s,a}}$，其中先验概率 $P(a|s)$ 由槽级置信度归一化得到。(2) 扩展阶段：当达到未扩展的叶节点 $s_t$ 时，为所有可行动作创建子节点，初始化访问计数和累积价值为零。(3) 模拟阶段：从扩展的子节点开始，使用温度缩放的softmax策略进行随机rollout直到所有槽被填充。轨迹分数为 $G = \frac{1}{T}\sum_{i=1}^{T} R(s_i, ilde{a}_i)$，最终价值为即时奖励与rollout回报的加权组合。(4) 反向传播阶段：将估计价值沿遍历路径向上传播，更新访问计数 $N_s, N_{s,a}$ 和累积价值 $W(s,a)$，计算经验平均动作价值 $Q(s,a) = W(s,a)/N_{s,a}$。经过 $N_{sim}$ 次模拟后，选择访问次数最多的动作 $a^* = \arg\max_{a} N_{s_t,a}$ 作为下一步。
技术新颖性
MCDIFFUSE的技术新颖性体现在三个方面：首先，它是第一个将MCTS应用于MDMs文本生成的工作，特别是针对槽顺序优化问题。与应用于自回归模型的MCTS不同，MCDIFFUSE需要处理槽间的离散决策而非token级别的连续生成。其次，MCDIFFUSE设计了专门的置信度感知rollout机制，使用温度缩放的随机采样而非确定性贪婪策略，以探索多样化的完成路径。最后，本文发现了探索广度比搜索深度更重要的洞察：大探索常数能有效克服模型置信度偏差，而增加模拟次数在低探索下反而可能导致过早收敛到局部最优。这与传统MCTS应用（如游戏AI）中增加模拟通常有效的经验形成对比。
实验结果
MCDIFFUSE在六个推理基准测试上进行了全面评估，结果表明其显著优于现有MDMs和自回归基线。在MBPP代码生成任务上，MCDIFFUSE相比ReFusion实现了19.45%的绝对准确率提升（从54.12%到73.57%），这是所有任务中最大的改进。在MATH500数学推理任务上，准确率从42.90%提升到47.80%，提升4.90%。在HumanEval代码任务上，提升最为显著，从62.05%跃升到78.37%，提升16.32%。在ARC Challenge常识推理上，提升0.70%（从87.98%到88.68%）。在GPQA-Diamond研究生级问答上，提升4.33%（从30.45%到34.78%）。在GSM8K数学推理上，MCDIFFUSE（87.91%）与Qwen2.5 7B（88.02%）相当，但优于ReFusion（85.64%）。平均而言，MCDIFFUSE达到68.52%的平均准确率，比最优ARM基线Qwen3 8B（65.33%）高出3.19个百分点。更重要的是，MCDIFFUSE在保持高准确率的同时，生成更紧凑的推理序列：在MATH500上平均token长度为152.2，而Qwen2.5 7B为436.0，token减少65.08%且准确率更高。统计分析表明，编码任务从MCTS规划中获益最大（3.13倍于非编码任务），这是因为代码组件间存在严格的结构依赖（如变量声明、函数定义、控制流）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K数学推理 | Pass@1 | 87.91% | 85.64% (ReFusion) | +2.27% |
| MATH500数学推理 | Pass@1 | 47.80% | 42.90% (ReFusion) | +4.90% |
| ARC Challenge常识推理 | Pass@1 | 88.68% | 87.98% (ReFusion) | +0.70% |
| GPQA-Diamond研究生问答 | Pass@1 | 34.78% | 30.45% (ReFusion) | +4.33% |
| MBPP代码生成 | Pass@1 | 73.57% | 54.12% (ReFusion) | +19.45% |
| HumanEval代码生成 | Pass@1 | 78.37% | 62.05% (ReFusion) | +16.32% |
局限与改进
尽管MCDIFFUSE取得了显著成果，但仍存在若干局限性。首先，MCDIFFUSE建立在ReFusion架构之上，继承了其固有的token长度限制（通常1024 tokens），这限制了其在需要长序列推理的任务上的表现。在MATH500上，当将自回归模型的上下文长度扩展到32768 tokens时，Qwen2.5 7B（63.40%）和Qwen3 8B（68.00%）均优于MCDIFFUSE（54.60%），表明扩展token长度对数学推理至关重要。其次，MCDIFFUSE的计算开销约为ReFusion的1.35倍（平均FLOPs从1.95E+16增加到2.63E+16），虽然这是可以接受的，但在资源受限场景下仍需权衡。第三，本文的实验主要基于7B参数的模型，尚不清楚MCDIFFUSE的改进是否能一致地扩展到更大的模型。最后，MCDIFFUSE的性能在不同任务间存在差异：在多选题（ARC、GPQA）上的改进较为 modest（<5%），而在代码生成任务上改进显著（>16%），这表明方法可能更适合具有强结构依赖的任务。
独立分析的弱点
从独立分析的角度，MCDIFFUSE存在几个值得关注的弱点及其改进方向。首先，MCDIFFUSE的探索机制依赖于固定的探索常数c，但最优c值可能因任务和模型而异。在实验中c=50在多数任务上表现良好，但在某些配置下c=100更好。改进方向是设计自适应探索常数，根据搜索状态动态调整c值，例如在搜索早期使用大c鼓励探索，在后期使用小c促进利用。其次，当前的rollout策略使用温度缩放的softmax采样，温度τ=0.5是固定的。一个改进方向是设计自适应温度调度，在rollout早期使用高温度探索多样化路径，在后期使用低温度聚焦于高价值轨迹。第三，MCDIFFUSE的槽划分是预定义的（如slot size=4），但最优的槽粒度可能因任务复杂度而异。改进方向是探索动态槽划分策略，根据代码结构或语义边界自适应调整槽大小。最后，当前方法在每次决策时需要运行256次MCTS模拟，计算成本较高。可以通过引入更高效的搜索剪枝策略或使用轻量级神经网络估计Q值来降低计算开销。
未来方向
本文作者和基于成果可延伸的未来研究方向包括多个有前景的路径。首先，作者明确指出扩展MCTS槽规划到其他扩散语言模型架构（如BD3LM、Eso-LMs）是重要方向，因为当前工作仅在ReFusion上验证。其次，本文发现探索广度比搜索深度更重要，这启发了更高效的探索策略研究，例如使用学习到的启发式函数引导搜索，或结合强化学习优化探索策略。第三，将MCDIFFUSE的思想扩展到视觉扩散模型的生成顺序优化是一个自然延伸，因为图像生成中patch或区域的生成顺序同样影响质量。第四，设计能处理更长序列的扩散模型架构是关键挑战，因为当前的token长度限制严重制约了复杂推理任务的表现。最后，将MCDIFFUSE的洞察应用于自回归模型的推理优化也值得探索，例如在chain-of-thought推理中规划推理步骤的顺序。
复现评估
从复现评估的角度，MCDIFFUSE具有较好的可复现性。论文提供了详细的算法伪代码（Algorithm 1-5），包括MCTS的四个阶段、rollout过程和反向传播的具体实现。所有实验使用4块NVIDIA H100 GPU进行，超参数设置明确：c=50、N_sim=256、τ=0.5、λ=0.3。数据集均为公开基准：GSM8K、MATH500、MBPP、HumanEval、ARC Challenge、GPQA-Diamond。基线模型ReFusion、LLaDA-8B、Dream 7B、Qwen2.5 7B、Qwen3 8B均为公开可用模型。论文提到代码和模型将在论文发表后开源。主要复现挑战在于需要足够的计算资源（4块H100 GPU）和理解MCTS与扩散模型交互的复杂性。此外，论文的实验设置使用了特定的评估框架（LM-Eval-Harness），复现时需要确保评估协议的一致性。
论文图表
