扩散LLM推理新突破:UCLA & Meta开源d1框架,类GRPO强化学习优于单独SFT

扩散LLM推理迎来新框架d1!UCLA和Meta提出类GRPO强化学习方法diffu-GRPO,结合监督微调,显著提升模型在数学和逻辑推理任务中的性能。

原文标题:扩散LLM推理用上类GRPO强化学习!优于单独SFT,UCLA、Meta新框架d1开源

原文作者:机器之心

冷月清谈:

UCLA和Meta AI的研究者提出了一个名为d1的两阶段后训练框架,旨在提升掩码扩散大语言模型(dLLM)的推理能力。该框架首先通过监督微调(SFT)在高质量推理轨迹上训练模型,然后引入了一种新颖的策略梯度方法diffu-GRPO,用于dLLM的强化学习。diffu-GRPO利用高效的一步对数概率估计,扩展了GRPO算法。研究者通过随机提示词掩码进行策略优化,从而减少了在线生成的需求并降低了计算时间。实验结果表明,d1框架在数学和逻辑推理基准测试中,始终优于基础模型以及单独使用SFT或diffu-GRPO的模型,展现出更强的整体性能提升。定性分析还观察到d1-LLaDA模型在长序列推理中展现出自修正机制和回溯行为。

怜星夜思:

1、d1框架中提到的“顿悟时刻”具体体现在哪些方面?除了自我修正和回溯行为,还可能存在哪些其他的表现形式?
2、文章中diffu-GRPO核心在于高效的对数概率估计。那么,这种估计方法在其他非自回归模型中是否通用?如果不行,改进的方向是什么?
3、既然d1框架结合了SFT和diffu-GRPO,那么在实际应用中,这两个阶段的训练数据应该如何选择和配比?有没有什么trade-off?

原文内容

机器之心报道

编辑:陈陈、杜伟

大语言模型的推理能力,不再是 AR(自回归)的专属。扩散模型现在也能「动脑子」,新框架 d1 让它们学会了解数学、懂逻辑、会思考。


当前,强化学习(RL)方法在最近模型的推理任务上取得了显著的改进,比如 DeepSeek-R1、Kimi K1.5,显示了将 RL 直接用于基础模型可以取得媲美 OpenAI o1 的性能。


不过,基于 RL 的后训练进展主要受限于自回归的大语言模型(LLM),它们通过从左到右的序列推理来运行。


与此同时,离散扩散大语言模型(dLLM)成为有潜力的语言建模的非自回归替代。不像以因果方式逐 token 生成文本的自回归模型那样,dLLM 通过迭代去噪过程生成文本,在多步骤操作中优化序列的同时并通过双向注意力利用过去和未来的上下文。其中,LLaDA 等开放的掩码 dLLM 实现了媲美同尺寸自回归模型的性能,而 Mercury 等闭源 dLLM 进一步展现了出色的推理延迟。


然而,顶级的开源 dLLM 并没有使用 RL 后训练,使得这一有潜力的研究方向还有很大的挖掘空间。这一范式转变引出了重要的问题:RL 后训练如何在非自回归上下文中高效地实现?


RL 算法适应掩码 dLLM 面临一些独特的挑战,原因在于自回归模型采用的已有方法(如 PPO、GRPO)通过计算生成序列的对数概率来估计和优化策略分布,导致无法直接应用于 dLLM。虽然这种计算在自回归模型中通过序列因式分解很容易实现,但 dLLM 由于它们的迭代、非序列生成过程而缺乏这种自然分解。


为了解决这些问题,来自 UCLA 和 Meta AI 的研究者提出了一个两阶段后训练框架 d1,从而可以在掩码 dLLM 中进行推理。在第一阶段,模型在高质量推理轨迹中进行监督微调;在第二即 RL 阶段,研究者引入了用于掩码 dLLM 的新颖策略梯度方法 diffu-GRPO,它利用提出的高效一步(one-step)对数概率估计在 GRPO 的基础上创建。


研究者表示,他们的估计器利用了随机提示词掩码,作为策略优化的一种正则化,使得可以扩展 per batch 的梯度更新数量并减少 RL 训练所需的在线生成数量。这将极大地降低计算时间。



  • 论文标题:d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning

  • 论文地址:https://arxiv.org/pdf/2504.12216

  • 项目主页:https://dllm-reasoning.github.io/

  • GitHub 地址:https://github.com/dllm-reasoning/d1


在实验部分,研究者使用 LLaDA-8B-Instruct 作为基础模型实例化 d1。他们将 d1-LLaDA 的性能与基础 LLaDA 模型以及仅使用 SFT 和仅使用 diffu-GRPO 训练的 LLaDA 模型进行比较。结果表明,d1 在四个数学和逻辑推理基准测试中始终优于基础模型,如下图 1 所示。d1-LLaDA 同样优于仅使用 SFT 方法和仅使用 diffu-GRPO 方法的模型。



方法概览


d1 是一个两阶段框架,通过依次结合监督微调(SFT)和在线强化学习(RL)来增强预训练掩码 dLLMs 的推理性能。


其中,在线强化学习(特别是 GRPO 算法)已被证明能有效提升离线训练语言模型的性能。然而,GRPO 的学习策略并不能直接泛化到 dLLMs。


