【他山之石】扩散LLM推理用上类GRPO强化学习!优于单独SFT,UCLA、Meta新框架d1开源
仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接
“他山之石,可以攻玉”,站在巨人的肩膀才能看得更高,走得更远。在科研的道路上,更需借助东风才能更快前行。为此,我们特别搜集整理了一些实用的代码链接,数据集,软件,编程技巧等,开辟“他山之石”专栏,助你乘风破浪,一路奋勇向前,敬请关注!
大语言模型的推理能力,不再是AR(自回归)的专属。扩散模型现在也能「动脑子」,新框架d1让它们学会了解数学、懂逻辑、会思考。
当前,强化学习(RL)方法在最近模型的推理任务上取得了显著的改进,比如DeepSeek-R1、KimiK1.5,显示了将RL直接用于基础模型可以取得媲美OpenAIo1的性能。
不过,基于RL的后训练进展主要受限于自回归的大语言模型(LLM),它们通过从左到右的序列推理来运行。
与此同时,离散扩散大语言模型(dLLM)成为有潜力的语言建模的非自回归替代。不像以因果方式逐token生成文本的自回归模型那样,dLLM通过迭代去噪过程生成文本,在多步骤操作中优化序列的同时并通过双向注意力利用过去和未来的上下文。其中,LLaDA等开放的掩码dLLM实现了媲美同尺寸自回归模型的性能,而Mercury等闭源dLLM进一步展现了出色的推理延迟。
然而,顶级的开源dLLM并没有使用RL后训练,使得这一有潜力的研究方向还有很大的挖掘空间。这一范式转变引出了重要的问题:RL后训练如何在非自回归上下文中高效地实现?
RL算法适应掩码dLLM面临一些独特的挑战,原因在于自回归模型采用的已有方法(如PPO、GRPO)通过计算生成序列的对数概率来估计和优化策略分布,导致无法直接应用于dLLM。虽然这种计算在自回归模型中通过序列因式分解很容易实现,但dLLM由于它们的迭代、非序列生成过程而缺乏这种自然分解。
为了解决这些问题,来自UCLA和MetaAI的研究者提出了一个两阶段后训练框架d1,从而可以在掩码dLLM中进行推理。在第一阶段,模型在高质量推理轨迹中进行监督微调;在第二即RL阶段,研究者引入了用于掩码dLLM的新颖策略梯度方法diffu-GRPO,它利用提出的高效一步(one-step)对数概率估计在GRPO的基础上创建。
研究者表示,他们的估计器利用了随机提示词掩码,作为策略优化的一种正则化,使得可以扩展perbatch的梯度更新数量并减少RL训练所需的在线生成数量。这将极大地降低计算时间。
论文标题:d1:ScalingReasoninginDiffusionLargeLanguageModelsviaReinforcementLearning
论文地址: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在以下两个层面的(对数)似然比:
token层面(用于优势权重计算);
序列层面(用于反向KL散度项)。
核心问题在于:研究者需要高效计算dLLMs生成内容的逐token对数概率和序列对数概率。
自回归(AR)模型,如Transformer,直接对每个token的对数概率进行建模,并且可以通过链式法则使用一次前向传递轻松计算出序列级别的对数概率
。同样,KL项可以分解为
。
与AR模型不同,dLLMs不遵循序列对数概率的顺序分解。同时,每个token的对数概率计算成本也很高,因为解码过程中需要多次调用掩码预测器f_θ。基于此,该研究提出了一个高效的对数概率估计器。
对于序列对数概率,该研究使用均场近似方法,将其分解为独立的每个token对数概率的乘积。
对于每个token的对数概率,该研究引入了一种估计方法,该方法仅调用一次f_θ。
基于新引入的对数概率估计器,该研究将GRPO扩展到掩码dLLMs,推导出diffu-GRPO的损失函数。
算法如下图所示。
实验结果
表1报告了基线模型LLaDA-8B-Instruct与采用不同后训练优化方案的模型,在四项任务上的零样本性能对比。
图3绘制了有效token的平均数量:
基于实验,该研究得出以下主要发现:
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模型展现出两种关键能力:自我修正机制和回溯行为。
本文目的在于学术交流,并不代表本公众号赞同其观点或对其内容真实性负责,版权归原作者所有,如有侵权请告知删除。
收藏,分享、在看,给个三连击呗!