2024/07/17 by Yuanzhi Zhu, Xingchao Liu, Zhu, Yuanzhi +3 · 1 voice · 11 citations
Computer Science · Engineering · Physics and Astronomy · #Generative Adversarial Networks and Image Synthesis #Lattice Boltzmann Simulation Studies #Model Reduction and Neural Networks #cs.CV
paper · pdf · doi:10.48550/arxiv.2407.12718
arxiv published 2024/07/17 · arxiv updated 2024/07/18
Diffusion models excel in high-quality generation but suffer from slow inference due to iterative sampling. While recent methods have successfully transformed diffusion models into one-step generators, they neglect model size reduction, limiting their applicability in compute-constrained scenarios. This paper aims to develop small, efficient one-step diffusion models based on the powerful rectified flow framework, by exploring joint compression of inference steps and model size. The rectified flow framework trains one-step generative models using two operations, reflow and distillation. Compared with the original framework, squeezing the model size brings two new challenges: (1) the initialization mismatch between large teachers and small students during reflow; (2) the underperformance of naive distillation on small student models. To overcome these issues, we propose Annealing Reflow and Flow-Guided Distillation, which together comprise our SlimFlow framework. With our novel framework, we train a one-step diffusion model with an FID of 5.02 and 15.7M parameters, outperforming the previous state-of-the-art one-step diffusion model (FID=6.47, 19.4M parameters) on CIFAR10. On ImageNet 64×64 and FFHQ 64×64, our method yields small one-step diffusion models that are comparable to larger models, showcasing the effectiveness of our method in creating compact, efficient one-step diffusion models.