论文

DMAD:将分布匹配转化为对抗蒸馏,加速视觉生成

DMAD: Distribution Matching as Adversarial Distillation for Fast Visual Generation

摘要

分布匹配蒸馏(DMD)通过分别估计目标分布与学生分布的得分之差,训练少步生成的学生模型,因此需要持续更新一个适应学生分布变化的辅助扩散模型,增加内存和计算开销。本文提出DMAD,把分布匹配转化为分类问题,直接学习所需的对数密度比。共享主干上的两个判别头分别把真实数据和教师样本与学生样本区分开,以其logits上的线性损失训练学生,不再拟合辅助得分模型。作者利用判别器logits与对数密度比之间的经典关系,证明在判别器达到最优时,这些损失可恢复DMD的分布匹配梯度。方法还依据真实数据判别头在真实样本与教师样本间的logit差距,在不同噪声水平自适应调整教师监督权重。DMAD在ImageNet-64×64上单步生成的FID为1.04,在COCO-10K上四步SDXL的FID为14.47,四步Wan2.1-T2V-14B的VBench总分为85.15,均为所比较少步方法和多步教师中的最佳结果。在MiniMax-H3-33B音视频联合生成上,排除平局后,四步学生相对DMD2和rCM的人类偏好率分别为79.1%和84.6%。代码、模型及演示见https://yzmblog.github.io/projects/DMAD。