FlowTurbo: Towards Real-time Flow-Based Image Generation with Velocity Refiner

We sincerely thank the reviewer for the positive comments on our work! We address the questions and clarify the issues accordingly as described below.

Q1: About experiments on other models.

[Reply] Thanks for your advice. Since SD3 is open-sourced a few months later than the submission, we didn't include experimental results on it in our original paper. In the following table, we compare the sampling quality (measured by FID 30K $\downarrow$ on MSCOCO) and the sampling speed of the default sampler of SD3 (Euler's method) and our FlowTurbo:

Besides, we also provide qualitative comparisons on SD3 and the newly released SOTA model FLUX.1-dev in Figure A2 and Figure A3 in the attached one-page PDF, where we demonstrate FlowTurbo can generate more details than the baseline.

Q2: About different model sizes.

[Reply] Thanks for pointing out this. In our original paper, we conducted experiments with models of various sizes including SiT (675M), InstaFlow (860M), Lumina-Next-T2I (2B). We also agree that it is interesting to see the performance on similar architectures with different model sizes such as SiT-XL and SiT-L. However, SiT only released their SiT-XL-2 model and we failed to find another publicly available family of models that satisfied the requirement. We will add the comparison once the SiT releases another model such as SiT-L or SiT-B.

Q3. About the refiner architecture

[Reply] When designing the architecture for the velocity refiner, we followed a simple rule to make the refiner have a similar architecture as the base velocity predictor but with much fewer parameters (~5% of the base model). The detailed architecture is described in Section 4.1 and Appendix C. For example, since SiT consists of multiple transformer blocks, we simply use a single block as the refiner. For text-to-image generation, we reduce the number of layers and channels of the UNet. In our early experiments, we did have tried another architecture for class-conditional image generation, where a SiT-S (a smaller version of the base velocity predictor SiT-XL) is adopted as the refiner (as shown in the following table):

We find that using a block of SiT-XL as the refiner is slightly better than the SiT-S. These results demonstrate that our framework is robust to the choice of model architectures for the velocity refiner. We will add these discussions and results to Appendix C of the revised paper.

Q4. About the application to video generation models

[Reply] Thanks for your advice. Since our FlowTurbo is designed to accelerate flow-based generative models, it can be applied to video generation models as long as they are trained via flow-matching (i.e., has a stable velocity during the sampling). Specifically, we experiment with the recently released model Open-Sora and provide the qualitative results in Figure A4 of the attached one-page PDF, where we show our FlowTurbo generates video with higher fidelity under the same computational budgets.

We sincerely thank the reviewer for the positive comments on our work! We address the questions and clarify the issues accordingly as described below.

Q1: About the design of sampling blocks.

