InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation
Authors: Xingchao Liu
Xiwen Zhang
Jianzhu Ma
Jian Peng
qiang liu
Venue: ICLR 2024 poster (2024)
Keywords: Diffusion Models
Generative Models
Acceleration
Primary area: generative models

Abstract: Diffusion models have revolutionized text-to-image generation with its exceptional quality and creativity. However, its multi-step sampling process is known to be slow, often requiring tens of inference steps to obtain satisfactory results. Previous attempts to improve its sampling speed and reduce computational costs through distillation have been unsuccessful in achieving a functional one-step model.
In this paper, we explore a recent method called Rectified Flow, which, thus far, has only been applied to small datasets. The core of Rectified Flow lies in its \emph{reflow} procedure, which straightens the trajectories of probability flows, refines the coupling between noises and images, and facilitates the distillation process with student models. We propose a novel text-conditioned pipeline to turn Stable Diffusion (SD) into an ultra-fast one-step model, in which we find reflow plays a critical role in improving the assignment between noise and images. Leveraging our new pipeline, we create, to the best of our knowledge, the first one-step diffusion-based text-to-image generator with SD-level image quality, achieving an FID (Fréchet Inception Distance) of $23.3$ on MS COCO 2017-5k, surpassing the previous state-of-the-art technique, progressive distillation, by a significant margin ($37.2$ $\rightarrow$ $23.3$ in FID). By utilizing an expanded network with 1.7B parameters, we further improve the FID to $22.4$. We call our one-step models \emph{InstaFlow}. On MS COCO 2014-30k, InstaFlow yields an FID of $13.1$ in just $0.09$ second, the best in $\leq 0.1$ second regime, outperforming the recent StyleGAN-T ($13.9$ in $0.1$ second). Notably, the training of InstaFlow only costs 199 A100 GPU days. Codes and pre-trained models are available at \url{github.com/gnobitab/InstaFlow}.

Published as a conference paper at ICLR 2024
INSTAFLOW: ONE STEP IS ENOUGH FOR HIGH- QUALITY DIFFUSION-BASER TEXT-TO-IMAGE GEN-
ERATION
Xingchao Liu'r Xiwen Zhang', Jianzhu Ma", Jian Peng", Qiang Liu'
Department of Computer Science, University of Texas Austin
2 Helixon Research
xcl cuucs utexas. edu, xiwenChelixen. com ma wajjianzhutttinngnnnn edu. en, jianpengéillinois edu,
edu, lqiang@cs utexas.ee..........a edu
ABSTRACT
Diffusion models have revolutionized text-to-image generation with its excep- tional quality and creativity. However, its multi-step sampling process is known to be slow, often requiring tens of inference steps to obtaín satisfactory results. Previous attempts improve its sampling speed and reduce computational costs through distillation have been unsuccessful in achieving functional one-step model. this paper, we explore a recent method called Rectified Flow 45 43 which, thus far, has only been applied to small datasets. The core of Rectified Flow lies in its reflow procedure, which straightens the trajectories of probability flows, refines the coupling between noises and images, and facilitates the distilla- tion process with student models. We propose novel text-conditioned pipeline to turn Stable Diffusion (SD) into ultra-fast one-step model, in which we find reflow plays a critical role improving the assignment between noises and im- pipeline, create, to the best of our knowledge, the
ages. Leveraging first one-step diffusion- based text-to-image generator with SD-level image qual- ity, achieving FID (Fréchet Inception Distance) of 23.3 on MS COCO 2017-5k, surpassing the previous state-of-the-art technique, progressive distillation (88) by a significant margin (37.2 + 23.3 FID). By utilizing an expanded network with 1.7B parameters, we further improve the FID to 22.4. We call our one- step models InstaFlow. On MS COCO 2014-30k, InstaFlow yields an FID of 13.1 in just 0.09 second, the best in < 0.1 second regime, outperforming the re- cent StyleGAN-T [73] (13.9 in 0.1 second). Notably, the training of InstaFlow only costs 199 A100 GPU days. Codes and pre-trained models are available at github. com/ gnobi tab Insstalstlfo
Stable Diffusion: N-step
LAFITE
30k on coco 2014
Step
Step
Step
Image
StyleGAN-T
InstaFlow-O
GLIDE
InstaFlow: One-Step
InstaFlow-1.78
N times faster
Dificion
Imagen
0.1 (og) Time to Generate One Image
3.0
10.0
Noise
Image
Figure 1: InstaFlow is high-quality one-step text-to-image model derived from Stable Diffusion 70 Within 0.1 second, it generates images with similar FID as StyleGAN-T 73 MS COCO 2014. The whole fine- tuning process to yield InstaFlow pure supervised learning and costs only 199 1000 GPU days.
INTRODUCTION
Stable Diffusion (70), StyleGAN-T 73 and GigaGAN 32 have demonstrated the remarkable ability to synthesize realistic, artistic, and detailed images based on textual descriptions. These advancements are made possible through the assistance of large-scale datasets 74] and models
Modern text-to- imoo-mmagee//a (T2I) generative models, such as DALL-E 666 67
Imagen 71 29
[32 66 70)
"This work done during internship Helixon Research

---

Published as a conference paper at ICLR 2024
(A) Images from one-slep Installow-OO99 (0.09s per image, 512 512)
V-rr
(B) One-step Instafllo-O099 SDXL-Refiner (1024 1024)
(C) Images from one-step InstaFlow-1.778 (0.12s per image, 512 x 512)
Figure 2: (A) Examples of 512 x 512 images generated from one-step InstaFlow -099 in 0.09s; (B) The images generated from one-step InstaFlow-O$3 -l99- can be further enhanced by SDXL-Retiner [62] to achieve higher resolution and finer details; (C) Examples of 512 512 images generated from step InstaFlow-1.72 0.12s. Inference time measured
machine with NVIDIA A 1000 GPU.
2

---

