💬 NLP

Diffusion Language Model Inference with Monte Carlo Tree Search

该论文介绍了 MEDAL，这是一个将蒙特卡洛树搜索（Monte Carlo Tree Search）集成到扩散语言模型中以优化去掩码轨迹的推理时扩展框架，在无需额外训练的情况下，实现了比现有启发式方法显著的性能提升。

原作者： Zheng Huang, Kiran Ramnath, Yueyan Chen, Aosong Feng, Sangmin Woo, Balasubramaniam Srinivasan, Zhichao Xu, Kang Zhou, Shuai Wang, Haibo Ding, Lin Lee Cheong

发布于 2026-02-03

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Zheng Huang, Kiran Ramnath, Yueyan Chen, Aosong Feng, Sangmin Woo, Balasubramaniam Srinivasan, Zhichao Xu, Kang Zhou, Shuai Wang, Haibo Ding, Lin Lee Cheong

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图写一个故事，但你开始时面对的是一页被黑色便利贴完全覆盖的纸。你的目标是一个接一个地揭开这些词，直到整个故事变得通顺。

这正是**扩散语言模型（Diffusion Language Models, DLMs）**的工作方式。与标准的 AI 那种像人类打字一样从左到右逐个单词生成的方式不同，DLM 会同时观察整个“被覆盖”的页面，并尝试猜测应该揭开哪些便利贴，以及这些便利贴下面隐藏着什么词。

问题在于？揭开这些便利贴的方式有无数种。如果你只是揭开那些目前看起来“最有可能”正确的词，你可能会陷入一个无法挽回的糟糕故事路径中。这就像是在没有考虑这个词如何影响整个段落的情况下，就匆忙选择了句子的第一个词。

这篇论文的作者提出了 MEDAL，一种更聪明的做法。他们将写作过程视为一种战略搜索，而不仅仅是一个简单的猜测游戏。

以下是他们解决方案的拆解，通过简单的类比来呈现：

1. “如果……会怎样”的探索者 (MCTS)

想象你是一位正在策划战役的将军。你不是盲目地根据直觉冲锋陷阵，而是派出几名侦察兵去探索地图上的不同路径。

论文的方法： 他们使用了**蒙特卡洛树搜索（Monte Carlo Tree Search, MCTS）**技术。你可以把它看作是一个“模拟引擎”。在 AI 决定揭开一批词之前，它会在脑海中运行成千上万次快速、微小的“如果……会怎样”的情景模拟。
目标： 它会问：“如果我现在揭开这个词，会让剩下的故事更容易写吗？还是会将我困在死胡同里？”
限制： 对整个故事进行这种模拟会耗时太长（就像为每一个动作都模拟一场完整的战争一样）。因此，MEDAL 仅在初始阶段（初始化阶段）使用这种强大的探索器来奠定坚实的基础。一旦路径确定，AI 就会切换到更快、更简单的方法来完成工作。

2. “置信度过滤器”（识别显而易见的事物）

“如果……会怎样”的探索者很聪明，但它无法检查字典中每一个可能的词。那是不可能的。

论文的方法： 他们使用了一个置信度引导过滤器（Confidence-Guided Filter）。想象一位图书管理员，他只允许你查看与主题最相关的 5 本书，而忽略其他数千本。
运作方式： AI 查看便利贴并说：“我有 90% 的把握这个词是‘猫’，但我只有 10% 的把握这个词是‘量子物理学’。”它会忽略低置信度的猜测，只对高置信度的猜测进行“如果……会怎样”的模拟。这使得搜索既快速又高效。

3. “信息增益”奖励（明智的选择）

当探索者选择一条路径时，它如何知道这条路径是否是好路径？

论文的方法： 他们使用了一种特殊的评分机制，称为信息增益（Information Gain）。
类比： 想象你正在玩拼图。如果你放下的一个碎片只能放在一个特定的位置，那是好事。但如果你放下的一个碎片不仅能填补空缺，还能帮助你确定另外五个碎片的去向，那简直太棒了。
结果： AI 获得的奖励不仅来自于猜对一个词，还来自于猜出一个能让后续部分更容易解决的词。它优先考虑那些能为未来减少困惑度的举动。

