AdvDMD:对抗性奖励与DMD结合实现高质量少步生成
AdvDMD: Adversarial Reward Meets DMD For High-Quality Few-Step Generation
AI总结:
本文提出AdvDMD方法,结合DMD蒸馏与强化学习,通过对抗训练的判别器作为奖励模型,提升少步生成质量并稳定训练过程。
摘要:
扩散模型在生成质量上表现优异,但需要大量的采样步骤。蒸馏方法,如分布匹配蒸馏(DMD),可以缓解这一问题,但在采样步骤受限时性能下降依然明显。强化学习(RL)已被用于改进蒸馏过程中的少步生成质量,甚至可能超越教师模型的性能。然而,现有方法本质上是组合性的,仅仅将RL过程与蒸馏过程整合,引入了不必要的复杂性。为了解决这一差距,我们提出了AdvDMD,一种无缝结合DMD蒸馏和RL的方法。具体而言,AdvDMD采用DMD2中对抗训练的判别器作为奖励模型,该模型对生成图像评分低,对真实图像评分高。它在去噪过程的中间和最终状态上进行训练,并与蒸馏模型在线更新,从而实现对采样轨迹的全面监督并缓解奖励黑客问题。我们采用统一的SDE反向模拟和不同的训练计划来训练DMD和RL,以实现更稳定和高效的训练。实验结果表明,4步AdvDMD在SD3.5上优于原40步模型,在DPG-Bench上,而在SD3上在GenEval上取得显著性能提升。在Qwen-Image上,我们的2步AdvDMD在TwinFlow上表现更优。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。