用多个小模型替代大模型,实现更快更高质量的单步图像生成
Multi-student Diffusion Distillation for Better One-step Generators
- 将一个大教师模型拆成多个小学生模型,各自负责部分条件数据
- 4个同尺寸学生比单个学生在相同条件下显著提升生成质量
- 适合需要实时生成的场景,如移动端图像生成
扩散模型虽能生成高质量样本,但需多步推理。为加速,扩散蒸馏可训练单步生成的学生模型,但其速度受限于教师模型规模。本文提出多学生蒸馏(MSD),将条件扩散教师模型蒸馏为多个单步生成器,每个负责特定条件数据子集。通过训练多个小型学生模型,实现更低延迟与更高生成质量。实验表明,使用4个同尺寸学生,相比单学生基线显著提升性能:在ImageNet-64x64上达到1.20的FID,在零样本COCO2014上达8.20,同时保持快速推理。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。