4. 拆解大任务 (Task Decomposition)

有时指令（提示词）非常复杂，以至于 AI 会感到不知所措，就像被要求一次性“写一部关于太空旅行的小说”。

论文的方法： 他们加入了任务拆解步骤。在写作之前，AI 被要求将大任务分解为更小、更易管理的步骤（例如：“1. 理解背景设定”、“2. 列出角色”、“3. 编写第一场戏”）。
结果： 这就像是一份路线图，引导 AI 循序渐进地完成那页复杂的便利贴任务，降低了迷失方向的可能性。

实验结果

作者在各种困难任务（如数学题、编程和阅读理解）上测试了这个“MEDAL”框架。

结果： 通过在开始阶段使用这种战略性的“如果……会怎样”搜索，结合智能过滤和任务拆解，AI 写出的故事和答案明显更好。
数据： 与其他方法相比，他们看到了高达 22% 的提升。
核心要点： 他们不需要重新训练 AI 或教给它新知识。他们只是给了它一种更好的在开始“写作”之前进行“思考”的策略。

总结： MEDAL 就像是给作家提供了一个“排练室”，让他们可以在正式动笔写下最终稿之前，快速尝试不同的开场白，看看哪一个能引向最好的故事。这种策略上的简单改变，让 AI 变得更加聪明且逻辑连贯。

技术摘要：MEDAL —— 基于蒙特卡洛树搜索的扩散语言模型推理

1. 问题陈述

扩散语言模型（DLMs）通过实现并行生成和提升全局连贯性，为自回归（AR）模型提供了一种极具吸引力的替代方案。然而，DLM 的推理过程面临着重大挑战：这是一个组合搜索问题，模型必须迭代地决定哪些掩码位置应当被取消掩码（unmask），以及选择哪些标记（tokens）进行提交。

现有的推理策略依赖于贪婪启发式算法（例如，优先取消掩码置信度最高的标记）或动态掩码调度。这些方法通常是近视的；一旦高置信度的标记被固定，后续步骤将被迫进行适应，从而频繁导致次优的轨迹。其他方法试图引导标记选择，但需要对辅助采样器进行额外训练，这增加了复杂性并限制了泛化能力。目前缺乏一种原则性的、无需训练的搜索机制来探索替代的取消掩码轨迹，以提高生成质量。

2. 方法论：MEDAL 框架

作者提出了 MEDAL（用于扩散语言模型推理的蒙特卡洛树搜索初始化），该框架将蒙特卡洛树搜索（MCTS）集成到 DLM 的推理阶段，而无需额外的模型训练。

核心组件

该框架通过三个阶段运行：

通过提示词引导的任务分解：
为了处理会导致高不确定性的复杂提示词，输入通过说明性示例进行增强，引导模型将任务分解为更小、更易处理的子任务（例如：理解输入、识别相关信息、合成答案）。这种结构化引导在搜索开始前降低了歧义性。
置信度引导的 MCTS 初始化：
MEDAL 并没有在整个生成过程中进行穷尽式的 MCTS 应用（这在计算上是难以承受的），而是将搜索限制在初始化阶段。
- 搜索空间缩减（置信度过滤）： 为了使搜索变得可行，作者引入了一种置信度引导的过滤机制。对于每个掩码位置，他们根据模型的 Logits 计算一个经过置信度调整的分数，其中包括熵惩罚（降低低置信度位置的权重）和前二差异（top-2 margin）（偏好在前两个候选标记之间具有明显区分度的位置）。这使得动作空间被限制在少量高概率的位置-标记对中。
- 信息增益奖励： MCTS 模拟不仅评估候选动作的即时标记选择，还评估它们提供的信息增益。奖励函数衡量的是在剩余掩码位置上预测熵的减少量。如果一个动作在解决掩码的同时，能增加模型对序列其余部分的预测置信度，则该动作会获得奖励。
- 初始化过程： MCTS 探索有前景的取消掩码轨迹，以构建一组鲁棒的候选部分取消掩码序列（长度为 $L_c$ ）。从中选择具有最高信息增益奖励的候选序列，作为最终生成的起点。
置信度引导的取消掩码（精炼）：
一旦选择了最优初始化，剩余的标记将使用高效的、置信度引导的启发式算法（基于置信度分数的采样）进行生成，不再进行进一步的树搜索。这平衡了早期阶段的高质量探索与全序列生成所需的效率。

