SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow

TL;DR

This paper introduces SlimFlow, a framework for training small, efficient one-step diffusion models using rectified flow, achieving high-quality image generation with significantly reduced model size and inference time.

Contribution

The paper proposes Annealing Reflow and Flow-Guided Distillation techniques to effectively compress diffusion models into smaller, one-step generators, overcoming initialization mismatch and distillation challenges.

Findings

Achieved state-of-the-art FID of 5.02 on CIFAR10 with 15.7M parameters.

Produced small models on ImageNet and FFHQ comparable to larger models.

Demonstrated effective model compression without significant quality loss.

Abstract

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$\times$64 and FFHQ 64$\times$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.