DMAD:将分布匹配作为对抗蒸馏,实现快速视觉生成
DMAD: Distribution Matching as Adversarial Distillation for Fast Visual Generation
精选理由
以对抗蒸馏替代辅助扩散模型,降低快速视觉生成成本并提升一至四步质量
AI 摘要
分布匹配蒸馏(DMD)通过分别估计的目标得分函数与学生得分函数之间的差异来训练少步数学生模型,因此必须额外维护一个拟合学生模型不断变化分布的辅助扩散模型,从而增加内存和计算开销。我们提出 DMAD(Distribution Matching as Adversarial Distillation),将分布匹配重新表述为分类问题,并直接学习所需的对数密度比。我们在共享骨干网络上设置两个判别器头,分别区分真实数据和教师样本与学生样本,并通过对其 logits 施加线性损失,在无需额外拟合辅助得分函数的情况下训练学生模型。我们证明,在判别器达到最优时,根据将判别器 logits 与对数密度比联系起来的经典恒等式,这些损失能够恢复 DMD 所依据的分布匹配梯度。我们进一步提出基于间隔的重加权方法,根据真实数据判别器头在真实样本与教师样本之间产生的经验 logit 间隔,自适应地调整不同噪声水平下的教师监督信号。在 ImageNet-64x64 上进行一步生成时,DMAD 的 Fréchet Inception Distance(FID)达到 1.04;在 COCO-10K 上使用四步 SDXL 时达到 14.47;在使用四步 Wan2.1-T2V-14B 时,VBench 总分达到 85.15;这些都是所比较的少步数方法和多步数教师模型中的最佳结果。在 MiniMax-H3-33B 上,我们的四步学生模型在联合音视频生成任务中,相较于 DMD2 和 rCM 的总体人类偏好率分别达到 79.1% 和 84.6%,不计平局。我们的代码、模型和演示可见于 https://yzmblog.github.io/projects/DMAD。