Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning

📄 arXiv: 2604.19009v1 📥 PDF

作者: Linwei Dong, Ruoyu Guo, Ge Bai, Zehuan Yuan, Yawei Luo, Changqing Zou

分类: cs.LG, cs.CV

发布日期: 2026-04-21

💡 一句话要点

提出GDMD，利用梯度强化学习指导扩散模型蒸馏，实现高质量少步生成。

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

关键词: 扩散模型 蒸馏 强化学习 梯度优化 少步生成

📋 核心要点

  1. 现有扩散模型蒸馏方法在加速生成的同时，通常会牺牲生成质量，这是核心问题。
  2. GDMD通过利用蒸馏梯度作为强化学习的奖励信号，优化蒸馏过程，提升生成质量。
  3. 实验表明，GDMD在少步生成任务上取得了SOTA结果，显著优于现有方法。

📝 摘要（中文）

扩散模型蒸馏，如分布匹配蒸馏（DMD），在少步生成方面展现了巨大潜力，但通常以牺牲质量为代价来换取采样速度。将强化学习（RL）整合到蒸馏中具有潜在价值，但直接融合这两个目标依赖于次优的原始样本评估。这种基于样本的评分与蒸馏轨迹存在内在冲突，并由于早期生成阶段的噪声产生不可靠的奖励。为了克服这些限制，我们提出了GDMD，一种新颖的框架，通过优先考虑蒸馏梯度而非原始像素输出作为优化的主要信号，重新定义了奖励机制。通过将DMD梯度重新解释为隐式目标张量，我们的框架使现有的奖励模型能够直接评估蒸馏更新的质量。这种梯度级别的指导充当自适应权重，将RL策略与蒸馏目标同步，有效地消除了优化发散。实验结果表明，GDMD为少步生成设定了新的SOTA。具体而言，我们的4步模型在质量上优于其多步教师模型，并在GenEval和人类偏好指标上大幅超过了之前的DMDR结果，展现出强大的可扩展性潜力。

🔬 方法详解

问题定义：现有的扩散模型蒸馏方法，例如DMD，虽然能够显著减少生成步骤，但往往会导致生成质量下降。直接将强化学习引入蒸馏过程，由于早期生成样本的质量较差，导致奖励信号不稳定，进而影响蒸馏效果。现有方法难以平衡生成速度和质量，且强化学习的引入容易导致优化发散。

核心思路：GDMD的核心思路是将DMD的梯度信息作为强化学习的奖励信号，而不是直接使用生成的像素样本。通过评估蒸馏更新的质量，可以更准确地指导蒸馏过程，从而在保证生成速度的同时，提升生成质量。这种方法避免了直接评估低质量样本带来的噪声干扰，并能更好地同步强化学习策略与蒸馏目标。

技术框架：GDMD框架主要包含以下几个部分：首先，使用DMD进行初步的蒸馏训练。然后，将DMD的梯度信息提取出来，作为强化学习的奖励信号。现有的奖励模型被用于评估这些梯度，从而指导强化学习策略的更新。强化学习策略的更新反过来影响蒸馏过程，使其朝着更好的方向发展。整个框架通过迭代优化，最终实现高质量的少步生成。

关键创新：GDMD最重要的创新点在于将蒸馏梯度作为强化学习的奖励信号。与传统的基于样本的奖励机制相比，基于梯度的奖励机制能够更准确地反映蒸馏更新的质量，从而避免了噪声干扰，提高了强化学习的效率。这种方法有效地解决了强化学习与蒸馏目标之间的优化发散问题。

关键设计：GDMD的关键设计包括：1) 如何有效地提取和利用DMD的梯度信息；2) 如何设计奖励模型，使其能够准确地评估蒸馏更新的质量；3) 如何将强化学习策略的更新融入到蒸馏过程中。论文中可能涉及对现有奖励模型的修改或定制，以适应梯度信号的特点。具体的损失函数可能包括DMD的蒸馏损失和强化学习的策略梯度损失，并通过自适应权重进行平衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

GDMD在少步生成任务上取得了显著的性能提升。具体而言，使用4步生成的GDMD模型在质量上超过了其多步教师模型，并在GenEval和人类偏好指标上大幅优于之前的DMDR结果。这些实验结果表明，GDMD能够有效地平衡生成速度和质量，并具有强大的可扩展性潜力。

🎯 应用场景

GDMD在图像生成、视频生成等领域具有广泛的应用前景。它可以用于快速生成高质量的图像和视频内容，例如在游戏开发、电影制作、广告设计等领域。此外，GDMD还可以应用于数据增强、图像修复等任务，提高模型的性能和鲁棒性。未来，GDMD有望成为一种通用的扩散模型蒸馏方法，推动生成模型的发展。

📄 摘要（原文）

Diffusion distillation, exemplified by Distribution Matching Distillation (DMD), has shown great promise in few-step generation but often sacrifices quality for sampling speed. While integrating Reinforcement Learning (RL) into distillation offers potential, a naive fusion of these two objectives relies on suboptimal raw sample evaluation. This sample-based scoring creates inherent conflicts with the distillation trajectory and produces unreliable rewards due to the noisy nature of early-stage generation. To overcome these limitations, we propose GDMD, a novel framework that redefines the reward mechanism by prioritizing distillation gradients over raw pixel outputs as the primary signal for optimization. By reinterpreting the DMD gradients as implicit target tensors, our framework enables existing reward models to directly evaluate the quality of distillation updates. This gradient-level guidance functions as an adaptive weighting that synchronizes the RL policy with the distillation objective, effectively neutralizing optimization divergence. Empirical results show that GDMD sets a new SOTA for few-step generation. Specifically, our 4-step models outperform the quality of their multi-step teacher and substantially exceed previous DMDR results in GenEval and human-preference metrics, exhibiting strong scalability potential.