与Gemini Diffusion共振!首个扩散式「发散思维链」来了

近年来,思维链在大模型训练和推理中愈发重要。近日,西湖大学 MAPLE 实验室齐国君教授团队首次提出扩散式「发散思维链」—— 一种面向扩散语言模型的新型大模型推理范式。该方法将反向扩散过程中的每一步中间结果都看作大模型的一个「思考」步骤,然后利用基于结果的强化学习去优化整个生成轨迹,最大化模型最终答案的正确率。不同于始终单向推理、线性生成的传统思维链(CoT),扩散式「发散思维链」允许模型以任意顺序非线性生成,且在生成过程中无需严格遵从语法结构和可读性要求,能够鼓励模型以更加发散、创造性的方法开展推理。


扩散式「发散思维链」目前已成功应用于两种具有代表性的扩散语言模型中。在连续时间扩散语言模型中,该方法可以直接优化由模型输出的得分函数所确定的策略分布;而在离散时间扩散语言模型中,团队将预测不同掩码 Token 的顺序当作模型决策的一部分,并基于 Plackett-Luce 模型设计去掩码策略。据此,团队成功训练有序掩码生成扩散语言模型(Large Language Diffusion with Ordered Unmasking, LLaDOU)。实验表明,仅用公开数据集和 16 张 H800,经扩散式「发散思维链」增强后的模型即可在数学推理和代码生成任务上超越现有扩散语言模型。


扩散式「发散思维链」对基础大模型的训练与推理给出了重要启示:传统的自回归思维链语言模型通过线性预测下一个 token 生成答案并非唯一的选择范式。团队的研究揭示了通过优化 token 生成的顺序进行非线性语言生成是发散式思维的重要特点,对于在生成过程中逐步构建从早期概念要素的形成、到最终连接成具有完整想法和语法结构的回答起到了关键作用。


相关研究成果已于 5 月 15 日公开。团队注意在此后谷歌发布了 Gemini Diffusion 语言模型,因而期待强化「发散思维链」可以应用到更多的扩散语言模型上成为标准训练过程的一部分。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第1张


  • 论文标题:Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models
  • arXiv 地址:https://arxiv.org/abs/2505.10446
  • GitHub 地址:https://github.com/maple-research-lab/LLaDOU


背景


近期,大型语言模型的推理能力引发了学术界的高度关注。一般而言,「推理」通常是指模型在生成最终应答前所经历的系统性思考过程。当前主流的大型语言模型普遍采用分步拆解问题的方法,构建一种具有因果顺序的线性思维链条,形成所谓的「思维链」推理范式。


值得注意的是,人类认知过程中的思维构建机制与此存在本质差异。在构思阶段,人类思维往往呈现非线性的发散特征,能够突破既有语言框架的约束,通过非线性、跳跃性的方式自发生成概念原型、词汇单元及初始设想。随着认知加工的持续深化,这些离散的思维片段经历系统性整合与结构化重组,最终形成逻辑连贯的完整表达体系。认知科学领域将此类思维模式定义为「发散思维」(Lateral Thinking),显著区别于传统思维链所采用的线性推理模式。


为模拟这一思考过程,西湖大学 MAPLE 实验室齐国君教授团队首次提出扩散式 「发散思维链」这一概念。如图所示,模型的思考过程从一段不包含任何信息的掩码序列开始,在思考过程中,模型会逐步生成推理所需要的关键信息,将掩码转换为具有实际语义内涵的文字内容,如数字和计算过程。最终,在整个扩散去噪流程结束后,模型将生成具有连贯语义内涵且包含正确答案的文字回复。通过仅基于结果的强化学习训练,团队鼓励模型探索多样化的、创造性的和非线性的思维路径,最终得出正确的答案。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第2张


扩散式「发散思维链」


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第3张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第4张


在这一过程中,为了得到最终的输出 x_N,模型天然需要生成一系列中间结果 x_1:N-1。这一过程与「思维链」(Chain-of-Thought, CoT)技术相似。然而,与 CoT 采用线性因果推理不同,扩散过程中的模型能够在思考过程中自由地生成任何有助于达到正确答案的中间内容,更符合发散思维的概念 —— 即通过间接、具有探索性的方法解决问题。正因如此,团队将由去噪过程中所有的中间结果组成的序列称为扩散式「发散思维链」(Diffusion Chain of Lateral Thoughts, DCoLT),并通过强化学习算法优化模型的这些中间扩散「推理」过程


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第5张


在下图中,团队以 GRPO 为例详细阐述了算法训练框架。类似地,其他强化学习算法也可应用于所提出的框架中。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第6张


连续时间扩散语言模型:DCoLT 强化的 SEDD


