Skip to yearly menu bar Skip to main content


Poster

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

Yuanzhi Zhu · Xingchao Liu · Qiang Liu

# 52
Strong blind review: This paper was not made available on public preprint services during the review process Strong Double Blind
[ ] [ Paper PDF ]
Thu 3 Oct 1:30 a.m. PDT — 3:30 a.m. PDT

Abstract: 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.

Live content is unavailable. Log in and register to view live content