Published as a conference paper at ICLR 2024
However, despite their impressive generation quality, these models often suffer from excessive infer- ence time and computational consumption [29
71
66:
67 70). This can be attributed to the fact that
most of these models are either rutt-regressive 8 14 or diffusion models (28 [0]. For instance,
Stable Diffusion, even when using a state-of-the-ar sampler [4 49 77 typically requires more
than 20 steps to generate acceptable images. As result, prior works 72 58 82] have proposed em- ploying knowledge distillation on these models to reduce the required sampling steps and accelerate their inference. Unfortunately, these methods struggle in the small step regime. In particular, one- step large-scale diffusion models have not yet been developed. The existing one-step large-scale T2I generative models are StyleGAN-T 773 and GigaGAN [32,, which rely on generative adversarial training and require careful tuning of both the generator and discriminator.
In this paper. we present a novel one-step generative model derived from the open-source Stable Diffusion (SD). We observed that a straightforward distillation of SD leads complete failure. The primary issue stems from the sub-optimal coupling of noises and images, which significantly ham- pers the distillation process. To address this challenge, we leverage Rectified Flow [45 45] approach to generative models and optimal transport that learn straight flow models amendable to fast simulation with few or one Euler steps. Rectified flow starts from matching data distribution with a potentially curved flow model (known as 1-flow in (45),, similar to DDIM [77]. probability flow ODEs [80] and other flow-based methods [40, 2 It then deploys an unique reflow proce- dure to straighten the trajectories of the flows, thereby reducing the transport cost between the noise distribution and the image distribution. This improvement in coupling significantly facilitates the distillation process. In this work, we take large-scale text-to-image models 1-flow, and focus on straightening them with reflow.
a recent
Consequently, we succeeded in training the first one-step SD model capable of generating high- quality images with remarkable details. Quantitatively, our one-step model achieves a state-of-the- art FID score of 23.4 on the MS COCO 2017 7 dataset (5,000 images) with an inference time of only 0.09 second per image. It outperforms the previous fastest SD model, progressive distillation (58), which achieved an one-step FID of 37.2. For MS COCO 2014 (30,000 images). our one-step model yields an FID of 13.1 in 0.09 second, surpassing one of the recent large-scale text-to-image GANs, StyleGAN-T 773 (13.9 in 0.1s). Notably, this is the first time a distilled one-step SD model per- forms on par with GAN, with pure supervised learning. Discussion of related works is deferred to Appendix A due to limited space.
2 METHODS
2.1 EFFICIENT INFERENCE IS NEEDED FOR LARGE- SCALE TEXT-TO-IMAGE GENERATION
Recently, various of diffusion-basee text-to-image generators (59 precedented performance. Among them, Stable Diffusion (SD) 770 open-sourced model trained on LAION-5B (74), gained widespread popularity from artists and researchers. It is based on latent diffusion model 770.. which is denoising diffusion probabilistic model (DDPM) (28 30 running in learned latent space. Because of the recurrent nature of diffusion models, it usually takes more than 100 steps for SD to generate satisfying images. To accelerate the inference, series of post-hoc samplers have been proposed 4D 49 77.. By transforming the diffusion model into a marginal- preserving probability flow, these samplers can reduce the necessary inference steps to as few as 20 steps [49]. However, their performance starts to degrade noticeably when the number of inference steps is smaller than 10. For the <10 step regime, progressive distillation [72] 58 is proposed to compress the needed number of inference steps to 2-4. Yet, it is still an open problem if is possible to turn large diffusion models, like SD, into an one-step model with satisfying quality.
29 70 have emerged with un-
2.2 RECTIFIED FLOW AND REFLOW Rectified Flow 45 43] is unified ODE-based framework for generative modeling and domain transfer. It provides an approach for learning a transport mapping T between two distributions mo and गi Idd from their empirical observations. In image generation, mo is usually standard Gaussian distribution and πı the image distribution.
Rectified Flow learns to transfer To to T via an ordinary differential equation (ODE), or flow model
dzz
dt v v(z,,)) initialized from Zo ~ no. such that Z, Ri
where Rd x [0, 1] + Rd is velocity field, learned by minimizing simple mean square objective:
1/'' dt x, - (X,,,) 112 at with X, = 5(xx. X1,
min E Xixx
(2)
3

---

Published as a conference paper at ICLR 2024
Original Model
Reflowed Model
Curved Trajectories; Bad Coupling
Straighter Trajectories; Good Coupling
Reflow
Distill
Distill
One-Step
One-Step
Generation
Generation
2 Failure: blurry, low-quality
Success: clear, high-quality
Figure 3: An overview our pipeline for learning one-step large-scale ext-to-image generative models. Direct distillation from pre-trained diffusion models, c.g., Stable Diffusion, fails because their probability flow ODEs have curved trajectories and incur bad coupling between noises and images. After fine-tuned with text- conditioned reflow, the trajectories more friendly to distillation. Consequently, the distilled model generates clear, high-quality images one step. The text prompt is "A dog head in the universe with planets and stars".
straightened and the coupling refined, thus the reflowed model
d where X; = ob Xo, X1, 1) is any fimediifferentiable interpolation between Xo and X1, with di X; 0,o(Xo, (x((x,X111 X), The y is any coupling (To. π1). A simple example y is the independent coupling To Ti. which can be sampled empirically from unpaired observed data from To and mT Usually, e is parameterized as deep neural network and equation 2 solved approximately with stochastic gradient methods. Different specific choices of the interpolation process X, result in different algorithms. As shown in Liu et al. [45), the commonly used denoising diffusion implicit model (DDIM) (77) and the probability flow ODEs of Song al. (80] correspond X, xo +
X,, with specific choices of time-differentiable sequences Ote, ,r (see Liu et al. 45 for details). In rectified flow, however, the authors suggested a simpler choice of
d de
X, (1 X) + tX
X. X1 Xo.
which favors straight trajectories that play crucial role in fast inference, discuss in sequel.
Straight Flows Yield Fast Generation practice, the ODE
in equation | need to be approximated by numerical solvers. The most common approach is the forward Euler method, which yields
Zg
ODE Trajectory
Suaight Trajessoey
2 - 2, + " AIE NO. 13/N,
(4) Figure 4: ODEs with straight trajecto- rics admits fast simulation.
where we simulate with step size of 1/N and completes the simulation with N steps. Ob- viously, the choice yields cost-accuracy trade-off: large N approximates the ODE better but causes high computational cost. For fast simulation, it is desirable to learn the ODEs that can be simulated accurately and fast with a small N. This leads to ODEs whose trajectory are straight lines. Specifically, we say that ODE is straight (with uniform speed) if
Straight flow: Z: = tZ + (1-t)2o = Zo + to to(Zo.0 0), Vt [0.
In this case, Euler method with even a single step (N = 1) yields perfect simulation; See Figure Hence, straightening the ODE trajectories is essential way for reducing the inference cost.
Straightening Text-Conditioned Probability Flows via CoxttConditoeeed Reflow Reflow 45 is
an iterative procedure straighten the trajectories of rectified flow without modifying the marginal
4

---

Published as a conference paper at ICLR 2024
Algorithm Training Text-Conditioned Rectified Flow from Stable Diffusion
1: Input: The pre-trained Stable Diffusion U15 A dataset of text prompts Dy.
2: for user-defined upper bound do
3: Initialize ck+1 from Uk.-
4: Train Uk++ by minimizing the objective equation
where the couplings
(Xo, X1 = (DE(v)](Xo) (T)) can be generated beforehand.
5: #NOTE: The trained Up is called k-Rectified Flow.
6: end for
Algorithm 2 Distilling Text-Conditioned k-Rectified Flow for One-Step Generation
1: Input: k-Rectified Flow U4A A dataset of text prompts Dy; A similarity loss D(;,)).
2: Initialize Ug from vi.
3: Train Ux by minimizing the objective equation
where the couplings (Xo, X1 = ODE[u,(XXo T)) can be generated beforehand.
4: #NOTE: The trained ëz is called k-Rectified Flow+Distill.
distributions, hence allowing fast simulation at inference time. In text-to-image generation, the velocity field should additionally depend on an imput text prompt T to generate corresponding images. We propose novel text-conditioned reflow objective,
Uh+1 arg min Exo~mo.T~D,
|1 [1 Xu) v(x,, 7) 11 dt
dt}
(5)
with X1 ( 7) and X, = tX1 + (1 t)
where Dy is a dataset of text prompts, [DE learned using the same rectified flow objective equation 2 but with the linear interpolation equation of (Xo, X1) pairs constructed from the previous ODE(WR).
X₂ + +o' us (X,, dtt and 0k+1 is
The key property of reflow is that it preserves the terminal distribution while straightening the parti- cle trajectories and reducing the transport cost of the transport mapping:
1) The distribution of correct image distribution πi U2 does so.
T) and ODE[ux](Xg, T) coincides; hence uh+1 generates the
2) The trajectories of ODE tend to be straighter than that of ODE(va). This suggests that it requires smaller Fuler steps N simulate ODE[**+1] than ODE(vx]- If mg is fixed point of reflow, that is, Ub+1 Vk. then ODE(Uk] must be exactly straight.
3) (Xo, ODE[EEEE [0 7)) forms better coupling than (Xo, ODE (Xp lower convex transport costs, that is, ([(c(ODE(vx+
T)) that it yields T) Xo)] E[c(ODE(Un) [Xx Xo)] for all convex functions c: Rdd + R. This suggests that the new coupling might be easier for the student network to learn.
In this paper, we set U1 to be the velocity field of a pre-trained probability flow ODE model (such as that of Stable Diffusion,, Ugp), and denote the following 22 (k 2) as k-Rectified Flow.
Text- Cond-Coodiioondd Distillation Theoretically, it requires an infinite number of reflow steps equa- tion 5 to obtain ODEs with exactly straight trajectories. However, it is practical to reflow too many steps due to high computational cost and the accumulation optimization and statistical error. Fortunately,
observed in Liu al. (45) that the trajectories of ODE vx] becomes nearly (though not exactly) straight with even one or two steps of reflows. With such approximately straight ODBs, one approach boost the performance of one-step models is via distillation:
T = arg min Ex
IT] (ODE px] (X, TT,, X, + v v(X, [T)],
(6)
where we learn single Euler step u(x T) to compress the mapping from Xu to ODE (X, TT)
by minimizing a differentiable similarity loss ((.)) between images. Learning one-step model with
distillation avoids adversarial training 73 32 19
special invertible neural networks 9 35 60
5

---

Published as a conference paper at ICLR 2024
Distillation and Reflow are Orthogonal Techniques It is important to note the difference be- tween distillation and reflow: while distillation tries to honestly approximate the mapping from Xo to ODE[ux] (XXo T). reflow yields smooth due to lower convex transport costs. Reflow an optional step before distillation, and they orthogonal to each other. In practice, find that it is essential to apply reflow to make the mapping ODE|uR](X, T)) sufficiently regular and smooth before applying distillation.
mapping ODE(U*+1 (Xo, T) that can be more regular and
Classifier- Free Guidance Velocity Field for Text-Conditioned Rectified Flow Classifier-Fres Guidance (26) has a substantial impact on the generation quality of SD. Similarly, propose the following velocity field on the learned text-conditioned Rectified Flow to yield similar effects as Calssifier-Free Guidance.
(Z,, T) - a T) + (1-))0(2, (ULLL)
(7)
where a trades off the sample diversity and generation quality. When a 1, v reduces back to the original velocity field v (T). We provide analysis in Section
3 PRELIMINARY RESULTS: REFLOW IS THE KEY TO IMPROVE DISTILLATION
In this section, we conduct experiments with Stable Diffusion 1.4 to examine the effective- ness of the Rectified Flow framework and the reflow procedure. The goal of the experi- ments in this section is 1) examine whether straightforward distillation can be effective for learning one-step model from pre-trained large-scale T2I prbobility flow ODEs; 2) ex- amine whether text-conditioned rellow can en- hance the performance of distillation. Our ex- periment concludes that: Reflow significantly eases the learning process of distillation, and distillation after reflow successfully produces one-step model.
2-Rectiffed
Distill
30
Sibble Diffision
-Distill
Distil
20
10
10a
Interence Time (s)
2-Reciified Flow
+Distill
Figure 5: Left: The inference time and FID-5k on MS COCO 2017 of all the models. Model distilled from 2- Rectified Flow has a lower FID and smaller gap with the teacher. Right: The images generated from different models with the same random noise and text prompt. 2- Rectified Flow refines the coupling between noises and images, making it better teacher for distillation.
3.1 GENERAL EXPERIMENT SETTINGS
In this section, we use the pre-trained Stable Diffusion 1.4 provided in the official open-sourced repository] to initialize the weights, since otherwise the convergence is unbearably slow. In our experiment, we set DT to be a subset of text prompts from laion2B-en (74), pre-processed by the same filtering as SD. ODE
implemented as the pre-trained Stable Diffusion with 25-step DPM- Solver [49 and a fixed guidance scale of 6.0. We set the similarity loss ID(,,) for distillation to be the LPIPS loss [00]. The neural network structure for both reflow and distillation are kept to the SD U-Net. We use batch size of 32 and 8 A100 GPUs for training with AdamW optimizer [48). The choice of optimizer follows the default protoco in HuggingFace for fine-tuning SD.
3.2 DIRECT DISTILLATION FAILS, WHILE REFLOW + DISTILLATION SUCCEEDS
Experiment Protocol Our investigation starts from directly distilling the velocity field 81 Use of Stable Diffusion 1.4 with equation 6 without applying any reflow. To achieve the best em- pirical performance,
conduct grid search on learning rate and weight decay the limit of our computational resources. For all the models, we train them for 100, 000 steps. We generate 32 x 100, 000 = 3, 200, 000 pairs of (Xo, ] as the training set for distillation. We compute the Fréchet inception distance (FID) on 5, 000 captions from MS COCO 2017 following the evaluation protocol in 58 then we show the model with the lowest FID Figure 5 To align the training cost between direct distillation and reflow+distillation for fair comparison, we train 2- Rectified Flow, 22 for 50, 000 steps with the weights initialized from pre-trained SD, then perform

---

Published as a conference paper at ICLR 2024
Method
Inf. Time
FID-3k
CLIP
Method
Inf. Time
FID-33k
SD 1.4 step]
0.88s
22.8
0,315
(Pre) 2-RF (25 step)
0.88s
22.1
0.313
SD 770
2.9s
9.62
PD (1 [5
0.09s
37.2
0.275
(Pre) 2-RF (25 step)
0.88s
13.4
SD 4+sitllll
0.09s
40.9
0.255
SD 1.4+Distill
0.09s
34,6
(Pre) 2-RF (1 step)
0.09s
68.3
0.252
(Pre) 2-RF+Distill
0.09s
31.0
0.285
(Pre) 2-RF+Distill
0.09s
20.0
(a) MS COCO 2017
(b) MS COCO 2014
Table I: Comparison of (a) FID and CLIP score on MS COCO 2017 with 5, 000 images following the evalua-
tion setup in [58] and (b) FID on MS COCO 2014 with 30, 000 images following the evaluation setup in [3]]
As in 32 [23 the inference time is measured on NVIDIA A GPU with batch size of 1. 'Pre' is added
to distinguish the models from Table 2 "RF" refers to Rectified Flow; 'DD' refers to Progressive Distilla-
tion [5872] 72.. denotes that the numbers are measured by 332
distillation for another 50, 000 training steps continuing from the obtained u2. To distill from 2- Rectified Flow, we generate 32 50, 000 = 1, 600, 000 pairs of (Xo, ODE(v2] (Xo)) with 25-step Euler solver. The results are also shown in Figure 5 for comparison with direct distillation. The guidance scale a for 2-Rectified Flow is set to 1.5. For more details, please refer to Appendix.
Statlo Duffuekon
2-Hterimtdd blow.
Sontl Dufiiioo
Observation and Analysis We observe that, after 100, 000 training steps, all the models from direct dis- tillation converge. However, as shown in Figure [ it is difficult for the student model (SD+Distill) to imitate the teacher model (25-step SD), resulting in a huge gap in FID between SD and SD+Distill. On the right side of Figure [ with the same random noise, SD+Distill gen- erates image with substantial difference from the teacher SD. From the experiments, we conclude that: direct dis- tillation from SD is a tough learning problem for the one-step student model, which is hard to mitigate by simply tuning the hyperparameters. In contrast, Rectified Flow refines the coupling between the noise dis- tribution and the image distribution, and eases the learning process for the student model when dis- tillation. It can be inferred from two aspects: (1) The gap between 2-Rectified Flow+Distill and 2-Rectified Flow is much smaller than SD+Distill and SD. (2) On the right side of Figure the im- age generated from 2- 2-Rectified Flow+Distill shares great resemblance with the original generation, showing that it easier for the student to imitate. This illustrates that 2-Rectified Flow is better teacher model to distill one-step student model than the original SD.
Phooi Valuee
log S03
2RRosifeel
Figure 6: The straightening effect of re- flow. Left: the straightness S(Z) on differ- ent models. Right: trajectories of randomly sampled pixels following SD 1.4+DPM- Solver and 2-Rectified Flow.
N-
N=1
+Distill
Training Cost Because our 2-Rectified Flow+Distill is fine-tuned from the publicly available pre-trained SD, training only costs Rs 24.65 A100 GPU days, which is negligible com- pared with other large-scale text-to-image models.
2-Roctifled Flow Stable Diftusion
For reference, the training cost for SD 1.4 from scratch is 6250 A 110 GPU days 70]; StyleGAN-T 1792 A100 GPU days 773: GigaGAN is 4783
A100 GPU days [32]. A lower-bound estimation of training the one-step SD in Progressive Distillation
is 108.8 A GPU days (58). More details can be
Figure 7: Visual comparison between SD and 2- Rectified Flow. N is number of inference steps.
found in the Appendix.
Comparison on MS COCO As shown in (a), on MS COCO 2017-5k, (Pre) 2-Rectified Flow can generate realistic images that yield similar FID with SD 1.4 (+DPMSolver (50) using 25 steps (22.1 <> 22.8). Within 0.09s, (Pre) 2-Rectified Flow+Distill gets an FID of 31.0, surpassing the previous best one-step SD model (FID=37.2) distilled from Progressive Distillation (58) with much
less training cost (the numbers for Progressive Distillation are measured from Figure 10 in [58] since the model is not publicly available). On MS COCO 2014-30k, (Pre) 2-Rectified Flow+Distill has noticeable advantage (FID=20.0) compared with direct distillation SD 1.4+Distill (FID=34.6) even when (Pre) 2- 2-Rectified Flow has worse performance than the original SD due to insufficient training, indicating the effectiveness of the reflow operation.
7

---

Published as a conference paper at ICLR 2024
Method
Inf. Time
FID-3K
CLIP
301
285
SD 1.5 (25 step) 70
0.88
20.1
0.318
A4
255
Puni-
Pas-30
6:41
150M
-18
10:31
8.30
2-RF (25 step)
0.88
21.5
0.315
AR
25.00
PD-SD (1 step)
0.09
37.2
0.275
255
1336
18
12:34
2-RF (1 step)
0.09
47.0
0.271
131
155
DALLE
530
00.35
InstaFlow-0.98
0.09
23,4
0.304
DVT
355
Io Iog
0.8
38
7.27
PD-SD (2 (tee) 558
0.13
26.0
0.297
GAN
355
512
LAFITE 100g
Mac
0.02s
13.
0.38
7S4
26.34
738
PD-SD (4 step)
0.21
26.4
0.300
GAN
512
GigutA
0.15s
18
9.0
2-RF' (2 step)
0.13
31.3
0.296
912
2.4
0130
InstaFlow-1.
0.12
22.4
0.309
513
0.00
6.58
13.05
1124
1.30
1183
(a) MS COCO 2017
(b) MS COCO 2014
Table 2: Comparison of (a) FID and CLIP score MS COCO 2017 with 5, 000 images following the evalua-
tion setup in (58) and (b) FID on MS COCO 2014 with 30, 000 images following the evaluation setup in [2]
'RF' refers to Rectified Flow; 'PD' refers Progressive Distillation [58) 72): 'AR' refers to Autoregressive, denotes that the numbers are measured by (32]
Straightening Effects of Reflow We empirically examine the properties of reflow in text-to-image generation. To quantitatively measure the straightness, we use the deviation of the velocity along
the trajectory following 45 1.
that is, S(Z) = fr E [I] (Z Zo) v(z,, t) [| dt. A smaller S(Z) means straighter trajectories, and when the ODE trajectories are all totally straight, S(Z) traeeerotiiiie
0. In Figure 6 reflow decreases the estimated S(Z), validating the straightening effect of reflow. Moreover, the pixels in SD travel in curved trajectories, while 2-Rectified Flow has much straighter trajectories. In Figure 7 qualitatively, since SD is curved, one-step generation leads to meaningless noises and SD+Distill fails. Thanks to reflow, one-step generation with 2-Rectified Flow shows recognizable images and distillation from it succeeds.
4 INSTAFLOW: SCALING UP FOR BETTER ONE-STEP GENERATION
Our preliminary results using SD 1.4 highlight the benefits of incorporating the reflow procedure in distilling one-step diffusion-based models. However, considering that the training process only consumes 24.65 A100 GPU days, there is a potential for further performance enhancement through scaling up. To this end, we extend the training duration with larger batch size, totaling 199 A100 GPU days. As a result, we achieve InstaFlow, the first one-step SD model capable of generating high-quality images with intricate details in just 0.09 second. Notably, this performance is on par with StyleGAN-T 73. one of the state-of-the-art GANs in the field.
DistaFlow-0.99. We switch to Stable Diffusion 1 1.5, and keep the same D7 as in Section C The ODE solver sticks to 25-step DPMSolver [49] for ODE |.o] Guidance scale for SD is slightly decreased to 5.0 because larger guidance scale makes the images generated from 2-Rectified Flow over-saturated. We still generate 1, 600, 000 pairs of data for reflow and distillation, respectively. We apply gradient accumulation to expand the batch size. We spend 75.2 A100 GPU days for reflow get 2-Rectified Flow, then another 108 A100 GPU days for distillation to get 2- Reettified Flow+Distill. The guidance scale a for 2-Rectified Flow is set to 1.5 during distillation. We name the distilled model InstaFlow- 0.9B since U-Net contains 0.9B parameters.
InstaFlow- 1.7B Expanding the model size a key step building modern foundation mod- els (70 62, 6 5 16. To this end, we stack two U-Nets in series, then remove unnecessary modules after a thorough ablation study (see Appendix for details). This gives us a large neural network, termed Stacked U-Net, with 1.7B parameters and an inference time of 0.12 second. Starting from 2- P-Rectified Flow obtained in InstaFlow-O.9. -.9B, we spend 39.6 A100 GPU days for distillation with Stacked U-Net. More training details of both models can be found in the Appendix.
Comparison with State-of-the-Arts on MS COCO We follow the experiment configuration in Seciton C In Table 2 (a), our InstaFlow- -0.9B gets FID-5k of 23.4 with an inference time of 0.09s, which is significantly lower than the previous state-of-the-art Progressive Distillation-SD (1 step. FID==77.2) 2) with similar distillation cost (108 108.8 1100 GPU days). The empirical result indicates that reflow helps improve the coupling between noises and images, and 2-Rectified Flow an easier teacher model distill from. By increasing the model size, InstaFlow-1.77 leads to lower FID-5k of 22.4 with an inference time of 0.12s. On MS COCO 2014, our InstaFlow-O.99 obtains an FID-30k of 13. 10 within 0.09s, surpassing StyleGAN-T 73 (13.90 in 0.1s).
8

---

Published as a conference paper at ICLR 2024
orggial
red how the
Ie haney soow
oll puinting
dee siverg of saaiena fentagy sowes slow,, sarkk, Night denvlee
Cogg oboswe csilggagcor
Figure 8: One-step generation from InstaFlow-O.P90 Left: With the same random noise, the pose and lighting are preserved across different text prompts. Right: Interpolation in the latent space of InstaFlow-O.99.
Few-step Generation with 2-Rectified Flow 2-Rectified Flow has straighter tra- jectories, which gives it the capacity generate with extremely few inference steps. We compare 2-Rectified Flow with SD 1.5-DPM Solver [50] on MS COCO 2017. The inference steps are set {1,2, 4,8) 8). Figure (A) clearly shows the advantage of 2-Rectified Flow when the number of inference steps S Guidance Scale a is widely known that guidance scale a is a important syper-parameter when using Stable Diffusion 126 70). Here, we investigate the influence of the guidance scale for 2-Rectified Flow, which has straighter ODE trajectories. {1.0, 1.5, 2.0, 2.5, 3.5, 4.0], which raises FID-5k and CLIP score MS COCO 2017 at the same time. The former metric indicates degradation image quality and the latter metric indicates enhancement in semantic alignment.
50 Saler
-0h LEN
Humber inturce
2.5 0316 3-23
1.38
(A)
(B)
Figure 9: (A) Comparison between SD 1.5-DPM Solver and 2-Rectified Flow (with standard Euler solver) in few-step in- ference. (B) The trade-off curve of applying different a as the guidance scale for 2- 2-Rectified Flow.
Figure 9 (B), a increases from
Fast Preview with One-Step InstaFlow potential use case of InstaFlow is to serve fast pre- viewer can accelerate the low-resolution filtering process and provide the user more generation possibilitics un- der the same computational budget. Then a powerful post- processing model the quality and increase the resolution. We verify the idea with SDXL-Refiner [62]. a recent model that can refine generated images. The one-step InstaFlows generate 512 x 512 images in 0.1s, then these images are interpolated 1024 and refined by SDXL-Refiner to get high-resolution images. Several examples shown in Figure
as previewers.
Oue-Stpp
improve
Step SDXL Retiner
Figure The images generated our one-step model can refined by SDXL-Refiner [62] to generate user- treferreerr high-resolution images higher efficiency.
InstaFlow-0199
S13 D.5-DPM Solver
LIMITATIONS AND CONCLUSIONS
In this paper, we introduce InstaFlow. state-of-the-art one- step text-to-image generator, which is derived from novel text-conditioned Rectified Flow pipeline with pure supervised learning. Although it may encounter challenges with complex compositions in the text prompts (see Figure TT for example), further training with longer dura- tion and larger datasets is likely mitigate them. InstaFlow significantly closes the gap between continuous time diffusion models and one-step generative models, inspiring algorithmic innovations and benefiting downstream tasks like 3D generation.
dog Make band volk soler Figure 11: One of the failure cases.

---

Published as a conference paper at ICLR 2024
SOCIETAL IMPACT
This work presents methodology for accelerating multi-step large-scale text-to-image diffusion models to one-step generators. On positive note, we believe that the efficiency of these one-step models can lead to energy conservation and environmental benefits, given the extensive utilization of such generative models. Conversely, faster generative models, when manipulated by bad actors, also simplify and speed up the creation of harmful information and fake news. While our work focuses on the scientific insights, these ultra-fast powerful generative models call for advanced techniques through research to ensure their alignment with human values and public interests.
REFERENCES
[1] Michael Albergo, Nicholas M Boffi, and Eric Vanden- Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023,
[2] Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochas- tie interpolants. In The Eleventh International Conference on Learning Representations, 2022.
[3] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song. Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, al. ediffi: Text-to-image diffusion models with an ensemble of expert denoisers. arviv preprint wrXiv:2211.013222 2022.
[4] Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference on Learning Representations, 2021.
[5] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks foundation models. arXiv preprint wwxXiv:2108.07255 2021.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
[7] Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming- Hsuan Yang, Kevin Murphy, William T Freeman, Michael Rubinstein, al. Muse: Text-to- image generation via masked generative transformers. In International conference ma- chine learning. PMLR, 2023.
[8] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun David Luan, and Ilya Sutskever. Generative pretraining from pixels. International conference on machine learn- ing, pp. 1691-1703. PMLR, 2020.
Ricky TQ Chen, Jens Behrmann, David Duvenaud, and Jömn-eerii Jacobsen. Residual flows for invertible generative modeling. Advances in Neural Information Processing Sys- tems, 32, 2019.
[10] Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degrraaation beyond linear dif- fusions: non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning, pp. 4462-4484, PMLR, 2023.
[11] Katherine Crowson, Stella Biderman, Daniel Kornis, Dashiell Stander, Eric Hallahan, Louis Castricato, and Edward Raff. Vqgan-clip: Open domain image generation and editing with natural language guidance. In European Conference on Computer Vision, pp. 88-105. Springer, 2022.
[12] Giamnis Daras, Mauricio Delbracio, Hossein Talebi, Alex Dimakis, and Peyman Milanfar. Soft diffusion: Score matching with general corruptions. Transactions on Machine Learning Research, 2023.
10

---

Published as a conference paper at ICLR 2024
[13] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-877 8794, 2021.
[14] Ming Ding, Zhuoyi Yang. Wenyi Hong, Wendi Zheng. Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, et al. Cogview: Mastering text-to-image generation via transformers. Advances in Neural Information Processing Systems,
2021.
[15] Ming Ding, Wendi Zheng, Wenyi Hong, and Jie Tang. Cogview2: Faster and better text-to- image generation via hierarchical transformers. Advances in Neural Information Processing Systems, 35:16890-16902, 2022.
[16] Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong. Tianhe Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint arXiv:2303.03.777 2023.
[17] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
[18] Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv Taig- man. aaceee- Scene-based text-to-image generation with human priors. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Pro- ceedings, Part XV, pp. 89-106. Springer, 2022.
[19] Ian Goodfellow, Jean Pouget-Abadic, Mchdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communica- tions of the ACM, 63(11):139) 144, 2020.
[20] Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data- free distillation of denoising diffusion models with bootstrapping. In ICML 2023 Workshop on Structured Probabilistic Inference {'&} Generative Modeling, 2023.
[21] Xizewen Han, Huangjie Zheng, and Mingyuan Zhou. Card: Classification and regression diffusion models. Advances in Neural Information Processing Systems, 35:18100-18115, 2022.
[22] Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris. Ganspace: Discov- ering interpretable gan controls. Advances in neural information processing systems, 33: 9841-9850, 2020.
[23] Eric Heitz, Laurent Belcour, and Thomas Chambon. Iterative -leendiendnn:: A minimalist deterministic diffusion model. In ACM SIGGRAPH 2023 Conference Proceedings, SIG- GRAPH '23, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400701597. doi: 10. N0.1445355884223551555 URL
3588432 3591540
[24] Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-or. Prompt-to-promp image editing with cross-attention control. In The Eleventh International Conference on Learning Representations, 2022.
[25] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint 1503.02553.055 2015.
[26] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NewrIPS 202/ Work- shop on Deep Generative Models and Downstream Applications, 2021.
[27] Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. In Advances in Neural Information Processing Sys- tems.
[28] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Ad- vances in Neural Information Processing Systems, 33:6840-6851, 2020.
11

---

Published as a conference paper at ICLR 2024
[29] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303, 2022.
[30] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling, Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp- 8867-8887. PMLR, 2022.
[31] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffu- sion for high resolution images. In International conference on machine learning. PMLR, 2023.
[32] Minguk Kang, Jun- Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 10124-10134, 2023.
[33] Tero Karras, Mïik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems.
[34] Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6007- 6017, 2023.
[35] Ivan Kobyzev, Simon JD Prince, and Marcus A Brubaker. Normalizing flows: An intro- duction and review of current methods. IEEE transactions on pattern analysis and machine intelligence, (3111):39643377) 2020.
[36] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: versa- tile diffusion model for audio synthesis. In International Conference on Learning Represen- tations.
[37] Wenbo Li, Pengchuan Zhang, Lei Zhang, Qiuyuan Huang, Xiaodong He, Siwei Lyu, and Jianfeng Gao. Object-driven text-to-image synthesis via adversarial training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 12182, 2019.
[38] Yanyu Li, Huan Wang, Qing Jin, Ju Hu, Pavlo Chemerys, Yun Fu, Yanzhi Wang. Sergey Tulyakov, and Jian Ren. Snapfusion: Text-to-image diffusion model on mobile devices within two seconds. arXiv preprint Xrxv:22306. 03000000000 2023.
[39] Tsung-Y Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
[40] Yaron Lipman, Ricky TQ Chen, Heli Ben- Ham-H Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2022.
[41] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations.
[42] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models manifolds. In International Conference Learning Representations, 2021.
[43] Qiang Liu. Rectified flow: marginal preserving approach to optimal transport. arXiv preprint arXiv:2209. 14577, 2022.
[44] Xingchao Liu, Chengyue Gong, Lemeng Wu, Shujian Zhang, Hao Su, and Qiang Liu. Fuse- dream: Training free text-to-image generation with improved clipt gan space optimization. arXiv preprint 122.0:21 2021.
12

---

Published as a conference paper at ICLR 2024
[45] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2022.
[46] Xingchao Liu, Lemeng Wu, Mao Ye, et al. Learning diffusion bridges constrained do- mains. In The Eleventh International Conference on Learning Representations, 2023.
[47] Xingchao Liu, Lemeng Wu, Shujian Zhang, Chengyue Gong, Wei Ping, and Qiang Liu. Flowgrad: Controlling the output of generative odes with gradients. Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 24335-243444 2023.
[48] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations.
[49] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems.
[50] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm- solver+4: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095, 2022.
[51] Haoming Lu, Hazarapet Tunanyan, Kai Wang, Shant Navasardyan, Zhangyang Wang, and Humphrey Shi. Specialist diffusion: Plug-and-play sample-efficient fine-tuning of text-to- image diffusion models learn any unseen style. In Proceedings of the IEFF/CVF Confer- ence Computer Vision and Pattern Recognition, pp. 14267-14276, 2023.
[52] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint 0tti:22111002888 2021.
[53] Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 2837-2845, 2021.
[54] Shitong Luo and Wei Hu. Score-based point cloud denoising. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4583-4592, 2021.
[55] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion- baated generative models for protein structures. In Advances in Neural Information Processing Systems.
[56] Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthesizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04378, 2023.
[57] Chenlin Meng, Yutong He, Yang Song. Jiaming Song. Jiajun Wu, Jun- Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2021.
[58] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 14297-14306. 2023.
[59] Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Megrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. In International Conference on Machine Learning, pp. 16784-16804. PMLR, 2022.
[60] George Papamakarios, Eric Nalisnick. Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1):617-22600 2021.
13

---

Published as a conference paper at ICLR 2024
[61] Or Patashnik, Zongze Wu, Eli Shechtman, Daniel Cohen-Or, and Dani Lischinski. Styleclip: Text-driven manipulation of stylegan imagery. In Proceedings of the IEEE/CVF International Conference Computer Vision. pp. 2085-20920 2021.
[62] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis, 2023.
[63] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakom. Diffusion autoencoders: Toward meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Putlern Recognition, pp. 10619-10629, 2022.
[64] Can Qin, Shu Zhang, Ning Yu, Yihao Feng, Xinyi Yang, Yingbo Zhou, Huan Wang, Juan Car- los Niebles, Caiming Xiong, Silvio Savarese, et al. Unicontrol: A unified diffusion model for controllable visual generation in the wild. arXiv preprint arXiv:2305. 147, 2023.
[65] Yiming Qin, Huangjie Zheng, Jiangchao Yao, Mingyuan Zhou, and Ya Zhang. Class- balancing diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pullern Recognition, pp. 18434-18443, 2023.
[66] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot texl-to-image generation. In International Conference Machine Learning, pp. 8821-8831 PMLR, 2021.
[67] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06123. 2022.
[68] Scott Reed, Zeynep Akata, Xinchen Yan, Lajanugen Logeswaran, Bernt Schiele, and Honglak Lee. Generative adversarial text to image synthesis. International conference on machine learning, pp. 1060-1069. PMLR, 2016.
[69] Scott E Reed, Zeynep Akata, Santosh Mohan, Samuel Tenka, Bert Schiele, and Honglak Lee. Learning what and where to draw. Advances in neural information processing systems, 29, 2016.
[70] Robin Rombach, Andreas Blattmann, Dominík Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models, 2021.
[71] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily I. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494, 2022.
[72] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations.
[73] Axel Sauer, Tero Karras, Samuli Laine, Andreas Geiger, and Timo Aila. Stylegan-t: Unlock- ing the power of gans for fast large-scale text-to-image synthesis. 2023.
[74] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wight- man, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion-5b: An open large-scale dataset for training next generation image-text models. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track.
[75] Yujun Shen and Bolei Zhou. Closed-fon form factorization of latent semantics in gans. In Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1532- 1540, 2021.
[76] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256-22666 2265. PMLR, 2015.
14

---

Published as a conference paper at ICLR 2024
[77] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations.
[78] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[79] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33:122388 2020.
[80] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations.
[81] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34: 1415-1428, 2021.
[82] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
[83] Ming Tao, Hao Tang, Fei Wu, Xiao-Yuan Jing, Bing- Mun Bao, and Changsheng Xu. Df-gan: A simple and effective baseline for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16515- 16525, 2022.
[84] Tengfei Wang, Yong Zhang. Yanbo Fan, Jue Wang, and Qifeng Chen. High-tidelity gan inver- sion for image attribute editing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11379- 1388, 2022.
[85] Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-gan: Training gans with diffusion. In The Eleventh International Conference on Learning Representations, 2022.
[86] Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, et al. Diffusion Diffusion-based molecule gener- ation with informative prior bridges. In Advances Neural Information Processing Systems.
[87] Lemeng Wu, Dilin Wang, Chengyue Gong, Xingchao Liu, Yunyang Xiong, Rakesh Ranjan, Raghuraman Krishnamoorthi, Vikas Chandra, and Qiang Liu. Fast point cloud generation with straight flows In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9445-9454, 2023.
[88] Zongze Wu, Dani Lischinski, and Eli Shechtman. Stylespace analysis: Disentangled controls for stylegan image generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12863-12872, 2021.
[89] Weihao Xia, Yulun Zhang, Yujiu Yang, Jing-Hao Xue, Bolei Zhou, and Ming-Hsuad Yang- Gan inversion: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(3):3121-3138, 2022.
[90] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In International Conference on Learning Representations, 2021.
[91] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Confer- ence on Learning Representations.
[92] Xingqian Xu, Zhangyang Wang, Eric Zhang, Kai Wang, and Humphrey Shi. Versa- tile diffusion: Text, images and variations all in one diffusion model. arXiv preprint orXiv:2211.08333 2022.
[93] Yanwu Xu, Yang Zhao. Zhisheng Xiao, and Tingbo Hou. Ufogen: You forward once large scale text-to-image generation via diffusion gans. arXiv preprint arXiv:2311.09257, 2023.
[94] Mao Ye, Lemeng Wu, and Qiang Liu. First hitting diffusion models for generating manifold, graph and categorical data. In Advances in Neural Information Processing Systems, 2022.
15

---

Published as a conference paper at ICLR 2024
[95] Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. arXiv preprint arXiv:3311.1882. 2023.
[96] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Va- sudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al. Scaling autoregressive models for content-rich text-to-image generation. Transactions on Machine Learning Re- search.
[97] Han Zhang, Tao Xu, Hongsheng Li, Shaoting Zhang, Xiaogang Wang, Xiaolei Huang, and Dimitris N Metaxas. Stackgan: Text to photo-realistic image synthesis with stacked genera- tive adversarial networks. In Proceedings of the IEEE international conference computer vision, pp. 5907-5915, 2017.
[98] Han Zhang, Jing Yu Koh, Jason Baldridge, Honglak Lee, and Yinfei Yang. Cross-modal contrastive learning for text-to-image generation. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 833-842, 2021.
[99] Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint wxXiv:2302.05533 2023.
[100] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
[101] Shu Zhang, Xinyi Yang, Yihao Feng. Can Qin, Chia-Chih Chen, Ning Yu. Zeyuan Chen, Huan Wang, Silvio Savarese, Stefano Ermon, Caiming Xiong, and Ran Xu. Hive: Harnessing human feedback for instructional visual editing. arXiv preprint urXiv:2303.09616 2023.
[102] Mingyuan Zhou, Tianqi Chen, Zhendong Wang, and Huangjie Zheng. Beta diffusion. arXiv preprint arXiv:2309.0788 2023.
[103] Yufan Zhou, Ruiyi Zhang, Changyou Chen, Chunyuan Li, Chris Tensmeyer, Tong Yu, Jiuxi- ang Gu, Jinhui Xu, and Tong Sun. Towards language-free free training for text-to-image genera- tion. pp. 17907-17917, 2022.
16

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
Diffusion Models and Flow-based Models
31 have achieved unprecedented results in various generative modeling tasks, including image/ video generation (27 [oI 86; 94
Diffusion models 76
28 80 81 79
78
33
13
21 65 102
51
92
12
71,, audio generation [36], point cloud gener-
ation (53) 54
46 87,
biological generation 19
53 86 30), etc.. Most of the works are based
on stochastic differential equations (SDEs), and researchers have explored techniques to transform them into cently, 45 40
preserving probability flow ordinary differential equations (ODEs) [8) Re-
I 23 propose to directly learn probability flow ODEs by constructing linear or non-linear interpolations between two distributions. These ODEs obtain comparable performance as diffusion models, but require much fewer inference steps. Among these approaches, Rectified Flow [45; 43 introduces a special reflow procedure which enhances the coupling between distribu- tions and squeezes the generative ODE to one-step generation. However, the effectiveness of reflow has only been examined on small datasets like CIFARIO, thus raising questions about its suitability on large-scale models and big data. In this paper, we demonstrate that the Rectified Flow pipeline can indeed enable high-quality one-step generation in large-scale text-to-image diffusion models, hence brings ultra-fast T2I foundation models with pure supervised learning.
Large-Scale Text-to-Image Generation Early research on text-to-imase -intge generation focused on scall-scale datasets, such as flowers and birds 668 69
Later, the field shifted its attention to more complex scenarios, particularly in the MS COCO dataset [39,, leading to advancements in training and generation 83 98
37.. DALL-E (66) was the pioneering transformer-basee model that showcased the amazing zero-shot text-to-image generation capabilities by scaling up the network size and the dataset scale. Subsequently, a series of new methods emerged, including autoregressive models [14
15
18 96.. GAN inversion [T] (4). GAN-based approaches [103), and diffusion mod-
els [59
31. Among them, Stable Diffusion is an open-source text-to-image generator based on latent diffusion models (70). It is trained on the LAION 5B dataset (74] and achieves the state-of-the-ar generalization ability. Additionally, GAN-based models like StyleGAN-T (73 and GigaGAN (32) are trained with adversarial loss to generate high-quality images rapidly. Our work provides a novel approach to yield ultra-fast, one-step, large-scale generative models without the delicate adversarial training.
67
64
Acceleration of Diffusion Models are known to be slow during inference due to the requirement of multiple iterations to reach the final result. To accelerate inference, there are two categories of algorithms. The first kind focuses on
Despite the impressive generation quality, diffusion models
fast post-hoc samplers B3 49 50 77, I0]. These fast samplers can reduce the number of inference steps for pre-trained diffusion models to 20-50 steps. However, relying solely on inference to boost performance has its limitations, necessitating improvements to the model itself (85] 90 33). Distillation (25) has been applied pre-trained diffusion models [52]. squeezing the number of in- ference steps to below 10. Progressive distillation [72) is a specially tailored distillation procedure for diffusion models, and has successfully produced 2/4-step Stable Diffusion
(58). Consistency models [82] are new family of generative models that naturally operate in a one-step manner. There are several works concurrent with InstaFlow for accelerating large-scale text-to-image mod- els: BOOT (20 designs a data-free distillation pipeline for tra-trained diffusion models with a bootstrapping method; Latent Consistency Model (56) adopts consistency distillation with a skip- ping scheme; UFOGen (93] combines diffusion models with GAN for high-quality distillation; Yin et al. (95] proposes Distribution Matching Distillation to align the distribution generated from the one-step model with the teacher Stable Diffusion. Instead of employing sophisticated distillation or GAN loss, InstaFlow uses Rectified Flow 45 + and its unique reflow procedure to straighten the ODE trajectories. With simple supervised learning on least-squares problem, it refines the cou- pling between the noise distribution and the image distribution, thereby improving the performance of direct distillation.
B NEURAL NETWORK STRUCTURE
The whole pipeline of our text-to-image generative model consists of three parts: the text encoder, the generative model in the latent space, and the decoder. We use the same text encoder and decoder as Stable Diffusion: the text encoder is adopted from CLIP ViT-L/14 and the latent decoder is
17

---

Published as a conference paper at ICLR 2024
In Block
Out Block
Downsample Blocks
Upsample Blocks
4 x 64 64 x64
4 x 64 X64
Input
Output
U-Net (Inference time=0.04s)
4 x 64 x64
4 x 64 x6
Input
Output
(b) Naïvely Concatenated U-Net (Inference time=0.08s)
xx 64 x64
4 x 64 X64
Input
Output
Our Stacked U-Net (Inference time=0.07s)
Figure 12: Different neural network structures for distillation and their inference time. The blocks with same colors can share weights.
18

---

Published as a conference paper at ICLR 2024
adopted from a pre-trained auto-encoder with downsampling factor of 8. During training, the parameters in the text encoder and the latent decoder are frozen. On average, to generate image on NVIDIA A100 GPU with batch size of 1, text encoding takes 0.01s and latent decoding takes 0.04s.
By default, the generative model in the latent space is U-Net structure. For reflow, we do not change any of the structure, but just fine-tune the model. For distillation, we tested three network structures, as shown in Figure 12] The first structure is the original U-Net structure in SD. The second structure is obtained by directly concatenating two U-Nets with shared parameters. We found that the second structure significantly decrease the distillation loss and improve the quality of the generated images after distillation, but it doubles the computational time.
To reduce the computational time, we tested family of networks structures by deleting different blocks in the second structure. By this, we can examine the importance of different blocks in this concatenated network in distillation, remove the unnecessary ones and thus further decrease infer- ence time. We conducted series of ablation studies, including:
1. Remove 'Downsample Blocks
(the green blocks on the left)'
2. Remove 'Upsample Blocks I (the yellow blocks on the left)"
3. Remove 'In+Out Block' the middle (the blue and purple blocks in the middle).
4. Remove 'Donnsample Blocks 2 (the green blocks on the right)"
5. Remove 'Upsample blocks (the yellow blocks on the right)"
The only one that would not hurt performance is Structure 3, and it gives us a 7.7% reduction in inference time (0.13- 0.12 0.0769 This third structure, Stacked U-Net, illustrated in Figure 12 (c).
0.13
C ADDITIONAL DETAILS AND RESULTS ON THE PRELIMINARY EXPERIMENTS
General Experiment Settings this section, we use the pre-trained Stable Diffusion 1.4 provided in the official open- nourcedd repository to initialize the weights, since otherwise the convergence is unbearably slow.
In our experiment, we set DJ to be a subset of text prompts from laion2B-en [74, pre-processed by the same filtering as SD. UDE[Usp] is implemented as the pre-trained Stable Diffusion with 25-step DPMSolver [49] and fixed guidance scale of 6.0. We set the similarity loss D(-,) .) for distillation to be the LPIPS loss [00]. The neural network structure for both reflow and distillation are kept to the SD U-Net. We use batch size of 32 and 8 A100 GPUs for training with AdamW optimizer [88]. Our training script is based on the official fine-tuning script provided by HuggingFace and the choice of optimizer follows the default protocol. We exponential moving average with factor of 0.9999, following the default configuration. We clip the gradient to reach a maximal gradient norm of 1. We warm-up the training process for 1,000 steps both reflow and distillation. BF16 format is adopted during training to save GPU memory. To compute the LPIPS loss, we used its official 0.4 version and its model based on AlexNet. The learning rate for reflow is 10-6
We measure the inference time of our models on a server with NVIDIA A100 GPU, and batch size of 1. We use PyTorch 2.0.1 and Hugging Face Diffusers 0.19.3. For fair comparison, we use the inference time of standard SD on our computational platform for Progressive Distillation-S5 as their model is not available publicly. The inference time contains the text encoder and the latent decoder, but does NOT contain NSFW detector.
19

---

Published as a conference paper at ICLR 2024
FID
weight decay
10-1
10-2
10-3
learning rate
10
44.04
45.90
44.03
10
137.05
134.83
139.31
10
~2297
297
~297
Table 3: FID of different distilled SD models measured with 5000 images MS COC02017.
C.I ADDITIONAL DETAILS AND RESULTS OF DIRECT DISTILLATION
Grid Search To achieve the best empirical performance, we conduct grid search on leam- ing rate and weight decay to the limit of our computational resources. Particularly, the learn- ing rates are selected from {10-5, 10-6, 10-7} and the weight decay coefficients are selected from {10-1 10-2 2,- 10-3). For all the 9 models, we train them for 100, 000 steps. We generate 32 x 100, 000 = 3, 200, 000 pairs of (Xo, ODE (eso)(Xo))) as the training set for distillation.
Additional Results shown in Figure 18 formance, we found that learning rate of 2 10-1 leads to unstable training and NaN errors. A small learning rate, like 10-6 and 10-7, results in slow convergence and blurry generation even after training 100, 000 steps.
We provide additional results on direct distillation of Stable Diffusion 1.4, 20 21 and Table Although increasing the learning rate boosts the per-
19
C.2 ADDITIONAL QUANTITATIVE COMPARISON
We provide additional quantitative results with parameter-sharine Stacked U-Net multiple reflow Table [] According to equation 5] the reflow procedure can repeated for multiple times. We repeat reflow for one more time to get 3-Rectified Flow (va), which is initialized from 2-Reciified Flow (12). 3-Rectified Flow is trained to minimize equation 5 for 50, 000 steps. Then we get its distilled version by generating 1, 600, 000 new pairs of (Xo, ODE[cy] (Xy)) and distill for another 50, 000 steps. We found that to stabilize the training process of R-ectifiedd Flow and its distillation, have to decrease the learning from 10- 10-?
C.3 ESTIMATION OF TRAINING COST
Measured on platform, when training with batch size of 4 and U-Net, one A100 GPU day can process 100, 000 iterations using L2 loss, 86, 000 iterations using LPIPS loss; when generating pairs with batch size of 16, one A100 GPU day can generate 200, 000 data pairs. We compute the computational cost according to this.
Estimated Training Cost of (Pre) 2-Rectified Flow+Distill (U-Net): 3, 200, 000/200, 000 (Data Generation) + 33/44x x 50, 000/ 000/100, 000 (Reflow) + 32/4 x 50, 000 000 (Distillation) R3 24.65 A100 GPU days.
Estimated Training Cost of Progressive Distillation: We refer to Appendix C.2.1 (LAION- 5B 512 x 512) of [58] and estimate the training cost. PD starts from 512 steps, and progressively applies distillation to step with batch size of 512. Quoting the statement 'For stage-two, we train the model with 2000-5000 gradient updates except when the sampling step equals to 1,2, or 4, where we train for 10000-50000 gradient updates', lower-bound estimation of gradient updates would be 2000 (512 to 256) 2000 (256 to 128) 2000 (128 to 64) + 2000 (64 to 32) 2000 (32 to 16) 5000 (16 to 8)+ + 10000 (8 to 4) + 10000 (4 to 2) + 50000 (2 to 1) = 85,000 iterations. Therefore, one-step 108.8 A100 GPU days. Note that ignored the computational cost of stage 1 of PD and steps DDIM with teacher' during PD, meaning that the real training cost is higher than 108.8 A100 GPU days.
PD at least requires 512/4 x 85000/ 100000
20

---

Published as a conference paper at ICLR 2024
Method
Inference Time
FID-5k
CLIP
SD 1.4-DPM Solver (25 step) 70 49
0.88s
22.8
0.315
(Pre) 2-Rectified Flow (25 step)
0.88s
22.1
0.313
(Pre) 3-Rectified Flow (25 step)
0.88s
23.6
0.309
Progressive Distillation-SD (1 step) 58
0.09s
37.2
0.275
SD 1.4+Distill (U-Net)
0.09s
40.9
0.255
(Pre) 2-Rectified Flow (1 step)
0.09s
68.3
0.252
(Pre) 2-Rectified Flow+Distill (U-Net)
0.09s
31.0
0.285
(Pre) 3-Rectified Flow (1 step)
0.09s
37.0
0.270
(Pre) 3-Rectified Flow+Distill (U-Net)
0.09s
29.3
0.283
Progressive Distillation-SD (2 step)
0.13s
26.0
0.297
Progressive Distillation-SD (4 step)
0.21s
26.4
0,300
SD 1.4+Distill (Stacked U-Net)
0.12s
52.0
0.269
(Pre) 2-Rectified Flow+Distill (Stacked U-Net)
0.12s
24.6
0.306
(Pre) 3-Rectified Flow+Distill (Stacked U-Net)
0.12s
26.3
0.307
Table 4: Comparison of FLD on MS COCO 2017 following the evaluation setup in [58] As in 32 77] the infer- ence time is measured on NVIDIA A100 GPU, with a batch size of FyTerch 2.0. and Huggingten ang Diffusers 0. 0.19. 3. 2-Rectified Distil outperforms Progressive Distillation within the infer- ence time using much less training cost. The numbers for Progressive Distillation measured from Figure 10 58
'Pre' is added distinguish the models from Table
Method
Inference Time
# Param.
FID-30k
SD [70]
2.9s
0,9B
9.62
(Pre) 2-Rectified Flow (25 step)
0.88s
0.9B
13.4
SD 1.4+Distill (U-Net)
0.09s
0,9B
34.6
(Pre) 2-Rectified Flow+Distill (U-Net)
0.09s
0.9B
20.0
SD 1.4+Distill (Stacked U-Net)
0.12s
0.9B
41.5
(Pre) 2-Rectified FlowiDistill (Stacked U-Net)
0.12s
0.9B
13.7
Table 5: Comparison of FID MS COCO 2014 with 30, 000 images. Note that the models distilled after
reflow has noticeable advantage compared with direct distillation when (Pre) 2-Rectified Flow has worse performance than the original SD due insufficient training. denotes the numbers are measured hy k2]
"Pre" is added to distinguish the models from Table As in StyleGAN-T 73 and GigaGAN 32 our generated images are downsampled to 256 x 256 before computing FID.
D ADDITIONAL TRAINING DETAILS ON INSTAFLOW
Implementation Details and Training Pipeline for InstaFlow-0.9B We switch to Stable Diffu- sion 1.5, and keep the same Dy in Section The ODE solver sticks 25-step DPMSolver [49 for ODE[Ug>]. Guidance scale is slightly decreased to 5.0 because larger guidance scale makes the im- ages generated from 2-Rectified Flow over-saturated. Since distilling from Rectifiied Flow yields satisfying results, 3-Rectifed Flow is not trained. We still generate 1, 600, 000 pairs of data for reflow and distillation, respectively. To expand the batch size be larger than 8 32, gradi- ent accumulation is applied. The overall training pipeline for 2- 2-Rectified Flow+Distill (U-Net) is summarized as follows:
1. Reflow (Stage 1): We train the model using the reflow objective equation with batch size of 64 for 70,000 iterations. The model is initialized from the pre-trained SD 1.5
weights. (11.2 A100 GPU days)
2. Reflow (Stage 2): We continue to train the model using the reflow objective equation with an increased batch size of 1024 for 25,000 iterations. The final model is 2-Rectified Flow. (64 A100 GPU days)
3. Distill (Stage 1): Starting from the Rectifiee Flow checkpoint, we fix the time t = 0 for the neural network, and fine-tune it using the distillation objective equation 6 with batch size of 1024 for 21,500 iterations. The guidance scale a of the teacher model, 2-Rectified Flow, is set 1.5 and the similarity loss D is 1.2 loss. (54.4 A 1100 GPU days)
21

---

Published as a conference paper at ICLR 2024
4. Distill (Stage 2): We switch the similarity loss IB to LPIPS loss, then continue to train the model using the distillation objective equation 6 and a batch size of 1024 for another 18,000 iterations. The final model is 2-Rectified Flow+Distill (U-Net). We name it InstaFlow-0.9R. (53.6 A100 GPU days)
The total training cost for InstaFlow-0.98 3, 200, 000/200, 000 (Data Generation) 11.2 + 64 54.4 + 53.6 = 199.2 A100 GPU days.
Implementation Details and Training Pipeline for InstaFlow-1.7B We adopt the Stacked U-Net structure in Section C but abandon the parameter-sharing strategy. This gives Stacked U-Net with 1.7B parameters, almost twice as large as the original U-Net. Starting from 2-Rectified Flow, 2-Rectified Flow+Distill (Stacked U-Net) is trained by the following distillation steps:
1. Distill (Stage I): The Stacked U-Net is initialized from the weights in the 2- 2-Rectified Flow checkpoint. Then we fix the time t 0 for the neural network, and fine-tune it using the distillation objective equation 6 with batch size of 64 for 110,000 iterations. The similarity loss D is L2 loss. (35.2 A10 GPU days)
2. Distill (Stage 2): We switch the similarity loss IDD to LPIPS loss, then we continue to train the model using the distillation objective equation 6 and batch size of 320 for another 2,500 iterations. The final model is 2-Rectified Flowi (Stacked U-Net). We name
InstaFlow-1.7B. (4.4 A100 GPU days)
Discussion (Experiment Observations) During training, we made the following observations: (1) the 2-Rectified Flow model did not fully converge and its performance could potentially benefit from even longer training duration; (2) distillation showed faster convergence compared to reflow; (3) the LPIPS loss had an immediate impact on enhancing the visual quality of the distilled one-step model. Based on these observations, we believe that with more computational resources, further improvements can be achieved for the one-step models.
Discussion 2 (One-Step Stacked U-Net and Two-Step Progressive Distillation) Although one- step Stacked U-Net and 2-step progressive distillation (PD) need similar inference time, they have two key differences: (1) 2-step PD additionally minimizes the distillation loss at 0.5, which may be unnecessary for one-step generation from / = O; (2) by considering the consecutive U-Nets as model, we are able to examine and remove redundant components from this large neural network, further reducing the inference time by approximately 8% (from 03s to 0.12s).
E ADDITIONAL QUALITATIVE RESULTS
including inspections on the latent spaces of InstaFlow-0.99 and InstaFlow-1177 and visual comparison of few-step generation and guidance scale a. We also show uncurated images generated from 20 random LAION text prompts with the same random noises for visual comparison. The images from different models are shown
We provide additional qualitative results in Figure 13
14 15 16
in Figure
24 25
Alignment between 2-Rectified Flow and the One-Step Models The learned latent spaces of generative models have intriguing properties. By properly exploiting their latent structure, prior works succeeded in image editing [89 34 24
47 101). semantic control [61] II 44], dis- etc.. In general, the latent spaces of one-step generators, like GANs, are usually easier to analyze and use than the multi-step diffusion models. One advantage of our pipeline is that it gives multi-step continuous flow and the corresponding one-step models simultancously. Figure 17 shows that the latent spaces of our distilled one-step models align with 2-Rectified Flow. Therefore, the one-step models can be good surrogates to un- derstand and leverage the latent spaces of continuous flow, since the latter one has higher generation quality.
84
57
entangled control direction discovery 22 63
8g
75
22

---

Published as a conference paper at ICLR 2024
medievall fantuay tonn, down, durk, highly detuiled"
small silting movie saiing papcora engine"
Juturistic lighthouse, flash Hight, typer realistic, opic composition, cinomatic'
photo of delictous
huurry plate,
impressioniam misterpiece
Figure 13: Latent space interpolation of our one-step InstaFlow 0.9B. The images are generated 0.09s, saving ~ 90% of the computational time from the 25 25-step SD-1.5 teacher model the inference stage.
original
+leather jacket
line art
original
-brown hoodie
abstrratt art
Chinese
black
[Avttur, Pandora
dootorr
[black]
Elf, Lord of Rings
'4 high-razohation conternd (photu/ of handuuee mus, lightingg
high-reselation [photo] of benwifil |] girl, lighting
Figure 14: Images generated from our onc-step InstaFlow-1.78 in 0.12s. With the same random noise, the pose and lighting preserved across different text prompts.
23

---

Published as a conference paper at ICLR 2024
N 2
N =
N = 8
N 25
SD 1.5-DPM Solver
2-Rectified Flow
'Masterpiece color pencil drawing horse; bright vivid color'
Figure 15: Visual comparison with different number of inference steps N. With the random seed, 2-
Rectified Flow can generate clear images when N while 1.5-DPM Solver cannot.
= 1.0
1.5
2.0
a 3.0
4.0
'A high-resolttion photo of an orange Porsche under sunshine"
'A minimalism Teddy bear in front of a wall of red roses"
Figure 16: Visual comparison with different guidance scale 2-Rectified Flow. When 1.0, the generated images have blurry edges twisted details; when a 2.0, the generated images gradually gets over-saturated.
2-RF stup)
InstaFlow-O.9B 9B
InsisFow-11778
2-RF (25 InstaFlew-0..0
IntiFFom- 1.7B
'A hyperroalistic photo of a for astromaui; perfect face, aristation"
Temple in forest,
concept
Figure With same random noise and text prompts, the one-step models generate similar images with continuous 2-Rectified Flow, indicating their latent space aligns. Therefore, the one-step models can good surrogates to analyze the properties of the latent space of the continuous flow.
24

---

Published as a conference paper at ICLR 2024
Figure 18: Uncurated samples from SD+Distill (U-Net) trained with a learning rate of 10 and a weight decay coefficient of 10-3
Figure 19: Uncurated samples from SD+Distill (U-Net) trained with a learning rate of 10 and a weight decay coefficient of 10 3
25

---

Published as a conference paper at ICLR 2024
THE AIEAA
Figure 20: Uncurated samples from SD+Distil (U-Net) trained with a learning rate of 10 and a weight decay coefficient of 10-3
Figure 21: Uncurated samples from SD+Distill (Stacked U-Net) trained with a learning rate of 10 and weight decay coefficient of 10
26

---

Published as a conference paper at ICLR 2024
D
Figure 22: Uncurated samples from Stable Diffusion 1.5 with 25-step DPMSolver 49
and guidance scale 5.0.
LITINISS
Figure 23: Uncurated samples from 2-Rectified Flow with guidance scale 1.5 and 25-step Euler solver.
27

---

Published as a conference paper at ICLR 2024
TAFINIES
CAOT THI
Figure 24: Uncurated samples from one-step InstaFlow-OO.9
TATINESS
Figure 25: Uncurated samples from one-step InstaFlow-L.7B
28

---

Published as a conference paper at ICLR 2024
Reconstructed (2-Rectified Flow)
Ninverse Nyorwars
Nasverra Nyorward
Nanvers Narward
Original Image
Ninseese 4
Misaorso Myooraard 25
Musserso
Nyorward 25
Nyarward 25
Ninverar Nyrrward
Masversa Narrwart
Naswrs Naorwara
Original Iraaee
Nawerre 4
Missrrse 2
Nosserse Mjarward 25
Nyowward
Nfoowsnn 25
Figure 26: Examples of image encoding and reconstruction with 2-Rectified Flow, Here, encode an image X1 to the latent noise space by simulating the inverse probability flow ODE, X, = X1 + J v(X,, t)dt. Then we reconstruct the image from the latent encoding Xo by simulating the forward probability flow ODE, X, Xo + for vxX,, t)dr. We show the reconstructed image X.. For both stages, we adopt Fuler solver and use Niveerse and N forward steps respectively. With as few as 4 steps, 2-Rectified Flow can encode and reconstruct the original image successfully. Since 2-Rectified Flow straighter, the encoding stage works well with even step.
29

---

Published as a conference paper at ICLR 2024
Original
Stormtrooper
Albert Einstein
Roman sculpture
Medieval knight
Condition: Pose
Van Gogh style Orieinal
Original
Emply text
Empty text
Condition: Depth
+ suit
Condition: Edge
+ a white horse
Figure 27: We surprisingly find that pre-trained InstaFlow fully compatible with ControlNets (99 pre-trained
with Stable Diffusion. The images shown here generated with one-step InstaFlow pre-trained ControlNets.
30


Reviews:
[{"review_id": "YDsyOMfoht", "reviewer": "Reviewer_VVpt", "summary": "This paper successfully demonstrates the use of RECTIFIED FLOW to linearize the model's sampling trajectory, followed by distillation to enhance the sampling speed of the ODE model. It proposes a method for distilling Text-Conditioned flow models, showcasing a variety of ablation studies and results across multiple settings.", "questions": "Given that RECTIFIED FLOW is trained based on its own trajectory, are there any issues that arise from this approach?\n\nMethods like DDIM inversion also seem like they could be applicable in a flow-based context. I am curious about the results in cases where a small number of steps, close to 2, are used.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "The results of this paper are truly captivating. It manages to generate images of impressive quality with just one or two steps. The quality showcased in the figures is highly satisfying. In addition, the paper provides a detailed account of various experiments and the corresponding performance metrics, which adds greatly to its value.", "weaknesses": "While the paper demonstrates impressive results, it appears to be a straightforward application of RECTIFIED FLOW. I was unable to discern any clear novelty in the algorithms or methods presented. If I'm wrong please kindly let me know the difference.\n\nThe model that claims to operate in 1 step actually resembles a configuration of two UNets linked together, and thus feels closer to a 2-step process. Additionally, when the refined model from SDXL is not applied, the results show a noticeable degradation in high-frequency details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5SUbNXhSEd", "reviewer": "Reviewer_ZDe9", "summary": "This paper extends a recently introduced framework of rectified flows to the distillation of the coupling learned by the pretrained diffusion model, e.g. Stable Diffusion (SD). \nWhile the recent work [1] reported the results of experiments on unconditioned generation (on CIFAR-10, LSUN, AFHQ, MetFace and CelebA-HQ datasets) as well as on img2img translation, the current submission focuses on text-conditional generation.\nThe paper reconfirms that a \"rectified\" ODE produces an easier target for 1-step distillation. \n\nThe main contribution is the InstaFlow model which essentially is a multi-step pipeline which takes a pretrained SD model as an input and outputs a 1-step generative model. \nIn addition, a novel type of architecture called Stacked U-Net is presented.\nAs the conducted evaluation shows, InstaFlow outperforms recent baselines such as Progressive Distillation of SD and StyleGAN-T in terms of FID and CLIP score.\n\n[1] Liu et al. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. In ICLR, 2023.", "questions": "1. Please, address the limitations discussed above.\n1. The training pipeline described in the Appendix D, looks pretty complicated. \n    1. Why was it necessary to change the batch size (step 2), and why wasn't more common learning rate tuning applied instead? \n    1. What is the reasoning behind switching from L2 to LPIPS objective (step 4) instead of training with a combination of L2 and perceptual loss from the very beginning of the distillation phase? \n    1. How were switching points for steps 2 and 4 selected?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The writing style of the paper is extremely clear.\nIt provides a very good introduction to the framework of rectified flows for a reader without deep knowledge of the topic, and overall the manuscript is quite self-contained. \nThe conducted experiments provide a sufficient support for the motivation of the method. I find the evaluation thorough enough.\nThe results achieved in the paper are definitely interesting for the broad community of ML researchers and practitioners due to the achieved combination of required computational resources for the model training and inference and its performance.", "weaknesses": "1. First of all, this submission is more an extension of the previous work [1] rather than an independent work. The novelty of the presented ideas is definitely limited: all parts of the pipeline were actually introduced previously, and the submitted work applies the same pipeline to the coupling learned by SD model instead of independent coupling of noise and images. The proposed results are definitely valuable for applications. However, they look more like a technical exercise on top of the [1]. Overall, I find this work too incremental although helpful for practitioners.\n1. While the idea of Stacked U-Net is interesting, the paper lacks the study if this type of architecture is actually better than increasing the depth (or the number of channels) of the conventional U-Net model.\n1. The paper provides the results for latent models only. Taking the empirical nature of this work into account, I suggest adding any of the open cascaded models to the comparison to see, e.g. DeepFloyd IF [2].\n\n[2] https://github.com/deep-floyd/IF", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NTFw9qhBG0", "reviewer": "Reviewer_W54n", "summary": "This paper proposes the InstaFlow, an application that applies the Rectified Flow to Stable Diffusion. The authors implement the rectified flow technique on Stable Diffusion and subsequently distill a one-step diffusion model from the rectified model. The rectified method makes the trajectories of Stable Diffusion straighter, thus making it much easier to distill the multi-step model to fewer or even one-step model.", "questions": "I don't have further questions regarding the experiments since they're satisfactory to me.\nHowever, given that this work involves practical applications, I encourage the authors to consider sharing the source code and pretrained models, as this would undoubtedly be of great benefit to the community.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written, and the experiments conducted are both sufficient and convincing. The InstaFlow achieves amazing performance (1-step inference with reasonable quality in approximately 0.09 seconds).", "weaknesses": "While this model demonstrates impressive performance, it does involve a trade-off between inference speed and generation quality. From the supplementary document of InstaFlow, we can still observe various artifacts, which may be inherited from the 2-Rectified Flow (e.g., many faces are already distorted in the rectified model). Nevertheless, as the authors also mentioned, this model can be used for generating quick reviews, and then larger models can be employed for further generating high-quality images.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JJmEHrFvtv", "reviewer": "Reviewer_mwbp", "summary": "The paper aims to address the limitations of diffusion models in text-to-image (T2I) generation, particularly their slow multi-step sampling process. The authors propose a novel one-step generative model derived from Stable Diffusion (SD) using a method called Rectified Flow. The core of Rectified Flow is its reflow procedure, which improves the coupling between noises and images and facilitates the distillation process.", "questions": "I tend to accept the paper, considering its theoretical and technical contributions. However, I have several questions about this paper and hope the authors answer them for me to make the final decision.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Well organized and clarified.\n2. The contribution is great. Making the conditional diffusions work with one or very few steps will greatly boost the development of diffusion community. \n3. The comparison experiments are carefully and fairly set up, and I appreciate that.", "weaknesses": "I have several questions about this paper, and I hope the authors to clearly clarify them.\n\n1. **Storage Overhead**: In my opinion, it seems that either the distillation process or the reflow process actually requires us to create a relatively large (noise, image) pair dataset in advance, which would cause the additional storage overhead.\n\n2. **Intrinsic Difference between the so-called distillation and reflow process**: The distillation step aims to make the model predict the same as target computed by the ODE process of Stable Diffusion at the zero-timestep. While the reflow process seems to only change to make the distillation applied to all possible timesteps. \n\n3. **Noise Scheduler**: The noise scheduler of SD requires a normal diffusion noise scheduler. The reflow requires the \"linear\" (I call it \"linear\" just for convenience) scheduler. Wouldn't that cause trouble? Besides, the SD Unet requires time embedding, what time embedding do you use?\n\n4. Why do you choose to predict \"x1-x0\" instead of \"x1-xt\"? Do you have any considerations about this?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)