InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation

Thank you for the positive review and valuable comments! We will address your concerns below.

Q#1 (Storage Overhead): We admit additional storage space is required to apply our current pipeline. However:

• Stable Diffusion is trained using more than 1B (text, image) pairs, while InstaFlow only uses 1.6M to achieve the current performance. The overhead is very small.
• Storage in the training stage is relatively cheap. In contrast, after training, the inference cost of the deployed model usually exceeds the training cost (consider the case of GPT and Midjourney).
• The distillation step can be further replaced by data-free distillation, like consistency distillation, to save the storage space.
• Generating the reflow dataset is fast, since only forward inference is involved.

Overall, we argue that the storage overhead is small compared with the original dataset, and the advantage in efficiency brought to the inference stage by our pipeline is significant and overweighs the storage overhead in training.

Q#2 (Intrinsic Difference between Distillation and Reflow): The intrinsic, fundamental difference is the coupling between the noise distribution and the image distribution. For distillation, the one-step generative model is represented by, $f_{distill}(X_0) = X_0 + v_\theta(X_0, 0)$

For reflow, the continuous-time generative model is represented by, $f_{flow}(X_0) = X_0 + \int_0^1 v_\theta(X_t, t) d t$

When distilling from a existing probability flow, the objective is, $E_{X_0 \sim \pi_0, X_1 := f_{flow}(X_0)} [D (f_{distill}(X_0), X_1)]$ The coupling between noise and image $(X_0, X_1)$ is determined by the flow $f_{flow}$, and the distilled one-step student tries to imitate these couplings. If the couplings $(X_0, X_1)$ are suboptimal, it could be too difficult for the student distilled model to learn, and make the distillation performance bad. Instead, reflow optimizes the whole velocity field, get new probability flow $f_{flow}^{new}$, and refines the coupling $(X_0, X_1^{new}:=f_{flow}^{new}(X_0))$. The new coupling induced from the new flow is easier to distill for the one-step student model. We explained these subtle differences at the top of Page 6 in our submission.

Q#3 (Noise Scheduler): Thank you for the question! We will answer according to our understanding, but further elaboration on the question is welcomed! Rectified flow trains with a new linear interpolated trajectory instead of the old DDPM noise schedule. In our practice, we found the pre-trained Stable Diffusion neural networks have an amazing ability to adapt to the new trajectory without catastrophic forgetting of the learned contents. At the beginning of the project, just like you, we were expecting potential issues, but it turns out the pre-trained networks have superior adaptivity. We think this is a valuable observation to the community as well, since it makes researchers confident to fine-tune with many other alternative trajectories. We will add that to the paper.

Q#4 (Time Embedding): We use the same time embedding as the original SD. The only difference is that their time steps are integers while we relax that to continuous values.

Q#5 (Network Prediction): We were following the practice in previous works. In fact, predicting $X_1 - X_0$ is equivalent to $X_1 - X_t$ up to a weight coefficient because, $X_1 - X_t = X_1 - (t X_1 + (1-t)X_0)= (1-t) X_1 - (1-t) X_0 = (1-t) (X_1 - X_0).$ Therefore, $||X_1 - X_t - f_\theta (X_t, t) ||^2 = || (1-t) (X_1 - X_0) - f_\theta (X_t, t) ||^2 = (1-t)^2|| X_1 - X_0 - \frac{f_\theta (X_t, t)}{1-t}||^2.$ So the only difference is the additional time-related weight coefficient $(1-t)^2$ if we parameterize $v_\theta(X_t, t) = \frac{f_\theta (X_t, t)}{1-t}$. The weighting is important as discussed in previous works (e.g., [1]), but we leave that for future investigation as uniform weighting works perfectly in our current practice.

[1] Denoising Diffusion Probabilistic Models, Ho et al.

Thank you for the positive review and valuable comments! We will address your concerns below.

Q#1 (Performance loss): We observe the training of 2-Rectified Flow is not totally converged and most text-to-image models are trained using more than 1B text prompts (InstaFlow only uses 1.6M) and thousands of GPU days (InstaFlow only uses 199). Therefore, we believe InstaFlow can be hugely improved by investing more computation and resources, though that is beyond our reach. The emphasis of our research is that text-conditioned reflow has a decisive impact on distilling Stable Diffusion, and we get state-of-the-art one-step Stable Diffusion from that insight.

Q#2 (Open-source): We will open-source the training/inference codes along with pre-trained checkpoints upon acceptance.

Thank you for the review, comments and valuable questions! We appreciate that the reviewer recognizes the excellent results of our paper. We will address your concerns below. Please kindly raise the rating if our response makes sense to you.

Q#1(Novelty): Besides the novelty/difference/interesting phenomenon mentioned in General Response #1, we would like to claim that scaling-up is never straightforward. Rather, it is a resource-consuming process with much back-and-forth. However, despite the difficulties, successful scaling-up is indispensable for the modern machine learning community to acknowledge the value of a specific algorithmic framework. The new, large-scale, extensive empirical comparison and state-of-the-art generation results in our work, using the novel text-conditioned reflow framework, sets a solid foundation for the value of the methodologies / theories and inspires practitioners / methodologists to improve the existing algorithms in the future. This, in our opinion, is of great interest to the ICLR community. We will open-source the codes and pre-trained models to further advance the research of generative models.

Q#2 (Network Architecture): Please refer to General response #2. Most of the results are generated by InstaFlow-0.9B using the original U-Net as SD.