GRPO 的目标函数(如公式 3 所示)需要同时计算当前策略 π_θ 和旧策略 π_θold 在以下两个层面的(对数)似然比:


  1. token 层面(用于优势权重计算);

  2. 序列层面(用于反向 KL 散度项)。


核心问题在于:研究者需要高效计算 dLLMs 生成内容的逐 token 对数概率和序列对数概率。


自回归(AR)模型,如 Transformer,直接对每个 token 的对数概率进行建模,并且可以通过链式法则使用一次前向传递轻松计算出序列级别的对数概率image.png。同样,KL 项可以分解为image.png


AR 模型不同,dLLMs 不遵循序列对数概率的顺序分解。同时,每个 token 的对数概率计算成本也很高,因为解码过程中需要多次调用掩码预测器 f_θ。基于此,该研究提出了一个高效的对数概率估计器。


对于序列对数概率,该研究使用均场近似方法,将其分解为独立的每个 token 对数概率的乘积。


对于每个 token 的对数概率,该研究引入了一种估计方法,该方法仅调用一次 f_θ。


基于新引入的对数概率估计器,该研究将 GRPO 扩展到掩码 dLLMs,推导出 diffu-GRPO 的损失函数。



算法如下图所示。



实验结果


表 1 报告了基线模型 LLaDA-8B-Instruct 与采用不同后训练优化方案的模型,在四项任务上的零样本性能对比。



图 3 绘制了有效 token 的平均数量:



基于实验,该研究得出以下主要发现:


diffu-GRPO 在所有 12 种设置中都一致优于基础的 LLaDA 和 SFT(监督式微调)。diffu-GRPO 和 SFT 都相较于 LLaDA-8B-Instruct 基线有所提升,但 diffu-GRPO 显示出更持续且幅度更大的增益。具体来说,diffu-GRPO 在所有 12 种设置中都优于 LLaDA-8B-Instruct 和 SFT,而 SFT 仅在其中的 7 种设置中优于 LLaDA-8B-Instruct,这表明 diffu-GRPO 相比于单独的 SFT 实现了更强的整体性能提升


LLaDA+diffu-GRPO 在所有设置中都优于基础的 LLaDA-8B-Instruct 模型,而 d1-LLaDA 在每种情况下都超过了 LLaDA+SFT。这表明,无论初始化是来自预训练模型还是经过 SFT 调整的检查点,diffu-GRPO 都能提供可靠的性能提升


d1 训练方案实现了最显著的性能提升。通过先进行监督微调(SFT)、再结合 diffu-GRPO 训练所形成的 d1-LLaDA 模型,产生了超越单一方法的叠加增益。这种组合式方法在 12 个实验设置中有 11 项优于纯 diffu-GRPO 方案,表明两个训练阶段存在协同效应。


定性结果表明,在 SFT 和 d1-LLaDA 生成中出现了顿悟时刻。尽管与 LLaDA-8B-Instruct 相比,生成序列长度为 128 和 256 的性能随着 SFT、diffu-GRPO 和 d1 有所提高,但从质的方面看,在生成的推理轨迹中并未观察到显著差异。然而当序列长度达到 512 时,该研究开始观察到 SFT 和 d1-LLaDA 模型展现出两种关键能力:自我修正机制和回溯行为。


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com

楼上说的有道理!从计算的角度,“顿悟”可能更多的是一种涌现现象。不过,从用户体验的角度,如果模型能把思考过程可视化出来,让用户也能感受到这个“顿悟”的过程,那就更棒了!

我认为可以考虑将这个思路推广到其他具有类似迭代优化过程的模型中。例如,在图像生成领域,也有通过迭代refinement来提高图像质量的模型,或许可以借鉴这种对数概率估计的策略。

数据配比肯定是个关键问题。SFT阶段需要高质量的推理数据,保证模型“见多识广”;diffu-GRPO阶段则需要reward signal来引导模型学习更有效的推理策略。如果SFT数据不足,模型可能学不到正确的推理模式;如果reward设计不好,可能导致模型“跑偏”。

好问题!文章提到了自我修正和回溯,我觉得“顿悟时刻”可能还体现在模型生成答案过程中突然找到了关键解题思路,或者直接输出了正确答案。就像人一样,遇到难题卡壳了,突然灵光一闪。如果能进一步研究模型内部的激活状态,也许能找到更明确的证据。

我觉着应该先用大量的SFT数据进行预训练,然后再用少量的、精心设计的reward数据进行强化学习。这样可以避免模型过早陷入局部最优解,提高模型的泛化能力。

抛砖引玉一下,大家有没有考虑过将这种思想应用到多模态领域?比如,视频生成、语音合成等。这些领域也存在类似的非自回归生成问题,或许可以从中找到一些灵感。

分享个小技巧,可以尝试Curriculum Learning的思想,先用简单的数据训练SFT,然后逐步增加难度,最后再进行强化学习。这样可以帮助模型更好地学习推理策略,避免“一口吃个胖子”。

我持保守态度。与其说是“顿悟”,不如说是强化学习训练后的结果。自我修正和回溯,本质上是模型在loss的驱使下,不断尝试,最终找到更优解的过程。也许未来我们可以用更复杂的实验来区分这两种情况。

这个问题问到了关键!diffu-GRPO的核心是针对dLLM的特性设计的,尤其是其迭代去噪的过程。直接套用到其他非自回归模型,比如并行解码的Transformer,可能效果不佳。改进方向可能是针对不同模型的生成特点,设计专门的对数概率估计方法。