首先团队考虑以 SEDD 为代表的连续时间扩散语言模型。这类模型通过如下线性常微分方程描述该演化过程。



与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第7张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第8张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第9张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第10张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第11张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第12张


扩展到整个序列时,其转移概率可以看作所有 token 转移概率的累乘,即可通过以下公式计算 DCoLT 生成过程中每一步动作对应的采样概率。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第13张


离散时间扩散语言模型:DCoLT 强化的 LLaDA


一些扩散语言模型直接在离散的时间步上执行多步生成过程。对于这些模型,需要为每个离散步骤定义其输出策略分布。在这其中,考虑最为常见的掩码扩散语言模型。


以 LLaDA 模型为例:生成过程从一个完全掩码序列开始,逐步去除掩码直至生成最终文本。在每个生成步骤中,模型接收一个带有掩码的序列作为输入,将其中部分掩码预测为有实际含义的文本内容。在整个生成过程进行时,掩码的数量会逐渐减少,直到模型最终输出完整的生成序列。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第14张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第15张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第16张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第17张


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第18张


从以上推导可以看出,某种意义上,LLaDOU 模型和基于 next token 预测的自回归(auto-regressive) 语言模型并没有本质区别。两者都是在给定了 prompt 和 context 作为前缀后,去预测后续的 token。区别仅在于,自回归模型要求预测的是紧邻的下一个 token;而 LLaDOU 模型允许通过一个 UPM 模块,从所有可能的后续位置,选择一个或多个 token 进行预测。后者相对于前者更加灵活,可以根据当前生成的结果,打破语言自左到右的自然顺序,在中间步骤,跳跃式地选择合适的 token 进行生成。当然,最终生成的完整结果,仍然满足各种语言语法结构的要求。


在同一时期,业界也推出了一些其他面向 diffusion model 的强化训练方法,如 d1 和 MMaDA。这些方法首先采样得到生成结果以及对应的奖励值,然后对生成结果或问题部分再次进行随机掩码处理,以估算每个 token 的生成概率,用于强化训练。这种情况下,实际采样生成的中间过程和计算概率时的再掩码过程并不一致,可能导致所强化的再掩码采样过程并不是模型真正的采样过程。不同于这些方法,团队直接基于采样过程中每一步所选中的 unmask token 计算概率,据此进行强化训练,保持训练和采样过程一致。同时,更重要的是,团队注意到每步如何选择要 unmask 的 token 也是扩散语言模型采样的关键步骤。基于此,本方法将 unmask token 生成的顺序也作为强化学习所优化策略的一部分,进一步提升扩散语言模型采样的性能。


实验结果


团队基于两个具有代表性的扩散语言模型 ——SEDD 和 LLaDA 开展实验进行验证。


首先,团队基于 SEDD 模型,在数独解题和数学推理两个任务上与其他方法展开了公平对比。DCoLT 取得了比 CoT 和 DoT 更好的实验结果。比如在 GSM8K-Aug 数据集上,同样是使用 SEDD 模型,DCoLT 取得了 57.0% 准确率,超越了 DoT,即使后者使用的训练数据中带有逐步骤的详细 CoT 标注。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第19张


而后,团队在 LLaDA 8B 权重的基础上训练 LLaDOU 模型,充分验证了这一思考技术在数学推理和代码生成任务上的能力。结果显示,该技术显著提升了模型对复杂数学逻辑问题的推理准确率,和生成代码的测试通过率。在相关的评测基准上,LLaDOU 超越了其他扩散语言模型,取得了最好的性能。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第20张


在下图中,团队用不同颜色展示了同一回答中不同 token 的先后生成顺序 —— 越浅的颜色代表 token 在更早的步数生成。可以看出,整个推理过程倾向于首先生成关键数字和计算符号,然后填充其他相关的文本内容,逐渐满足语法约束。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第21张


在这里,团队也以视频形式展示了 LLaDOU 在解决数学问题的完整生成过程。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第22张


总结


这篇文章介绍了由西湖大学 MAPLE 实验室提出的一种全新的大模型推理范式,扩散式「发散思维链」。该框架将反向扩散过程中的中间结果看作模型的推理过程,并将模型最终输出结果的正确性作为奖励开展强化学习训练,大幅提升了大模型的推理能力,在数学推理、代码生成等任务上取得了超越其他扩散语言模型的性能。扩散式「发散思维链」这一理论打破了大模型推理过程的固有范式,为复杂推理问题提供了创新性的方法解决方案,值得我们进一步挖掘。


文章来自于“机器之心”,作者“机器之心”。


与Gemini Diffusion共振!首个扩散式「发散思维链」来了  第23张