[Reply] Sorry for the confusion. As illustrated in Figure 1, the velocity during the sampling process stabilizes in the final few steps. At this stage, we employ a lightweight refiner to adjust the velocity offset. Additionally, our pseudo corrector is crafted to efficiently achieve second-order convergence, which requires second-order intermediate results (achieved by several Heun's steps at the start) for initialization. Consequently, the sequence of the three block types should be: Heun's block, pseudo corrector block, and refiner block. As the results shown below, changing the order of the blocks will lead to worse sampling quality:

Regarding the number of steps for each block, we have included comparisons in Table 3 and Table 4. Generally, we apply the refiner during the last 20-40% of the sampling steps, where the velocity is relatively stable. We can always adjust the ratio between Heun's steps and pseudo-corrector steps to balance sampling quality and complexity (as seen in Table 3d). Moreover, since reconfiguring the steps doesn't require additional training, we can select an optimal configuration based on computational resources and qualitative results. We will incorporate this detailed analysis and discussion into our revised paper to offer better guidance on designing the configuration of our FlowTurbo.

Q2: About the visualization.

[Reply] Many thanks for your valuable advice. We agree that visualizing the curvature is more suitable and accurate. We provide the comparisons of the curvatures of different models including DiT, SiT, SD3, FLUX.1-dev (a newly released SOTA flow model), Open-Sora (a flow-based video generation model) in the Figure A1 in the attached PDF. The visualization of curvature basically aligns with our original PCA visualization, where we find the flow-based models enjoys much smaller curvature, indicating the sampling trajectory is straighter. We will replace the visualization in Figure 1 with curvature in our revised paper.

Q3: About the related work.

[Reply] Thanks for pointing out this. We agree that these two papers are related to our paper and we will discuss them in the related work (L107) as follows:

There are also some works to accelerate the sampling of flow-based models, such as computing the optimal sampling stepsizes [1] and solver distillation approaches [2].

In our FlowTurbo, we do not consider the variation of stepsizes during the sampling, which makes our method orthogonal to [1]. In other words, we can combine [1] and our FlowTurbo to improve the sampling efficiency. Following [1], we perform experiments on CelebA-HQ and the results are as follows:

In the following table, we compare our FlowTurbo with Bespoke Non-Stationary (BNS) Solvers [2] on ImageNet. Note that BNS only provides the results on ImageNet 64x64 while we adopt the experimental while our experiment setting follows SiT (ImageNet 256x256) and we cannot find a good publicly available flow-based model trained on ImageNet 64x64. The different resolution and base models choices might slightly affect the results. We also report the number of function evaluations (NFE) in the table, where we use $f_1$ and $f_2$ ($f_2\approx 0.1 f_1$) to represent the velocity predictor and refiner, respectively.

We also provide the result on MS COCO text-to-image generation:

These results clearly show that our FlowTurbo is better in both sampling quality and training efficiency. Besides, BNS solvers struggle to generalize unseen NFEs while our FlowTurbo allows flexible adjustment of the sampling configuration after the velocity has been trained. We will add the above discussion and comparisons in our revised paper.

Q4: Applying FlowTurbo for diffusion models.

[Reply] That's a good question. In our early experiments, we did try to transfer a similar idea of FlowTurbo to diffusion models. Specifically, we adopted PixArt as the base model (diffusion transformer) and tried to learn a refiner for the $\epsilon_\theta$. However, just as Figure 1 suggests, the $\epsilon_\theta$ of diffusion models changes a lot during the sampling and it is really hard to regress the offset of $\epsilon_\theta$. We have also tried to convert the $\epsilon_\theta$ to an equivalent velocity through Equation (32):

$\boldsymbol{v}=\frac{\dot{\alpha}_t}{\alpha_t}x + (\dot{\sigma}_t-\frac{\dot{\alpha}_t\sigma_t}{\alpha_t})\boldsymbol{\epsilon},$ but the refiner still cannot converge. We think some techniques such like [3] to transform the sampling path to be straighter could be helpful, and we will leave this for future work.

[3] Shaul et. al, Bespoke Solvers for Generative Flow Models, arxiv 2023, ICLR 2024

We sincerely thank the reviewer for the positive comments on our work! We address the questions and clarify the issues accordingly as described below.

Q1: About prior works on flow-based methods.

[Reply] Thanks for your suggestions. We agree that these works are important in the area of flow-based models, and we will include the mentioned papers in the revised paper. Specifically, we will modify L31 to:

Alongside the research on diffusion models, flow-based models have garnered increasing attention due to their versatility in modeling data distributions, and have been widely applied in various domains including image generation [5], video generation [6, 7], video prediction [2], human motion synthesis [4], depth estimation [8], text generation [9], etc.

And for the acceleration methods [10-13] for flow-based models, we will add the discussion of these methods in our related work (L107):

There are also some work aim to obtain better flow paths which are straighter and more effective for sampling. These work includes some useful technique such as transformed sampling paths [10], minimizing the trajectory curvature [11], multisample flow matching [12], and minibatch optimal transport [13].

Q2: About the configuration of the sampling blocks.

[Reply] Sorry for the confusion. According to the observation in Figure 1, the velocity during the sampling would become stable at the final few steps, where we adopted a lightweight refiner to regress the velocity offset. Besides, our pseudo corrector is designed to efficiently achieve 2-order convergence, which requires a 2-order intermediate result as initialization. This explains why we need several Heun's steps at the beginning. We also provide some comparisons of different orders and schemes (including repeating the blocks) as follows:

As for the number of different steps, we have included some comparisons in Table 3 and Table 4. Generally speaking, we choose to apply the refiner at the last $20\sim 40$% sampling steps where the velocity is relatively stable, and we can always adjust the ratio between Heun's steps and pseudo-corrector steps to control the trade-offs between sampling quality and complexity (Table 3d). Besides, since adjusting the configuration does not require further training, we can always select a good configuration according to the computational budgets and qualitative results. We will add these detailed analyses and discussion to our revised paper to provide better guidelines on how to design the configuration of our FlowTurbo.

Q3: About the evaluation.

[Reply] Thanks for the comments. We provide the results in Table 1a and 1b with a column of NFEs and the baseline methods compiled as follows. We use $f_1$ to denote the function evaluation of the original velocity predictor and $f_2$ to denote the evaluation of our velocity refiner.

Table 1a: Class-conditional Image Generation, ImageNet (256x256)

Table 1b: Text-to-image Generation, MS COCO 2017 (512x512)

These results clearly show that our method still surpasses the compiled baseline by large margins, demonstrating the effect of our velocity refiner and the pseudo corrector. We will modify these in the revised paper.

Q4: About the description of the linear interpolation.

[Reply] Thanks for pointing out this. We agree that linear interpolation connects noise distribution to the distribution concentrated around a particular data point in the conditional flow-matching defined in [14]. We will modify the description in L37, L89, L141 in the revision.

Q5: About the typos.

[Reply] Thanks for your careful reading. We will modify these typos accordingly in the revised paper.

We thank the reviewer for the valuable comments. We address the questions and clarify the issues accordingly as described below.

Q1: About comparisons to distillation-based methods.

[Reply] Thanks for your suggestions. We have detailedly discussed the difference between our FlowTurbo and distillation-based methods in our original paper (L62-L68, L100-L104, L205-L214). Although the previous diffusion distillation methods can be directly adapted to flow-based models, they do not investigate the unique properties of the flow sampling trajectory. Our method, however, aims to offer a training-efficient solution for accelerating flow-based image generation. As Reviewer PyDm suggests, the low additional training costs can make our FlowTurbo have a high impact on the acceleration of generative models.

We also agree that it would be better to include more distillation methods for comparison in Table 2. We summarize the comparisons between our FlowTurbo and [1-3] on ImageNet in the following table. Note that all these methods only perform experiments on ImageNet 64x64, while our experiment setting follows SiT (ImageNet 256x256) and we cannot find a good flow-based model trained on ImageNet 64x64. The different resolution and base model choices might slightly affect the results.

Besides, we also provide comparisons with more distillation methods [4-5,7-8] on text-to-image. We adopt the FID 30K on COCO as the evaluation metric. We do not include [6] since we cannot find the same evaluation metric in their paper and the models have not been released.

From the above results, we demonstrate that our FlowTurbo can achieve competitive sampling quality with previous diffusion distillation methods while requiring much fewer training costs. We will add these results to our revised paper to make the comparison clearer.

[7] Luo, Simian, et al. "Latent consistency models: Synthesizing high-resolution images with few-step inference." arXiv preprint arXiv:2310.04378 (2023).

[8] Luo, Simian, et al. "LCM-LoRA: A universal stable-diffusion acceleration module." arXiv preprint arXiv:2311.05556 (2023).

Q2: About the implementation.

[Reply] Thanks for the valuable comments. The motivation to adopt the velocity refiner is quite clear according to Figure 1, where we show the velocity tends to become more and more stable during the sampling. Heun's method is a well-known ODE sampler that is adopted in our baseline SiT, and the pseudo-corrector is proposed to reduce the inference costs while attaining the same convergence order as Heun's method. Therefore, the sampling paradigm of FlowTurbo is quite intuitive and reasonable. As for the scalability of our method, we have conducted experiments with SiT (675M), InstaFlow (860M), and Lumina-Next-T2I (2B) in our paper. Besides, we also conduct experiments on the provide the newly released Stable-Diffusion-3-Medium (8B) and compare our FlowTurbo with the default sampler Euler's method:

The sampling quality is measured by FID 30K $\downarrow$ on MSCOCO. Besides, we also provide some qualitative comparisons on SD3 and the newly released SOTA model FLUX.1-dev (12B) in Figure A2 and Figure A3 in the attached one-page PDF, where we demonstrate that FlowTurbo can generate more visual details.

We believe the results on SD3 and FLUX.1-dev, together with the experiments on various flow-based models in our original paper can demonstrate the effectiveness of our FlowTurbo at scale.

Q3: Applying FlowTurbo for diffusion models.

[Reply] Thanks for pointing out this. In our early experiments, we did try to transfer FlowTurbo to diffusion models. Specifically, we adopted PixArt as the base model (diffusion transformer) and tried to learn a refiner for the $\epsilon_\theta$. However, just as Figure 1 suggests, the $\epsilon_\theta$ of diffusion models changes a lot during the sampling and it is really hard to regress the offset of $\epsilon_\theta$. Alongside attempting to directly regress the offset of $\epsilon_\theta$, we have also tried to convert the $\epsilon_\theta$ to an equivalent velocity through Equation (32):

$\boldsymbol{v}=\frac{\dot{\alpha}_t}{\alpha_t}\boldsymbol{x} + (\dot{\sigma}_t-\frac{\dot{\alpha}_t\sigma_t}{\alpha_t})\boldsymbol{\epsilon},$ but the refiner still cannot converge. To better understand the difference between diffusion and flow-based models, we plot the curvature of the sampling trajectories in Figure A1 in the attached PDF, where we show the flow-based models enjoy lower curvature (indicating the sampling path is straighter). We will add more theoretical and empirical analysis in the revision.