A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models

📄 arXiv: 2607.14522v1 📥 PDF

作者: Zikun Zhang, Jiayuan Sheng, David D. Yao, Wenpin Tang

分类: cs.LG

发布日期: 2026-07-16

备注: 36 pages, 5 figures

💡 一句话要点

提出连续时间强化学习框架以优化离散扩散模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 连续时间 离散扩散模型 策略优化 中间奖励 掩蔽扩散模型 轨迹子采样 自然语言处理

📋 核心要点

  1. 现有的GRPO方法仅依赖终端奖励，无法有效利用中间奖励信号，限制了优化的灵活性。
  2. 提出的连续时间RL框架允许在去噪过程中整合中间奖励，提升了策略优化的效果和灵活性。
  3. 在数学推理和编码任务上，实验结果显示该方法在低维熵正则化优化问题和dLLMs的后训练中表现优异。

📝 摘要（中文）

本文通过随机控制方法将强化学习（RL）形式化为具有离散状态空间和可能任意动作空间的连续时间问题，利用受控连续时间马尔可夫链（CTMC）建模状态动态。我们推导了相应的策略梯度方法，提出了连续时间的近端策略优化（PPO）和群体相对策略优化（GRPO）变体。作为主要应用，我们开发了一个完整的连续时间RL框架，用于微调基于分数的离散扩散模型。该框架支持奖励驱动的优化，无需对奖励信号的可微性要求，允许在去噪轨迹中整合中间奖励或优势信号。针对掩蔽扩散模型（MDMs），我们的框架提供了丰富的策略参数化，具有可解析的概率比，统一了MDMs中的探索与策略优化视角。我们还提出了轨迹子采样技术，以高效估计计算上昂贵的轨迹似然，降低计算成本。

🔬 方法详解

问题定义：本文旨在解决现有GRPO方法对奖励信号的依赖限制，特别是无法有效利用中间奖励的问题。现有方法在优化过程中缺乏灵活性，导致性能下降。

核心思路：通过将强化学习形式化为连续时间问题，利用受控连续时间马尔可夫链（CTMC）建模状态动态，允许在去噪轨迹中整合中间奖励信号，从而提升策略优化的灵活性和效果。

技术框架：整体架构包括状态动态建模、策略优化和奖励信号整合三个主要模块。首先，利用CTMC建模状态变化；其次，推导连续时间的PPO和GRPO方法；最后，结合中间奖励信号进行策略优化。

关键创新：最重要的技术创新在于提出了允许中间奖励信号的连续时间RL框架，与现有方法相比，显著提升了策略优化的灵活性和效果。

关键设计：在技术细节上，设计了适用于掩蔽扩散模型的策略参数化，确保概率比的可解析性，并提出了轨迹子采样技术以降低计算成本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，所提出的方法在低维熵正则化优化问题和dLLMs的后训练中均表现出色，相较于基线方法，性能提升显著，尤其在数学推理和编码任务上取得了优异的效果。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、图像生成等基于扩散模型的任务。通过优化离散扩散模型，该框架能够提升模型在复杂任务中的表现，具有重要的实际价值和未来影响。

📄 摘要（原文）

We formulate reinforcement learning (RL) in continuous time with discrete state spaces and possibly arbitrary action spaces via a stochastic control approach, where the state dynamics are modeled as a controlled continuous-time Markov chain (CTMC). We consider policy optimization problems and derive the corresponding policy gradient methods, leading to continuous-time variants of proximal policy optimization (PPO) and group relative policy optimization (GRPO). As a primary application, we develop a complete continuous-time RL framework for fine-tuning score-based discrete diffusion models. The proposed framework enables reward-driven optimization without requiring differentiability on the reward signals. In contrast to the existing GRPO-based approaches that only rely on terminal rewards, our formulation allows intermediate reward or advantage signals to be incorporated throughout the denoising trajectory. Importantly, when specialized to masked diffusion models (MDMs), our framework encompasses a rich class of policy parameterizations over the vocabulary simplex with analytically tractable probability ratios, providing a unified perspective on exploration and policy optimization in MDMs. For masked diffusion large language models (dLLMs), we further propose trajectory subsampling techniques to efficiently estimate computationally prohibitive trajectory likelihoods, reducing the computational cost of computing per-position probability ratios. We showcase the effectiveness of our methods on both low-dimensional entropy-regularized optimization problems and RL post-training of dLLMs on mathematical reasoning and coding tasks.