arXiv:2505.10446cs.CL2025-05NeurIPS被引 52

用强化学习优化扩散模型的思维链,提升推理准确率。

Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language Models

  • 将扩散过程的每一步视为思维动作,用奖励驱动优化整个推理路径。
  • 在GSM8K等数据集上,推理准确率提升最高达19.5%。
  • 适合研究扩散语言模型与复杂推理任务的学者使用。

我们提出扩散思维链(DCoLT),一种针对扩散语言模型的推理框架。DCoLT将反向扩散过程中的每个中间步骤视为潜在的‘思考’动作,并通过基于结果的强化学习(RL)优化整个推理轨迹,以最大化最终答案的正确性。与传统线性因果思维链不同,DCoLT支持双向、非线性的推理,且对中间步骤的语法正确性无严格要求。我们在两种代表性扩散语言模型上实现DCoLT:一是连续时间离散扩散模型SEDD,其通过具体得分构建概率策略,以最大化序列中所有中间步骤的RL奖励;二是离散时间掩码扩散模型LLaDA,发现预测和解掩码的顺序对基于排列的解掩码策略模块(UPM)至关重要,该模块由Plackett-Luce模型定义。在数学与代码生成任务上的实验表明,仅使用公开数据和16块H800 GPU,DCoLT增强的模型优于经SFT或RL训练的其他模型。值得注意的是,经过DCoLT增强的LLaDA在GSM8K、MATH、MBPP和HumanEval上的推理准确率分别提升了+9.8%、+5.7%、+11.4%和+19.5%。

原文摘要 · Abstract (English)

We introduce the Diffusion Chain of Lateral Thought (DCoLT), a reasoning framework for diffusion language models. DCoLT treats each intermediate step in the reverse diffusion process as a latent "thinking" action and optimizes the entire reasoning trajectory to maximize the reward on the correctness of the final answer with outcome-based Reinforcement Learning (RL). Unlike traditional Chain-of-Thought (CoT) methods that follow a causal, linear thinking process, DCoLT allows bidirectional, non-linear reasoning with no strict rule on grammatical correctness amid its intermediate steps of thought. We implement DCoLT on two representative Diffusion Language Models (DLMs). First, we choose SEDD as a representative continuous-time discrete diffusion model, where its concrete score derives a probabilistic policy to maximize the RL reward over the entire sequence of intermediate diffusion steps. We further consider the discrete-time masked diffusion language model -- LLaDA, and find that the order to predict and unmask tokens plays an essential role to optimize its RL action resulting from the ranking-based Unmasking Policy Module (UPM) defined by the Plackett-Luce model. Experiments on both math and code generation tasks show that using only public data and 16 H800 GPUs, DCoLT-reinforced DLMs outperform other DLMs trained by SFT or RL or even both. Notably, DCoLT-reinforced LLaDA boosts its reasoning accuracy by +9.8%, +5.7%, +11.4%, +19.5% on GSM8K, MATH, MBPP, and HumanEval.

扩散模型推理增强强化学习语言模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。