Q#3 (Performance Loss): We observe the training of 2-Rectified Flow is not totally converged and most text-to-image models are trained using more than 1B text prompts (InstaFlow only uses 1.6M) and thousands of GPU days (InstaFlow only uses 199). Therefore, we believe InstaFlow can be hugely improved by investing more computation and resources, though that is beyond our reach. The emphasis of our research is that text-conditioned reflow has a decisive impact on distilling Stable Diffusion, and we get state-of-the-art one-step Stable Diffusion from that insight.

Q #4 (Flow Trajectory): Thank you for the question! We will answer according to our understanding, but further elaboration on the question is welcomed! InstaFlow abandons the original DDPM trajectory of the Stable Diffusion, and re-trains a new linear interpolated trajectory. In our practice, we found the pre-trained Stable Diffusion neural networks have an amazing ability to adapt to the new trajectory without catastrophic forgetting of the learned contents. At the beginning of the project, just like you, we were expecting potential issues, but it turns out the pre-trained networks have superior adaptivity. We think this is a valuable observation to the community as well, since it makes researchers confident to fine-tune with many other alternative trajectories. We will add that to the paper. Thank you for asking again.

Q #5 (DDIM Inversion): Thank you for the suggestion! We put additional results in Figure 26 of the revised Appendix. 2-Rectified Flow can efficiently encode the image to the latent noise space with even 1 step.

Q #6 (Ethical Implications): Thank you for the suggestion! We have added a discussion of societal impact in the revised version of the paper.

Further discussion are welcomed!

Thank you for the review, comments and valuable questions! We appreciate that the reviewer recognizes the contribution of our paper to the broad ML community. We will address your concerns below. Please kindly raise the rating if our response makes sense to you.

Q#1 (Novelty): Besides the novelty/difference/interesting phenomenon mentioned in General Response #1, we would like to claim that scaling-up is a resource-consuming process with much back-and-forth, and successful scaling-up is indispensable for the modern machine learning community to acknowledge the value of a specific algorithmic framework. In our work, we use new large-scale models, new large-scale datasets and new text-conditioned rectified flow framework, which are not presented in the previous works. Moreover, we extensively provide new empirical comparison and state-of-the-art results on large-scale text-to-image generation, which is a significant improvement over the previous work that has experiments only on toy dataset CIFAR10. The framework and the results, in our opinion, are of great interest to the ICLR community. We will open-source the codes and pre-trained models to further advance the research of generative models.

Q#2 (Network Structure): Thank you for proposing the alternatives! We discuss them below:

• Increasing Depth: Stacked U-Net is actually increasing the depth of the original U-Net by adding more layers. Please refer to Figure 12 c.
• Increasing the number of channels: In our early experiments, we tried to double the number of channels in each layer of the original U-Net, and found no obvious improvement over the original U-Net. This suggests the importance of increasing depth. We will add discussion of this result in the future versions.

Q#3 (Deep Floyd): Thank you for the suggestions! We are definitely interested in applying our text-conditioned rectified flow to more advanced models like DeepFloyd and SDXL. However, for this work, our focus is to fairly compare different distillation methods to provide scientific insights on using text-conditioned reflow before distillation in large-scale text-to-image generation. Limited by the computational resources accessible to us, we leave applications to these advanced models for future development.

Q#4 (Training Pipeline): Thank you for your questions! The pipeline is a little bit complicated because we are limited by resources. While the original Stable Diffusion training adopts 32 * 8 A100 GPUs, we only have 8 A100 GPUs available. Therefore, the training pipeline is a mixture of empirical observations and trade-off between training time and training quality. Even for the original Stable Diffusion, their training pipeline is empirically divided into multiple stages (https://huggingface.co/runwayml/stable-diffusion-v1-5) due to the complexity of training large-scale models. For your questions, we will address the individual concerns below:

• Large batch size vs. small learning rate: We started from training with a batch size of 64, and found the loss converged after around 70,000 iterations (step 1). Then we actually tried both increasing batch size and decreasing learning rate, and found increasing batch size can further reduce the loss while decreasing learning rate has no obvious effect. Increasing batch size is also a common practice when training modern large-scale models, as adopted in GPT-3, PaLM, GLM-130B, etc..

• Distillation loss: Optimization on the L2 loss in the latent space of VAE is faster. It also consumes less GPU memory and thus allows larger batch size, because LPIPS loss is computed over the image and the gradient needs to be back-propogated through the VAE decoder. However, optimizing the LPIPS loss has an instant improvement on the generated image quality of the distilled model, so step 4 is indispensable. We actually tried the combination of L2 and LPIPS loss as you suggested, and found the visual quality is worse than using LPIPS loss only. Overall, step 3 is an efficient warm-up stage for step 4 given limited resources.

• Switching point: As discussed in sub-question 1, we switch from step 1 to step 2 after observing the (moving average) loss converged with small batch size. In step 2, we train the model up to our computational budget. The loss for training 2-Rectified Flow with a batch size of 1024 is not converged and may improve further given longer training time. We switch from step 3 to step 4 after the quality of the generated images saturates with L2 loss.

If there are enough resources, the training pipeline could be simplified to step 2+step 4 only (large batch size + LPIPS loss for distillation), and the training duration can be expanded for further improvement. The training pipeline could be suboptimal due to resource limitation but will be a good starting point for future research with our open-source.