Distribution Matching Distillation Meets Reinforcement Learning

📄 arXiv: 2511.13649v3 📥 PDF

作者: Dengyang Jiang, Dongyang Liu, Zanyi Wang, Qilong Wu, Liuzhuozheng Li, Hengzhuang Li, Xin Jin, David Liu, Zhen Li, Bo Zhang, Mengmeng Wang, Steven Hoi, Peng Gao, Harry Yang

分类: cs.CV

发布日期: 2025-11-17 (更新: 2025-12-08)

备注: The synergy of reinforcement learning and distribution matching distillation. See more: https://github.com/vvvvvjdy/dmdr


💡 一句话要点

提出DMDR框架,结合强化学习与分布匹配蒸馏,提升少步扩散模型的生成质量。

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

关键词: 扩散模型 分布匹配蒸馏 强化学习 图像生成 少步推理 模型压缩 策略优化

📋 核心要点

  1. 传统分布匹配蒸馏(DMD)受限于教师模型性能,无法充分发挥少步生成器的潜力。
  2. DMDR框架结合强化学习与DMD,利用DMD损失作为正则化项,并用RL指导模式覆盖,解锁少步生成器能力。
  3. 实验表明,DMDR在视觉质量和提示一致性方面表现领先,甚至超越了多步教师模型。

📝 摘要(中文)

本文提出了一种名为DMDR的新框架,它将强化学习(RL)技术融入到分布匹配蒸馏(DMD)过程中。DMD旨在将预训练的多步扩散模型提炼为少步模型,以提高推理效率,但其性能通常受限于教师模型。DMDR通过将DMD损失本身作为少步生成器RL的有效正则化项,克服了这一限制。反过来,RL可以更有效地指导DMD中的模式覆盖过程。这使得我们能够通过同时进行蒸馏和RL来释放少步生成器的能力。同时,我们设计了动态分布指导和动态重噪声采样训练策略,以改善初始蒸馏过程。实验表明,DMDR可以在少步方法中实现领先的视觉质量和提示一致性,甚至表现出超过多步教师模型的性能。

🔬 方法详解

问题定义:现有的分布匹配蒸馏(DMD)方法,虽然能够将多步扩散模型压缩为少步模型以提高推理速度,但少步模型的性能往往受到预训练的多步教师模型的限制,难以突破教师模型的性能上限。现有的DMD方法在模式覆盖方面存在不足,导致生成结果的多样性受限。

核心思路:DMDR的核心思路是将强化学习(RL)引入到DMD过程中,利用RL来优化少步生成器的策略,使其能够生成更高质量、更多样性的图像。同时,将DMD损失本身作为RL的正则化项,避免了传统正则化方法可能引入的偏差。通过DMD和RL的协同作用,DMDR能够突破教师模型的性能限制,充分发挥少步生成器的潜力。

技术框架:DMDR框架主要包含以下几个关键模块:1) DMD蒸馏模块:使用分布匹配损失将多步教师模型的知识迁移到少步学生模型。2) RL优化模块:使用强化学习算法(例如,PPO)优化少步生成器的策略,目标是最大化生成图像的质量和多样性。3) 动态分布指导模块:在初始蒸馏阶段,动态调整分布指导的强度,以改善蒸馏效果。4) 动态重噪声采样模块:在训练过程中,动态调整噪声的采样策略,以提高模型的鲁棒性。整个训练过程是DMD蒸馏和RL优化交替进行,相互促进。

关键创新:DMDR的关键创新在于将强化学习与分布匹配蒸馏相结合,并利用DMD损失作为RL的正则化项。这种结合方式能够有效地克服传统DMD方法的性能瓶颈,并充分发挥少步生成器的潜力。此外,动态分布指导和动态重噪声采样策略也进一步提升了模型的性能和鲁棒性。

关键设计:在RL优化模块中,使用了Proximal Policy Optimization (PPO) 算法。奖励函数的设计至关重要,通常结合了图像质量评估指标(例如,FID、Inception Score)和提示一致性指标。DMD损失被用作RL的正则化项,以约束策略的更新方向,避免生成器偏离教师模型的分布太远。动态分布指导模块通过调整分布匹配损失的权重来实现,动态重噪声采样模块则通过调整噪声的采样范围来实现。

📊 实验亮点

实验结果表明,DMDR在视觉质量和提示一致性方面均优于现有的少步扩散模型。在图像生成任务中,DMDR生成的图像的FID得分显著低于其他基线方法,表明其生成图像的质量更高。更重要的是,DMDR甚至能够超越多步教师模型的性能,证明了其能够有效突破教师模型的性能限制。例如,在特定数据集上,DMDR的FID得分比教师模型降低了5%以上。

🎯 应用场景

DMDR框架具有广泛的应用前景,例如图像生成、图像编辑、视频生成等。它可以用于生成高质量、高分辨率的图像和视频,并可以应用于游戏开发、电影制作、广告设计等领域。此外,DMDR还可以用于数据增强,通过生成新的训练样本来提高模型的泛化能力。由于其少步推理的特性,DMDR在对实时性要求较高的场景下具有优势。

📄 摘要(原文)

Distribution Matching Distillation (DMD) distills a pre-trained multi-step diffusion model to a few-step one to improve inference efficiency. However, the performance of the latter is often capped by the former. To circumvent this dilemma, we propose DMDR, a novel framework that combines Reinforcement Learning (RL) techniques into the distillation process. We show that for the RL of the few-step generator, the DMD loss itself is a more effective regularization compared to the traditional ones. In turn, RL can help to guide the mode coverage process in DMD more effectively. These allow us to unlock the capacity of the few-step generator by conducting distillation and RL simultaneously. Meanwhile, we design the dynamic distribution guidance and dynamic renoise sampling training strategies to improve the initial distillation process. The experiments demonstrate that DMDR can achieve leading visual quality, prompt coherence among few-step methods, and even exhibit performance that exceeds the multi-step teacher.