用对抗奖励提升扩散模型少步生成质量,4步效果超原版40步。
AdvDMD: Adversarial Reward Meets DMD For High-Quality Few-Step Generation

- 将对抗判别器作为奖励模型,统一训练扩散蒸馏与强化学习。
- 4步生成在DPG-Bench上超越原40步模型,2步在Qwen-Image上优于TwinFlow。
- 在线更新奖励模型,避免奖励欺骗,适合高效高质生成场景。
扩散模型虽生成质量高,但需大量采样步数。蒸馏方法如分布匹配蒸馏(DMD)可减少步数,但在少步条件下性能仍下降明显。已有强化学习(RL)方法尝试提升少步生成质量,甚至超越教师模型,但多为组合式设计,引入额外复杂性。为此,本文提出AdvDMD,将DMD蒸馏与强化学习无缝融合。具体地,利用DMD2中对抗训练的判别器作为奖励模型,对生成图像打低分,真实图像打高分,并在去噪过程的中间与最终状态上联合训练,且随蒸馏模型在线更新,实现对采样轨迹的全局监督,有效防止奖励欺骗。采用统一的SDE反向模拟和差异化训练策略,提升训练稳定性与效率。实验表明,4步的AdvDMD在SD3.5上超越原始40步模型,在DPG-Bench上表现优异;对SD3在GenEval上也有显著提升。在Qwen-Image上,2步AdvDMD优于TwinFlow。
原文摘要 · Abstract (English)
Diffusion models offer superior generation quality at the expense of extensive sampling steps. Distillation methods, with Distribution Matching Distillation (DMD) as a popular example, can mitigate this issue, but performance degradation remains pronounced when sampling steps are limited. Reinforcement learning (RL) has been leveraged to improve the few-step generation quality during distillation, with the potential to even surpass the performance of the teacher model. However, existing approaches are combinatorial in nature, merely integrating an RL process with the distillation process, which introduces unnecessary complexities. To address this gap, we propose AdvDMD, a method that seamlessly unifies DMD distillation and RL. Specifically, AdvDMD employs the adversarially trained discriminator from DMD2 as the reward model, which assigns low scores to generated images and high scores to real ones. It is trained on both intermediate and final states of the denoising process and updated online with the distilled model, enabling a holistic supervision of the sampling trajectories and mitigating reward hacking. We adopt a unified SDE backward simulation and a different training schedule for DMD and RL to enable a more stable and efficient training. Experimental results demonstrate that the 4-step AdvDMD outperforms the original 40-step model for SD3.5 on DPG-Bench, while achieving significant performance gains for SD3 on the GenEval. On Qwen-Image, our 2-step AdvDMD achieves superior performance over TwinFlow.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。