DiffusionOPD——扩散模型中在线策略蒸馏的统一视角

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

一、论文标题与基本信息

****论文标题:****DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models

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

****机构:****复旦大学(Fudan University)与阿里巴巴通义万相团队(Wan Team, Alibaba Group)

****发表时间:****2026年5月14日

****arXiv编号:****2605.15055

****领域:****机器学习(cs.LG)、计算机视觉与模式识别(cs.CV

****项目主页:****https://quanhaol.github.io/DiffusionOPD-site/

****代码链接:****https://github.com/ali-vilab/DiffusionOPD

二、研究背景

2.1 强化学习在扩散模型中的应用现状

近年来,强化学习(Reinforcement Learning, RL)已成为提升基于扩散模型的文生图系统性能的重要范式。自DDPO(Denoising Diffusion Policy Optimization)将经典RL算法引入扩散模型训练以来,一系列工作如DiffusionNFT、Flow-GRPO、GRPO-Guard等相继涌现,证明了RL能够在特定奖励信号下显著提升模型表现。这些奖励信号涵盖了美学质量、文字渲染准确性、构图对齐度等多个维度,使得扩散模型在单一任务上能够获得针对性的性能提升。例如,通过美学奖励优化,模型可以生成视觉上更具吸引力的图像;通过OCR奖励优化,模型在文字渲染方面的准确性可以得到大幅改善;通过GenEval等组合性评估奖励,模型对复杂文本指令的理解和执行能力也能得到增强。这些单任务RL方法的成功为扩散模型的后期对齐提供了有效的技术路径。

2.2 单任务优化的局限性

然而,单任务RL优化的局限性在实际应用中日益凸显。在真实场景中,用户往往期望一个统一的文生图模型能够同时满足多种需求------既要画面美观,又要忠实于文本指令,还要保证文字渲染的正确性。单任务优化只能针对某一特定维度进行改进,训练出的模型在其他维度上可能表现平平甚至有所退化。例如,一个经过美学奖励充分优化的模型,可能在文字渲染的准确性上不如预期;而一个专注于OCR能力的模型,其生成图像的构图和美感可能有所欠缺。这种"顾此失彼"的现象使得单任务RL难以满足实际部署中对模型综合能力的需求。因此,如何将多种任务能力高效地整合到同一个扩散模型中,成为了一个亞待解决的核心问题。

2.3 多任务RL面临的挑战

将RL扩展到多任务设定面临着根本性的困难。现有的多任务RL方法主要遵循两种范式:联合优化(Joint Optimization)和级联RL(Cascade RL),两者各有其固有的缺陷。

联合优化方法试图在统一的框架中同时训练多个任务。具体而言,它将多个任务特定的奖励信号合并为一个加权组合奖励,然后使用现有的RL算法(如DiffusionNFT、GRPO等)进行联合优化。这种范式面临两个核心挑战。第一是跨任务干扰(Cross-task Interference):不同任务的优化方向往往不一致甚至相互矛盾。例如,提升美学质量可能倾向于生成更加平滑、风格化的图像,而提升OCR准确性则需要模型在文字区域保持更高的清晰度和精确度,这两种优化方向在像素层面可能产生冲突。第二是任务难度不平衡(Task-difficulty Imbalance):不同任务的学习难度差异显著,简单的任务往往会在训练过程中主导梯度更新,使得困难任务的信号被淹没,导致模型在简单任务上过拟合而在困难任务上欠拟合。

级联RL方法则采用顺序训练的策略,即按照一定顺序依次在不同任务上训练模型。虽然这种方式在每个训练阶段内避免了直接的梯度冲突,但在实践中存在诸多问题。首先,训练流程极为繁琐,需要设计多个训练阶段、精心安排训练调度,并为每个阶段调整独立的超参数。其次,级联RL容易遭受灾难性遗忘(Catastrophic Forgetting)的困扰:当模型在后续任务上进行训练时,往往会丧失之前已学到的能力。例如,先训练美学能力再训练OCR能力后,模型的美学表现可能明显退化。这种能力退化使得级联RL难以真正实现多任务能力的有效整合。

2.4 OPD在大语言模型领域的成功

在线策略蒸馏(On-Policy Distillation, OPD)最初在大语言模型(LLM)领域取得了显著成功。OPD的核心思想是:先让多个任务特定的教师模型各自独立地进行充分训练,然后通过在线策略蒸馏的方式将这些教师的能力整合到一个统一的学生模型中。在LLM设定下,学生模型按照自身的策略自回归地生成token序列,随后教师模型在学生访问到的每一个token状态上提供分布级别的监督信号。由于LLM的token分布是离散的,每一步的条件KL散度可以直接以闭式计算,因此OPD目标可以通过直接反向传播来优化,避免了传统RL中高方差的策略梯度估计。这种"先独立探索、后能力整合"的范式有效地解耦了单任务优化与多任务融合,为多任务对齐提供了一个优雅且高效的解决方案。OPD在LLM领域的成功自然引发了如下问题:能否将这一范式推广到扩散模型领域?这正是本文所致力于回答的核心问题。

三、核心方法详解

3.1 DiffusionOPD的整体框架

DiffusionOPD提出了一种基于在线策略蒸馏的多任务扩散模型训练新范式。其整体框架分为两个清晰的阶段。

第一阶段(Stage 1)是单任务教师独立训练。针对不同的优化目标(如GenEval、OCR、Aesthetic等),分别使用现有的扩散模型RL方法训练对应的专家教师模型。每个教师模型仅负责单一任务,因此可以充分专注于该任务的优化,完全避免了跨任务干扰。在具体实现中,GenEval任务采用DiffusionNFT算法,OCR与美学任务采用GRPO-Guard算法。由于每个教师独立训练,可以针对各自任务选择最合适的RL算法和超参数配置,从而最大化每个教师的单任务性能。

第二阶段(Stage 2)是在线策略蒸馏多任务能力到学生模型。从一个预训练扩散模型初始化统一的学生模型,然后通过在线策略蒸馏整合所有教师的能力。在训练过程中,学生模型针对不同任务,基于自身的策略生成去噪轨迹(rollout trajectories)。在学生生成的每个去噪状态上,由对应任务的教师模型提供监督信号。关键在于,学生模型无需重新对所有任务进行从零探索,而是可以直接学习各任务教师已经掌握的策略与能力,从而实现高效的多任务能力融合。这种设计将"探索"与"整合"完全解耦:探索的负担由各教师模型独立承担,学生模型只需要做整合工作。

3.2 理论推导:从离散Token到连续状态马尔可夫过程

将OPD从LLM推广到扩散模型面临一个根本性的理论挑战:LLM处理的是离散的token序列,而扩散模型的去噪过程是连续状态空间中的马尔可夫过程。为了弥合这一鸿沟,作者进行了精巧的理论推导。

首先,作者将OPD的原始表述从自回归token转移提升为一般的离散时间马尔可夫链。设去噪轨迹为 x_{t_0}, x_{t_1}, ..., x_{t_N},学生和教师分别定义各自的单步转移核 p_S 和 p_T。OPD目标可以自然地推广为:在学生自身采样的轨迹上,对所有去噪步的学生-教师转移核之间的KL散度求和并取期望。这一推广保留了两项关键结构性质:(i)轨迹由学生模型采样(在线策略属性);(ii)每步KL必须可以闭式计算(避免REINFORCE技巧)。

对于flow-matching模型,作者遵循Flow-GRPO的方法,通过Euler-Maruyama方法对逆时SDE进行离散化。在给定的时间调度 1=t_0>t_1>...>t_N=0 上,每一步去噪转移都对应一个高斯转移核。具体而言,学生模型的单步转移核为 p_S(x_{t_{j+1}}|x_{t_j}) = N(mu_S(x_{t_j}), sigma_bar_j^2 * I_d),其中均值 mu_S 由学生速度 v_j^S 和当前状态 x_{t_j} 共同决定,方差 sigma_bar_j^2 = sigma_{t_j}^2 * (-Delta_t_j) 由SDE扩散系数和时间步长决定。教师模型的转移核 p_T 采用完全相同的公式结构,仅将学生速度替换为冻结的教师速度 v_j^T。

3.3 闭式KL目标

由于学生和教师的转移核都是高斯分布,且它们共享完全相同的协方差矩阵 sigma_bar_j^2 * I_d(因为协方差仅依赖于调度和噪声水平,与模型参数无关),两个高斯分布之间的逆KL散度可以写成完全解析的闭式表达式。经过推导,per-step KL目标简化为:

L_OPD(theta) = E_{x_{0:N} ~ p_S} sum_{j=0}\^{N-1} \|\|mu_S(x_{t_j}) - mu_T(x_{t_j})\|\|\^2 / (2 \* sigma_bar_j\^2)

这是一个纯粹的均值匹配(mean-matching)损失,完全没有Monte Carlo方差。其物理含义非常直观:在学生模型自身生成的去噪轨迹上的每一个状态,让学生模型的转移均值向教师模型的转移均值靠近。由于协方差相同,KL散度仅取决于两个均值之间的差异。这一闭式目标可以直接通过标准的反向传播进行优化,无需任何策略梯度的方差缩减技巧。

3.4 闭式KL与PPO风格策略梯度的讨论

一个自然的替代方案是:将每步的逆KL散度视为一个密集奖励,把教师模型当作"过程级奖励模型",然后使用PPO风格的策略梯度来优化学生。然而,作者在理论上严格证明了这种做法虽然梯度期望与闭式KL的梯度完全相同,但PPO的梯度中会多出一项与高斯噪声成正比的score-function项。这一额外项的期望为零(因此是无偏的),但其方差不为零,使得PPO估计天然比闭式KL更加"嘘杂"。

更关键的是,PPO形式依赖于log概率和重要性比率的计算,这在ODE确定性采样器下根本无法定义,因此PPO仅支持SDE采样器。而闭式KL目标则天然地同时适用于随机SDE和确定性ODE采样器,提供了更好的通用性。实验也验证了这一理论分析:在相同噪声水平下,闭式KL目标始终优于PPO风格策略梯度。

3.5 统一SDE和ODE采样器

DiffusionOPD的闭式KL目标提供了一个统一SDE和ODE采样器的优雅框架。在SDE情况下(噪声水平 a > 0),目标表现为均值匹配(mean-matching),学生需要在有噪声的去噪状态上向教师的转移均值靠近。当噪声水平设为0时,SDE采样器退化为ODE采样器,此时闭式KL目标进一步退化为转移匹配(transition matching),即学生和教师之间的确定性转移应当完全一致。这一统一性意味着DiffusionOPD可以无缝地在不同类型的采样器下工作,无需修改目标函数的形式。实验表明,更低的噪声水平(即更接近ODE)能够带来更快的收敛速度和更高的性能上限。

3.6 训练配方

在具体实现上,DiffusionOPD的训练配方包含以下关键要素。首先,在蒸馏阶段,学生模型从预训练扩散模型初始化,各教师模型保持冻结。训练时,学生模型针对不同任务分别生成去噪轨迹,在每个去噪状态上收集对应教师的速度预测作为监督信号,然后通过闭式KL目标(即均值匹配损失)进行优化。多个任务的损失通过加权求和的方式进行组合。在采样器选择上,作者发现使用较低噪声水平的SDE采样器(接近ODE)能够取得最佳效果。此外,作者还指出,DiffusionOPD框架不仅限于闭式逆KL目标------一旦学生生成了在线策略轨迹,教师可以使用更广泛的蒸馏目标(如DMD、TDM等)来监督学生访问到的去噪状态。因此,DiffusionOPD应当被理解为一个统一的在线策略蒸馏框架,而非仅仅是一个逆KL方法。

四、实验结果分析

4.1 多任务RL基准对比

DiffusionOPD在多个权威基准上进行了全面评估,包括GenEval(组合性生成评估)、OCR(文字渲染准确性)和美学评分(Aesthetic Score)。实验结果表明,DiffusionOPD在所有评估的基准上均取得了当前最优(State-of-the-Art, SOTA)的结果。

与联合多任务RL基线(如Multi-Task GRPO-Guard、Multi-Task NFT)和级联RL基线(如Cascade NFT)相比,DiffusionOPD展现出两个显著优势。第一是收敛速度的大幅提升:从训练曲线可以清晰看到,DiffusionOPD在训练早期就快速超越所有基线方法,且性能差距随训练推进持续扩大。第二是更高的性能上限:即使在充分训练后,DiffusionOPD的最终性能也显著优于所有基线,表明其不仅能更快地学习,还能达到更高的能力水平。这一结果有力地验证了"解耦探索与整合"范式的优越性。

定性对比结果同样令人印象深刻。在可视化对比中,DiffusionOPD生成的图像在美学质量、文字渲染准确性和文本指令忠实度方面均表现优异,而基线方法往往在某些维度上存在明显不足。例如,Multi-Task GRPO-Guard在文字渲染方面表现较弱,Multi-Task NFT在美学质量上有所欠缺,而Cascade NFT则可能遗忘早期任务的能力。DiffusionOPD成功地将各教师的专业能力整合到了统一的学生模型中,实现了真正的多任务综合优化。

4.2 蒸馏方法消融实验

为了验证DiffusionOPD作为蒸馏框架的有效性,作者进行了一组控制变量实验:固定同一批专家教师,分别使用DiffusionOPD、DMD、TDM和SFT四种蒸馏方法将教师能力蒸馏到同一个学生模型。实验结果表明,在相同的教师和采样设置下,DiffusionOPD在收敛速度和最终性能上限上都明显优于其他蒸馏方法。这表明,DiffusionOPD的在线策略蒸馏范式(即学生基于自身轨迹接收教师监督)比传统的离线蒸馏方法更适合多任务能力整合的场景。SFT方法表现最差,说明简单的监督微调不足以有效整合多个教师的能力。

4.3 损失形式与采样器消融

作者还对损失函数形式和采样器噪声水平进行了详细的消融实验。实验结果清晰地表明:(1)闭式KL目标始终优于PPO风格策略梯度,验证了理论分析中关于梯度方差的结论;(2)更低的噪声水平能够带来更快的收敛速度和更高的性能上限,当噪声水平设为0(即ODE采样器)时,性能达到最佳。这一发现具有重要的实践指导意义:在蒸馏阶段使用确定性ODE采样器不仅可以获得更好的性能,还能消除采样随机性带来的额外方差。

五、与相关工作的对比

5.1 与单任务RL方法的对比

单任务RL方法(如DDPO、Flow-GRPO、GRPO-Guard、DiffusionNFT等)是DiffusionOPD的基础构建模块。这些方法在各自的目标任务上能够取得显著的性能提升,但它们只能优化单一奖励信号,无法同时满足多种需求。DiffusionOPD并不与这些方法竞争,而是将它们作为教师模型的训练算法。事实上,DiffusionOPD的第一阶段正是利用这些单任务RL方法来训练高质量的教师模型,第二阶段再通过在线策略蒸馏将多个教师的能力整合到统一的学生中。因此,DiffusionOPD可以视为单任务RL方法的多任务扩展框架,任何单任务RL方法的改进都可以自然地被DiffusionOPD所吸收和利用。

5.2 与多任务RL方法的对比

DiffusionOPD与现有的多任务RL方法(联合优化和级联RL)形成了鲜明对比。联合优化方法试图在单一模型上同时优化多个奖励,但面临跨任务干扰和任务难度不平衡的挑战。DiffusionOPD通过将单任务训练完全解耦,从根本上消除了跨任务干扰------每个教师只关注自己的任务,不存在梯度冲突的问题。级联RL方法虽然避免了同时优化多个任务,但其顺序训练的方式导致训练流程复杂且容易发生灾难性遗忘。DiffusionOPD通过蒸馏的方式整合能力,学生模型同时从所有教师接收监督信号,不存在"先学后忘"的问题。实验结果充分验证了DiffusionOPD相对于这两类基线方法的全面优势。

5.3 与扩散蒸馏方法的区别

传统的扩散蒸馏方法(如DMD、TDM、SDSR等)主要关注"步数蒸馏"(Step Distillation),即将多步教师模型压缩为少步学生模型以提升推理效率。这些方法的目标是减少采样步数,而非整合多种任务能力。DiffusionOPD则关注一个完全不同的问题维度:如何将多个任务特定的教师模型的能力整合到一个统一的学生模型中。虽然DiffusionOPD在蒸馏阶段也可以使用DMD、TDM等现有的蒸馏目标作为监督信号,但其核心贡献在于提出了在线策略蒸馏的范式------学生基于自身的去噪轨迹接收教师的监督,而非在离线数据上进行蒸馏。这种在线策略属性使得教师只需要在学生实际访问的状态上提供监督,大大提高了蒸馏的效率和效果。

六、优缺点分析

6.1 优点

第一,解耦单任务探索与多任务整合。DiffusionOPD将多任务RL分解为两个独立的阶段,每个教师可以专注于自己的任务而不受其他任务的干扰,学生模型则专注于能力整合而不需要从零开始探索。这种解耦设计从根本上避免了联合优化中的跨任务干扰和级联RL中的灾难性遗忘问题。

第二,闭式目标降低梯度方差。通过利用高斯转移核协方差相同的性质,DiffusionOPD推导出了完全解析的闭式KL目标,可以直接通过反向传播优化,无需依赖高方差的策略梯度估计。理论分析和实验结果均表明,闭式KL目标优于PPO风格的策略梯度替代方案。

第三,统一框架的通用性。DiffusionOPD提供了一个统一的在线策略蒸馏框架,同时适用于SDE和ODE采样器,且不局限于特定的蒸馏目标函数。任何现有的蒸馏方法(如DMD、TDM等)都可以在DiffusionOPD的框架下使用,体现了良好的扩展性和兼容性。

第四,显著的训练效率和性能优势。实验表明,DiffusionOPD在收敛速度和最终性能上全面超越现有的多任务RL基线,在多个基准上取得了SOTA结果。

6.2 缺点

第一,需要训练多个教师模型的计算开销。DiffusionOPD的第一阶段需要为每个任务分别训练一个教师模型,当任务数量较多时,这一阶段的总体计算开销不容忽视。虽然教师训练可以并行进行,但每个教师本身的RL训练仍然需要大量的计算资源。

第二,对教师质量的依赖。学生模型的性能上限受限于教师模型的质量。如果某个教师模型训练不充分或存在偏差,这些缺陷可能会通过蒸馏传递到学生模型中。因此,如何保证每个教师模型都达到足够高的质量,是DiffusionOPD实际应用中需要关注的问题。

第三,蒸馏阶段的任务平衡。虽然DiffusionOPD避免了联合优化中的梯度冲突,但在蒸馏阶段仍需要对不同任务的损失进行加权组合。如何选择合适的权重以平衡不同任务的学习进度,仍是一个需要实践调优的问题。

七、个人见解和未来展望

7.1 对RL与扩散模型结合方向的思考

DiffusionOPD代表了RL与扩散模型结合方向上的一个重要进展。过去两年,这一方向的研究主要集中在如何将经典RL算法适配到扩散模型的连续生成过程中,以及如何设计有效的奖励函数来引导模型优化。DiffusionOPD从一个更高的视角审视了这一问题:它关注的不是如何在扩散模型上做RL,而是如何高效地将多个RL优化的结果整合到一个统一模型中。这种"元优化"的视角可能为未来的研究开辟新的方向。例如,是否可以设计自动化的教师选择和质量评估机制,使得整个流程更加自动化?是否可以在蒸馏阶段引入自适应的任务权重调整策略?这些问题都值得进一步探索。

7.2 对OPD范式扩展到更多模态的展望

DiffusionOPD的核心贡献在于将OPD从离散token的LLM设定推广到了连续状态空间的扩散模型设定。这一理论推广的思路具有广泛的适用性。在视频生成领域,多个教师模型可以分别优化视频的时序一致性、画面质量和文本对齐度等不同维度,然后通过DiffusionOPD整合到一个统一的多能力视频生成模型中。在3D生成领域,类似的多任务整合需求同样存在。在音频生成领域,不同教师可以分别优化音质、音色和内容准确性。可以预见,OPD范式有望成为多模态生成模型多任务对齐的通用框架。

此外,DiffusionOPD的理论框架将OPD从自回归过程推广到了一般的连续状态马尔可夫过程,这一推广本身也具有独立的理论价值。它为理解扩散模型中的策略优化和知识蒸馏提供了新的数学工具和视角,可能启发更多关于连续空间中策略学习和蒸馏的理论研究。

7.3 对工业部署的意义

从工业部署的角度来看,DiffusionOPD具有显著的实际价值。在当前的文生图模型开发流程中,不同团队往往分别负责不同能力的优化(如美学团队、OCR团队、安全团队等),最终需要将这些能力整合到一个统一的模型中。DiffusionOPD提供的"先独立训练、后蒸馏整合"的范式与这种团队分工模式高度契合。各团队可以独立训练自己的教师模型,然后通过DiffusionOPD进行高效整合,大大简化了多能力模型的生产流程。

此外,DiffusionOPD的模块化设计使得新增任务变得非常便捷:只需训练一个新的教师模型,然后将其纳入蒸馏阶段即可,无需重新训练整个系统。这种增量式的扩展能力对于快速迭代的工业环境尤为重要。随着用户对AI生成内容质量要求的不断提升,DiffusionOPD为构建兼具多种优秀能力的统一生成模型提供了一个高效、灵活且理论扎实的解决方案,有望在未来的工业实践中发挥越来越重要的作用。

相关推荐
盼小辉丶5 小时前
PyTorch强化学习实战(12)——Double DQN(DDQN)
人工智能·pytorch·深度学习·强化学习
君为先-bey2 天前
JointDiT:使用扩散变换器增强RGB-深度联合建模
人工智能·深度学习·计算机视觉·扩散模型·图像生成
君为先-bey2 天前
VideoReward: 人类反馈优化视频生成文献深度阅读分析
人工智能·音视频·扩散模型
君为先-bey2 天前
CineMaster: 3D感知电影级视频生成框架文献深度阅读分析
3d·音视频·扩散模型
未知方程 无解3 天前
我用RTX4090显卡训练了一个AI绘画模型
ai绘画·扩散模型·rtx4090·rxt4090显卡
摘星编程3 天前
豆包Seedream 4.0多图融合实力派:田园犬+三花猫多场景创作,AI绘画新时代来了!
ai绘画·扩散模型·图像生成·多图融合·seedream4.0
亲爱的阿瞎3 天前
p12 3.3 学习状态值函数_cdn
强化学习
亲爱的阿瞎3 天前
p09 2.4 random stochastic e_cdn
强化学习
亲爱的阿瞎4 天前
p08 2.3 贝尔曼方程_cdn
强化学习