DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models

📄 arXiv: 2605.15055v1 📥 PDF

作者: Quanhao Li, Junqiu Yu, Kaixun Jiang, Yujie Wei, Zhen Xing, Pandeng Li, Ruihang Chu, Shiwei Zhang, Yu Liu, Zuxuan Wu

分类: cs.LG, cs.CV

发布日期: 2026-05-14

💡 一句话要点

DiffusionOPD：扩散模型中基于在线策略蒸馏的多任务统一框架

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 强化学习 在线策略蒸馏 多任务学习 文本到图像生成

📋 核心要点

  1. 现有扩散模型强化学习方法局限于单任务，多任务联合优化存在任务间干扰和不平衡问题。
  2. DiffusionOPD 采用在线策略蒸馏，先独立训练教师模型，再将知识提炼到统一的学生模型。
  3. 实验表明，DiffusionOPD 在训练效率和性能上优于多奖励和级联强化学习，并在多个基准测试中达到 SOTA。

📝 摘要（中文）

强化学习已成为改进基于扩散的文本到图像模型的强大工具，但现有方法主要局限于单任务优化。将强化学习扩展到多任务具有挑战性：联合优化会受到跨任务干扰和不平衡的影响，而级联强化学习则繁琐且容易发生灾难性遗忘。我们提出了DiffusionOPD，这是一种基于在线策略蒸馏（OPD）的扩散模型多任务训练新范式。DiffusionOPD首先独立训练特定于任务的教师模型，然后沿着学生模型自身的 rollout 轨迹将它们的能力提炼到一个统一的学生模型中。这解耦了单任务探索与多任务集成，并避免了从头开始联合解决所有任务的优化负担。从理论上讲，我们将OPD框架从离散 token 提升到连续状态马尔可夫过程，推导出封闭形式的每步 KL 目标，该目标通过均值匹配统一了随机 SDE 和确定性 ODE 细化。我们正式地并通过实验证明，与传统的 PPO 风格策略梯度相比，这种解析梯度提供了更低的方差和更好的泛化性。大量的实验表明，DiffusionOPD 在训练效率和最终性能方面始终优于多奖励强化学习和级联强化学习基线，同时在所有评估的基准测试中都取得了最先进的结果。

🔬 方法详解

问题定义：现有基于强化学习的扩散模型训练方法，在扩展到多任务时面临挑战。直接联合优化多个任务容易导致任务间的相互干扰和训练不平衡，而级联强化学习方法则流程繁琐，且容易出现灾难性遗忘现象。因此，如何高效且有效地训练能够处理多个任务的扩散模型是一个亟待解决的问题。

核心思路：DiffusionOPD 的核心思路是将多任务学习分解为单任务学习和知识蒸馏两个阶段。首先，针对每个任务独立训练一个教师模型，使其充分掌握该任务的知识。然后，利用在线策略蒸馏（OPD）的方法，将这些教师模型的知识迁移到一个统一的学生模型中。学生模型在自身的 rollout 轨迹上学习，从而避免了直接联合优化多个任务的困难。

技术框架：DiffusionOPD 的整体框架包含以下几个主要阶段： 1. 教师模型训练：针对每个任务，使用强化学习方法（如 PPO）独立训练一个教师模型。 2. 学生模型初始化：初始化一个学生模型，作为多任务学习的载体。 3. 在线策略蒸馏：学生模型在自身的 rollout 轨迹上进行学习，通过最小化与教师模型输出的差异来学习教师模型的知识。具体而言，论文推导了一个封闭形式的每步 KL 散度目标函数，用于指导学生模型的训练。 4. 模型评估与迭代：定期评估学生模型在各个任务上的性能，并根据评估结果调整训练策略，进行迭代优化。

关键创新：DiffusionOPD 的关键创新在于将在线策略蒸馏框架从离散 token 扩展到连续状态的马尔可夫过程，并推导出了一个封闭形式的每步 KL 散度目标函数。该目标函数统一了随机 SDE 和确定性 ODE 细化过程，并能够提供比传统 PPO 风格策略梯度更低的方差和更好的泛化性。此外，DiffusionOPD 通过解耦单任务探索和多任务集成，有效避免了联合优化多个任务的困难。

关键设计： 1. KL 散度目标函数：论文推导了一个封闭形式的每步 KL 散度目标函数，用于衡量学生模型和教师模型输出之间的差异。该目标函数基于均值匹配，能够有效地指导学生模型的训练。 2. Rollout 轨迹生成：学生模型在自身的 rollout 轨迹上进行学习，这意味着学生模型需要能够生成高质量的 rollout 轨迹，以便进行有效的知识蒸馏。 3. 任务权重设置：在多任务学习过程中，需要合理设置每个任务的权重，以平衡不同任务之间的学习进度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DiffusionOPD 在多个文本到图像生成基准测试中取得了最先进的结果。实验表明，DiffusionOPD 在训练效率和最终性能方面始终优于多奖励强化学习和级联强化学习基线。具体性能提升数据未知，但论文强调了其在多个评估基准上的优越性。

🎯 应用场景

DiffusionOPD 有潜力应用于各种需要多任务处理的扩散模型应用场景，例如图像编辑、风格迁移、超分辨率重建等。通过将多个任务的知识集成到一个统一的模型中，可以显著提高模型的通用性和效率。该方法还可以应用于机器人控制、自然语言处理等领域，具有广阔的应用前景。

📄 摘要（原文）

Reinforcement learning has emerged as a powerful tool for improving diffusion-based text-to-image models, but existing methods are largely limited to single-task optimization. Extending RL to multiple tasks is challenging: joint optimization suffers from cross-task interference and imbalance, while cascade RL is cumbersome and prone to catastrophic forgetting. We propose DiffusionOPD, a new multi-task training paradigm for diffusion models based on Online Policy Distillation (OPD). DiffusionOPD first trains task-specific teachers independently, then distills their capabilities into a unified student along the student own rollout trajectories. This decouples single-task exploration from multi-task integration and avoids the optimization burden of solving all tasks jointly from scratch. Theoretically, we lift the OPD framework from discrete tokens to continuous-state Markov processes, deriving a closed-form per-step KL objective that unifies both stochastic SDE and deterministic ODE refinement via mean-matching. We formally and empirically demonstrate that this analytic gradient provides lower variance and better generality compared to conventional PPO-style policy gradients. Extensive experiments show that DiffusionOPD consistently surpasses both multi-reward RL and cascade RL baselines in training efficiency and final performance, while achieving state-of-the-art results on all evaluated benchmarks.