$R_{dm}$: Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation
作者: Linqian Fan, Peiqin Sun, Tiancheng Wen, Shun Lu, Chengru Song
分类: cs.CV, cs.LG
发布日期: 2026-03-30
💡 一句话要点
提出Rdm框架,将分布匹配重构为扩散蒸馏的奖励,提升生成质量与效率。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 扩散模型 扩散蒸馏 强化学习 分布匹配 生成模型
📋 核心要点
- 扩散模型采样速度慢,扩散蒸馏受限于教师模型,难以突破性能瓶颈。
- 将分布匹配视为奖励,统一扩散匹配蒸馏和强化学习,实现更灵活的优化。
- 引入组归一化分布匹配(GNDM)稳定训练,结合重要性采样提升采样效率。
📝 摘要(中文)
扩散模型在生成任务上表现出色,但其迭代采样过程缓慢。扩散蒸馏技术虽然能实现高质量的少步生成,但传统目标函数通常将学生模型限制于教师模型,阻碍了性能提升。本文提出一种新范式,将分布匹配重新概念化为奖励,记为$R_{dm}$。这种统一视角弥合了扩散匹配蒸馏(DMD)和强化学习(RL)之间的算法差距,带来多项优势:(1)增强优化稳定性:引入组归一化分布匹配(GNDM),利用组均值统计来稳定$R_{dm}$估计,建立更鲁棒有效的优化方向。(2)无缝奖励集成:基于奖励的公式自然支持自适应权重机制,灵活结合DMD与外部奖励模型。(3)提高采样效率:通过与RL原则对齐,该框架易于整合重要性采样(IS),显著提升采样效率。实验表明,GNDM优于vanilla DMD,FID降低1.87。多奖励变体GNDMR在美学质量和保真度之间取得平衡,HPS达到30.37,FID-SD降至12.21。
🔬 方法详解
问题定义:扩散模型虽然生成效果好,但采样速度慢,难以满足实时性要求。扩散蒸馏旨在加速采样过程,但现有方法通常直接模仿教师模型的输出,限制了学生模型的潜在性能,难以超越教师模型。此外,将蒸馏和强化学习简单相加,优化不稳定,难以取得理想效果。
核心思路:将分布匹配过程重新理解为一种奖励信号,即$R_{dm}$。学生模型的目标不再是直接模仿教师模型的输出,而是最大化与教师模型输出分布的匹配程度,从而获得奖励。这种思路将扩散匹配蒸馏和强化学习统一起来,允许更灵活的优化策略,并能方便地整合外部奖励信号。
技术框架:整体框架包含一个学生扩散模型,通过最大化奖励函数来训练。奖励函数由两部分组成:分布匹配奖励$R_{dm}$和可选的外部奖励。$R_{dm}$通过比较学生模型和教师模型的输出分布来计算。为了稳定训练,引入了组归一化分布匹配(GNDM)。训练过程中,可以使用重要性采样来提高采样效率。
关键创新:核心创新在于将分布匹配重新概念化为奖励。这种视角转变使得可以利用强化学习的工具和技术来优化扩散蒸馏过程,从而突破传统蒸馏方法的性能限制。GNDM通过利用组均值统计,稳定了$R_{dm}$的估计,解决了训练不稳定的问题。
关键设计:GNDM使用组归一化来稳定分布匹配奖励的估计。具体来说,将一批样本分成若干组,计算每组的均值和方差,然后使用这些统计量来归一化每个样本的分布匹配奖励。损失函数包括分布匹配损失和可选的外部奖励损失。重要性采样通过对不同样本赋予不同的权重来提高采样效率,权重基于样本的奖励值计算。
🖼️ 关键图片
📊 实验亮点
实验结果表明,提出的GNDM方法优于传统的扩散匹配蒸馏(DMD),FID降低了1.87。多奖励变体GNDMR在美学质量和保真度之间取得了更好的平衡,HPS达到了30.37,FID-SD降至12.21,超过了现有的基线方法,证明了该方法的有效性。
🎯 应用场景
该研究成果可应用于图像、音频、视频等多种生成任务,尤其适用于对实时性要求较高的场景,如游戏、虚拟现实、内容创作等。通过提高生成速度和质量,可以显著提升用户体验,并降低计算成本。未来,该方法有望进一步扩展到其他生成模型和任务中。
📄 摘要(原文)
Diffusion models achieve state-of-the-art generative performance but are fundamentally bottlenecked by their slow iterative sampling process. While diffusion distillation techniques enable high-fidelity few-step generation, traditional objectives often restrict the student's performance by anchoring it solely to the teacher. Recent approaches have attempted to break this ceiling by integrating Reinforcement Learning (RL), typically through a simple summation of distillation and RL objectives. In this work, we propose a novel paradigm by reconceptualizing distribution matching as a reward, denoted as $R_{dm}$. This unified perspective bridges the algorithmic gap between Diffusion Matching Distillation (DMD) and RL, providing several key benefits. (1) Enhanced optimization stability: we introduce Group Normalized Distribution Matching (GNDM), which adapts standard RL group normalization to stabilize $R_{dm}$ estimation. By leveraging group-mean statistics, GNDM establishes a more robust and effective optimization direction. (2) Seamless reward integration: our reward-centric formulation inherently supports adaptive weighting mechanisms, allowing flexible combination of DMD with external reward models. (3) Improved sampling efficiency: by aligning with RL principles, the framework readily incorporates importance sampling (IS), leading to a significant boost in sampling efficiency. Extensive experiments demonstrate that GNDM outperforms vanilla DMD, reducing the FID by 1.87. Furthermore, our multi-reward variant, GNDMR, surpasses existing baselines by achieving a strong balance between aesthetic quality and fidelity, reaching a peak HPS of 30.37 and a low FID-SD of 12.21. Overall, $R_{dm}$ provides a flexible, stable, and efficient framework for real-time high-fidelity synthesis. Code will be released upon publication.