论文

DMD联合扩展分数估计网络容量并提供步数条件,改善少步蒸馏

A Better Critic Matters for Distribution Matching Distillation

摘要

分布匹配蒸馏(DMD)利用教师得分函数与在线评估模型对学生分布得分函数的估计之间的差异,训练少步学生模型。我们考察评估模型需要具备怎样的容量、获得哪些信息,才能提供可靠训练信号。在容量方面,ImageNet 上涵盖七种 SiT 模型规模的受控实验表明,增大教师而保持小型评估模型不变,可能降低学生质量并导致训练崩溃。让评估模型随教师同步扩展——我们称为“联合扩展”——在本实验中避免了崩溃,使学生能够受益于更大的教师。在信息方面,我们不再要求评估模型拟合不区分生成步骤的学生输出混合分布,而是引入步骤感知评估模型(STAC),以学生生成步骤为条件,估计各步骤对应的得分函数。联合扩展与 STAC 相结合,在 ImageNet 类别条件生成、使用 SDXL 和 FLUX.2 的文生图,以及使用 Wan2.1 的文生视频上,均持续改善了少步蒸馏效果。

DMD联合扩展评论器容量并提供步数条件改善少步蒸馏:论文配图
图1:使用与教师同样大的评价器,蒸馏少步数学生模型。上方:STAC用32B教师与32B评价器蒸馏出的四步FLUX.2 [klein] 4B学生样本。下方:STAC用14B教师与14B评价器蒸馏出的四步Wan2.1 1.3B学生生成六段视频,各展示三帧。设置为FLUX.2 1024×1024、Wan2.1 480p。