理论保证

论文为该方法提供了理论依据。它将推理误差框架化为模型误差与联合依赖项之和。研究表明，MCTS 初始化通过最小化一个可计算的依赖误差代理（熵间隙），在探索的选项中理论上选择了能实现最小累计依赖误差上界的候选前缀。

3. 核心贡献

新颖的公式化： 论文将 DLM 推理视为一个搜索问题，并引入了 MEDAL，这是第一个将 MCTS 集成到 DLM 推理中的框架，实现了超越贪婪启发式的原则性探索。
新颖的设计： 作者设计了一种特定的推理方法，结合了 MCTS 引导的初始化模块和任务分解模块。关键创新包括用于管理搜索空间的置信度引导过滤机制，以及用于引导树搜索向提升全局模型置信度方向发展的信息增益奖励。
广泛的实验： 该方法在多个基准测试（GSM8K, ARC-C, HumanEval, MMLU, DROP, Countdown）上进行了评估，并使用了多种骨干 DLM（LLaDA, LLaDA1.5, Dream）。

4. 实验结果

评估表明，MEDAL 显著优于现有的推理策略：

性能提升： 在多个基准测试和骨干模型中，MEDAL 相比现有策略实现了高达 22.0% 的性能提升。例如，在采用 LLaDA 骨干模型的 DROP 数据集上，性能提升了 22.0%。
与基准对比： 该方法持续优于原始模型以及使用 Best-of-5 (bst5) 解码策略的模型。值得注意的是，经 MEDAL 增强后的 DLM（LLaDA）在大多数数据集上达到了与自回归 Llama-3 8B 模型相当或更优的结果。
消融实验： 移除任何组件（MCTS、任务分解或置信度过滤）都会导致性能下降，证实了每个部分的必要性。
超参数敏感性：
- 将任务分解为 3 个子任务 效果最佳。
- 在 MCTS 期间选择 5 个候选标记 在探索与噪声之间达到了最佳平衡。
- 增加 MCTS 初始化长度 ( $L_c$ ) 会提升性能，但在约 20 步 左右达到饱和点，此后收益相对于计算成本而言开始递减。
智能体工作流： 在将 MEDAL 与 ADAS 智能体系统集成的案例研究中，该方法进一步提升了在 DROP 和 MMLU 上的表现，表明其在复杂设置下增强了推理和规划能力。

5. 意义与主张

论文声称 MEDAL 建立了基于搜索的 DLM 推理的新范式。其主要意义在于：

无需训练的扩展性： 它实现了推理侧的扩展（Inference-time scaling），即随着搜索预算的增加，生成质量随之提升，且无需额外的训练或辅助采样器。
全局连贯性： 通过使用 MCTS 在过程早期探索替代轨迹，该方法避免了贪婪解码的近视缺陷，从而实现了更好的全局连贯性和推理能力。
实用性： 通过将 MCTS 限制在初始化阶段并对剩余部分使用高效启发式算法，该框架在处理大规模 DLM 推理时仍保持了计算上的可行性。

作者总结道，尽管 DLMs 具有强大的生成能力，但其在推理过程中的可控性和推理能力一直受到限制。MEDAL 通过提供一种原则性的搜索机制，有效地利用模型自身的分布来引导标记选择，从而解决了这一问题。