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

摘要

Diffusion and flow-based models nowadays have notable success in generating diverse and high-quality images. However, their iterative sampling process and substantial model size pose challenges in fast generation and downstream applications. In this paper, we propose SlimReFlow, a new method to obtain efficient one-step diffusion models as an extension of Rectified Flows. We first propose Annealing Rectifying to avoid the training of 1-flows, which results in 2-flow models directly from the data pairs generated by pre-trained diffusion models. Then we introduced a new distillation loss with additional supervision from the 2-flow models to get better one-step distilled flows. We illustrate the versatility of our method by applying it to various diffusion and flow-based models, including Rectified Flows and EDM, as shown through the use of data pairs. Through extensive examination of various model sizes and dataset choices, we demonstrate that our approach can significantly lower the number of parameters in the models while maintaining the quality of one-step image generation. Our method achieves an FID of 5.02 with 15.7M parameters and 4.53 with 27.9M parameters on CIFAR-10 32$\times$32 in one-step generation. On the FFHQ 64$\times$64 dataset, we record FIDs of 7.70 and 7.21 with 15.7M and 27.9M parameters, respectively. Additionally, on the ImageNet 64$\times$64 dataset, our method secures an FID of 12.34 using only 80.7M parameters.