arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2410.23274cs.LGcs.AIcs.CV

面向更优单步生成器的多学生扩散蒸馏

Multi-student Diffusion Distillation for Better One-step Generators

Harvard University(哈佛大学);NVIDIA(英伟达);Vector Institute(向量研究院)

Yanke Song, Jonathan Lorraine, Weili Nie, Karsten Kreis, James Lucas

AI总结:

本文提出多学生蒸馏框架MSD,将条件扩散教师拆分为多个负责不同条件数据子集的单步学生,以更小模型实现更快推理,并在ImageNet和COCO上取得优异FID。

摘要:

扩散模型以冗长的多步推理过程为代价实现了高质量样本生成。为克服这一问题,扩散蒸馏技术生成能够在单步内匹配甚至超越教师模型的学生生成器。然而,学生模型的推理速度受限于教师架构的规模,导致计算密集型应用无法实现实时生成。在本工作中,我们提出多学生蒸馏(MSD),这是一个将条件教师扩散模型蒸馏为多个单步生成器的框架。每个学生生成器负责条件数据的一个子集,从而在相同容量下获得更高的生成质量。MSD训练多个蒸馏学生,使模型可以采用更小的尺寸,因此推理更快。此外,在相同架构下,MSD相较单学生蒸馏提供了轻量级的质量提升。我们通过使用分布匹配和对抗蒸馏技术训练多个同等尺寸或更小的学生进行单步蒸馏,证明了MSD的有效性。使用更小的学生时,MSD在单步生成中以更快的推理获得具有竞争力的结果。使用4个同等尺寸的学生时,MSD显著优于单学生基线对应方法,并在单步图像生成上取得出色的FID分数:ImageNet-64x64上为1.20,零样本COCO2014上为8.20。

英文摘要:

Diffusion models achieve high-quality sample generation at the cost of a lengthy multistep inference procedure. To overcome this, diffusion distillation techniques produce student generators capable of matching or surpassing the teacher in a single step. However, the student model's inference speed is limited by the size of the teacher architecture, preventing real-time generation for computationally heavy applications. In this work, we introduce Multi-Student Distillation (MSD), a framework to distill a conditional teacher diffusion model into multiple single-step generators. Each student generator is responsible for a subset of the conditioning data, thereby obtaining higher generation quality for the same capacity. MSD trains multiple distilled students, allowing smaller sizes and, therefore, faster inference. Also, MSD offers a lightweight quality boost over single-student distillation with the same architecture. We demonstrate MSD is effective by training multiple same-sized or smaller students on single-step distillation using distribution matching and adversarial distillation techniques. With smaller students, MSD gets competitive results with faster inference for single-step generation. Using 4 same-sized students, MSD significantly outperforms single-student baseline counterparts and achieves remarkable FID scores for one-step image generation: 1.20 on ImageNet-64x64 and 8.20 on zero-shot COCO2014.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