Improving the Training of Rectified Flows
Authors: Sangyun Lee
Zinan Lin
Giulia Fanti
Venue: NeurIPS 2024 poster (2024)
Keywords: generative modeling
rectified flow
diffusion model
Primary area: generative_models
TL;DR: We propose improved techniques for training rectified flows, allowing them to compete with knowledge distillation methods even in the low NFE setting

Abstract: Diffusion models have shown great promise for image and video generation, but sampling from state-of-the-art models requires expensive numerical integration of a generative ODE.
    One approach for tackling this problem is rectified flows, which iteratively learn smooth ODE paths that are less susceptible to truncation error.
    However, rectified flows still require a relatively large number of function evaluations (NFEs).
    In this work, we propose improved techniques for training rectified flows, allowing them to compete with knowledge distillation methods even in the low NFE setting.
    Our main insight is that under realistic settings, a single iteration of the Reflow algorithm for training rectified flows is sufficient to learn nearly straight trajectories; hence, the current practice of using multiple Reflow iterations is unnecessary.
    We thus propose techniques to improve one-round training of rectified flows, including a U-shaped timestep distribution and LPIPS-Huber premetric.
    With these techniques, we improve the FID of the previous 2-rectified flow by up to 75\% in the 1 NFE setting on CIFAR-10.
    On ImageNet 64$\times$64, our improved rectified flow outperforms the state-of-the-art distillation methods
    such as consistency distillation and progressive distillation in both one-step and two-step settings and rivals the performance of improved consistency training (iCT) in FID.
    Code is available at https://github.com/sangyun884/rfpp.

Improving the Training of Rectified Flows
Sangyun Lee Carnegie Mellon University sangyunl@andres Gandrew. cmu.edi
Zinan Lin
Giulia Fanti Carnegie Mellon University gf efanti Gandrew. cmu edu
Microsoft Research
zinanl cinanlin@mi www.iiiiiiiiooooooooommmmmooo
200.::2405.0332000 [cs.CV] 8 Oct 2024
Abstract
Diffusson models have shown great promise for image and video generation, but sampling from state-of-the-art models requires expensive numerical integration of generative ODE. One approach for tackling this problem is rectified flows, which iteratively learn smooth ODE paths that are less susceptible to truncation However, rectified flows still require relatively large number of function evaluations (NFEs). In this work, we propose improved techniques for training rectified flows, allowing them to compete with knowledge distillation methods even in the low NEE setting. Our main insight is that under realistic settings, a single iteration of the Reflow algorithm for training rectified flows is sufficient to learn nearly straight trajectories; hence, the current practice of using multiple Reflow iterations is unnecessary. We thus propose techniques improve one-round training of rectified flows, including a U-shaped timestep distribution and LPIPS- Huber premetric. With these techniques, we improve the FID of the previous 2-rectified flow by up to 75% in the NFE setting on CIFAR- 10. On ImageNet 64 64, our improved rectified flow outperforms the state-of-the methods such as consistency distillation and progressive distillation in both one-step and two-step settings and rivals the performance of improved consistency training (iCT) in FID. Code is available at withttt:: h/pp:://gt thub. cww/cooccocccoo. saas//nnnnn/r/fpppp
distillation
Introduction
Diffusion models (Soht-Dickstein al. 2015 Ho al. 2020 Song and Ermon 2019 Song et al. 2020b) have shown great promise in image |Ramesh al. 2022 and video [Ho et al. 2022 generation. They generate data by simulating stochastic denoising process where noise is gradually transformed into data. To sample efficiently from diffusion models, the denoising process is typically converted into a counterpart of Ordinary Differential Equations (ODEs) [Song et al. 2020b called probability flow ODEs (PF-ODEs).
Despite the success of diffusion models using PF-ODEs. drawing high-quality samples requires numerical integration of the PF-ODE with small step sizes, which is computationally expensive. Today, two prominent classes approaches for tackling this issue are: (1) knowledge distillation (e.g., consistency distillation [Song et al.) 2023), progressive distillation [Salimans and Ho 2022))
and (2) simulation-free flow models (e.g., rectified flows [Liu et al. 2022], flow matching [Lipman
al. 2022
In knowledge distllation-based methods [Luhman and Luhman 2021 Salimans and Ho 2022 Song
et al., 2023 Zheng et al.) 2022) student model is trained to directly predict the solution of the PF-ODE. These models are currently state-of-the-ard in the low number of function evaluations (NFEs) regime (c.g. 1-4).
Another promising direction is simulation-free flow models such rectified flows [Liu et al. 2022
Liu 2022 a generative model that learns a transport map between two distributions defined via neural ODEs. Diffusion models with PF-ODEs are special case. Rectified flows leam smooth
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

ODE trajectories that are less susceptible to truncation error. which allows for high-quality samples with fewer NFEs than diffusion models. They have been shown outperform diffusion models in the moderate to high NFE. regime [Lipman et al., 2022 Liu et al. 2022 Esser et al. 2024]. but they
still require a relatively large number of NFEs compared to distillation methods.
Compared to knowledge distillation methods (Luhman and Luhman 2021 Salimans and Ho 2022 Song et al. 2023 Zheng al. [022] rectified flows have several advantages. First, they can be generalized to map two arbitrary distributions to one another, while distillation methods are limited
to Gaussian noise distribution. Also, as a neural ODE, rectified flows naturally support inversion
from data to noise, which has many applications including image editing [Hertz et al. [002 Kim
2022 Wallace et al.] [023] Couairon et al. [022] Mokady et al.) [023] Su et al. [022] Hong et al.
al.
2023 and watermarking Wen et al.. [023]] Further, the likelihood of rectified flow models can
be evaluated using the instantaneous change of variable formula (Chen et al. 2018], whereas this
is not possible with knowledge distillation- baasd methods. In addition, rectified flows can flexibly adjust the balance between the sample quality and computational cost by altering NFEs, whereas distillation methods either do not support multi-step sampling or do not necessarily perform better with more NFEs (e.g. > 4) [Kim et al. 2023]
Given the qualitative advantages of rectified flows, a natural question is, can rectified flows compete with distillation-based methods such as consistency models [Song et al. 2023] in the low NFE setting? Today, the state-of state-of-the-art techmiques for training rectified flows use the Reflow algorithm to improve low NFE performance (Liu et al.. [022 2023]. J Reflow is a recursive training algorithm where the rectified flow is trained on data-noise pairs generated by the generative ODE of the previous stage model. In current implementations of Reflow, to obtain a reasonable one-step generative performance, Reflow should be applied at least twice, followed by an optional distillation stage to further boost performance [Liu et al.) 2022 2023]. Each training stage requires generating a large number of data-noise pairs and training the model until convergence, which is computationally expensive and leads to error accumulation across rounds. Even with these efforts, the generative performance of rectified flow still lags behind the distillation methods such as consistency models [Song et al. 2023
We show that rectified flows can indeed be competitive with the distillation methods in the low NFE setting by applying Reflow with our proposed training techniques. Our techniques are based on the observation that under realistic settings, the linear interpolation trajectories of the pre-trained rectified flow rarely intersect with each other. This provides several insights: 1) applying Reflow once sufficient to obtain straight- line generative ODE in the optima, 2) the training loss of 2-rectified flow has zero lower bound, and 3) other loss functions than the squared E2 distance can be used during training. Based upon this finding, we propose several training techniques to improve Reflow, including: (1) a U-shaped timestep distribution, (2) an LPIPS- huber premetric, which we find to be critical for the few-step generative performance. After being initialized with pre-trained diffusion models such EDM [Karras et al.) 20022). J. our method only requires one training stage without additional Reflow or distillation stages, unlike previous works [Liu et al. 2022 2023).
Our evaluation shows that on several datasets (CIFAR-10 Krizhevsky et al., 2009). ImageNet 64x64 (Deng et al., 2009), ). our improved rectified flow outperforms the ofate-of-the-art distilla- tion methods such as consistency distillation (CD) [Song et al.] 2023) and progressive distillation (PD) [Salimans and Ho 2022) in both one-step and two-step settings, and it rivals the performance of the improved consistency training (iCT) [Song et al. 2023 in terms of the Frechet Inception Distance [Heusel et al. 2017 (FID). Our training techniques reduce the FID of the previous 2-rectified flow [Liu et al. 2022 by about 75% (12.21 + 3.07) on CIFAR-10. Ablations on three datasets show that the proposed techniques give a consistent and sizeable gain. We also showcase the qualitative advantages of rectified flow such as few-step inversion, and its application to interpolation and image-to- to-image translation.
2 Background
2.1 Rectified Flow
Rectified flow (see also flow matching [Lipman et al. 2022 and stochastic interpolant (Albergo and Vanden 2022 ) is generative model that smoothly transitions between two distributions
Px and Pz by sol ving ODEs [Luu et al.) 2022]. For x ~ Px and z ~ Pas we define the interpolation between x and z as X; (1 t)x tz for t E [0, 1]. Liu et al. [0022] showed that for zo Px, the

---

B x :3
X
X
Px
Pz
Px
Px
Px
Ps
Px
Ps
Lincar interpolation induced by Piz
(b) Generative process of 1-rectiffed now
Lincar interpolation induced by pla
Generative process of 2-recifiied flow
Figure Rectified flow process (figure modified from Liu trajectories so there are no intersecting trajectories (a) + (b). Then, we take noise samples from Ps. and their generated samples from Px. and linearly interpolate them (c). In Reflow, rectified flow applied again (c) + (d) to straighten flows. This procedure is repeated recursively.
[202 Rectified flow rewires
following ODE yields the same marginal distribution as x, for any t:
dz: dt
ve(zi) == - - =z))-
Since x2 z, Eq. ( transports Px Pz- We can also transport pz to Px by drawing z from Pz and solving the ODE backwards from t= = 1 to = O. During training, we estimate the conditional expectation [E[x[x; = 2e] with a vector-valued neural network xg trained on the squared &g loss:
min E x.zxp.. E Eipp, (t)l (x,,,, =)
where Paz is the joint distribution of and Xe is parameterized by 0, and (t) is weighting function. Pt is chosen to be the uniform distribution on [0, 1] in Liu et al. [2022 2023]. In the optimum of Eq. 2 Xy becomes the conditional expectation (MMSE) estimator, which is then plugged into the ODE ( to generate samples. Instead predicting the conditional expectation directly, Liu al. 2022 choose to parameterize the velocity v, with neural network Ve and train it on
is minimum mean squared error
min Ex Exup, Ill(z x) vo(xi, 1)
which is equivalent to Eq. 2 with (() t.2 See Appendix, A
In this paper, we consider the Gaussian marginal case, i.e., P2 (0, In this case, if we define x
and z as independent random variables (i.e., Pxx(x, z) Px (x) (x))p(()) (z)) and use specific nonlinear
interpolation instead of the linear interpolation for Xt. Eq. (2) becomes the weighted denoising objective of the diffusion model [Vincent, 2011] and Eq. becomes the probability flow ODE (PF-ODE) [Song et al. 2020b].
2.2 Reflow
Algorithm 1 Reflow Procedure
1: First iteration:
2: 0 = arg ming Exxerp Braps [w(1)[[x xx(x,, 1)]
Train I-rectified flow
3: for k = K do
4:
(ze X0, (2e:t) ()) with 01 = z 5: xx, = (2(2) (x TTx(2)
D Generate synthetic pairs for next coupling Train (k: 1) -rectiffee flow
6: Oct1 arg ming Ex, Ex Bemps [w]t][[x xx(x,,))]l
7: end for
The independent coupling Pxu (x, z) Px (x)))2(()) is known to lead to curved ODE trajectories, which require large number of function evaluations (NFE) to generate high-quality samples (Poola dian et al. 2023 Lee et al.) 2023 Reflow [Liu et al. 2022 is recursive training algorithm to find better coupling that yields straighter ODE trajectories. Starting from the independent cou- pling Pha (x, z) (z), the Reflow algorithm generates p++ - z) from px (x, z) by first

---

x x²)
M
Generated from
Genernod from
[| of -XII]
2+1xx-x²)
Preeuuncy
Genenaive (I-Rectifial Fluw)
90 038 313 1 14s 12)
(9)
Figure 2: An illustration of the intuition in Sec. 3 (a) If two linear interpolation trajectories intersect,
z" zz is parallel to x' - -x"". This generally maps z" to atypical (e.g., one with high autocorrelation or norm that too large to be on Gaussian annulus) realization of Gaussian noise, so the 1-rectified flow cannot reliably map z" to x" MMx (b) Generated samples from the pre-trained I-rectified flow starting from z ~ N (0, I) (right), which is the standard setting, and z" = z + (x' x"). where x', x" are sampled from I-rectified flow trained on CIFAR-10 (left). Qualitatively, we see that the left samples have very low quality. (c) Empirically, we show the & norm of 2" = z + +x'xx' -x") compared to 2', which is sampled from the standard Gaussian. z" generally lands outside the annulus of typical Gaussian noise. (d) z + (x' - x") has high autocorrelation while the autocorrelation of Gaussian noise is nearly zero in nigh-dimensional space.
generating synthetic (x, z) pairs from PÊ then training rectified flow the generated synthetic pairs (Figure (b) - (d)). We call the vector field resulting from the k-th iteration of this procedure k-rectified flow. Pseudocode for Reflow is provided in Algorithm.
Convergence: Liu et al. [0022 show that Reflow trajectories are straight the limit as 00. Hence, to achieve perfectly straight ODE paths that allow for accurate one-step generation, Reflow may need to be applied many times until equilibrium, with each training stage requiring many data-noise pairs, training the model until convergence, and a degradation in generated sample quality, Prior work has empirically found that Reflow should applied least twice (i.e. 3-rectified flow) for reasonably good one-step generative performance (Liu et al. 2022 2023). This has been a major downside for rectified flows compared to knowledge distillation methods, which typically
require only one distillation stage [Luhman and Luhman 2021
Song et al.
2023 Zheng et al. 2022
Applying Reflow Once Sufficient
In this section, we argue that under practical settings, the trajectory curvature of the optimal 2-rectified flow is actually close to zero. Hence, prior empirical results requiring more rounds of Reflow may be the result of suboptimal training techniques, and we should focus on improving those training techniques rather than stacking additional Reflow stages.
First, note that the curvature of the optimal 2-rectified flow is zero if and only if the linear interpolation trajectories of I-rectified flow- dowggenertted pairs do not intersect, or equivalently, E|x|xx (1 t)x' tz'] = x' for all pairs (x', zz) [Liu et al. 2022].
To begin with, consider the manifold My of the synthetic distribution (x) fp' z)dz.
Consider two points x' and x" from the manifold, and two noises z' and z" that are mapped to x' and x" by 1-rectified flow. Here, we say pairs (x', z') and (x".z") intersect if 3t E (o. 1] s.t. - (1-t)x' + tz' = (1 - t)x" + tz". For example, in Figure aa), we observe that the two trajectories intersect an intermediate l.
For an intersection to exist at t must hold that 1) 1-rectified flow maps z" to x". and 2) x" z' + 1-f Ix(('' xx) by basic geometry. However, note that for realistic data distributions and if

---

1-rectified flow is sufficiently well-trained, z" = z'' - -x") is not common noise realization (e.g., it is likely to have nonzero autocorrelation or norm that is too large to be on Gaussian annulus), as shown visually in Fīgure As 1 1-rectified flow is almost entirely trained on common Gaussian noise inputs, cannot generally map an atypical z" to MMx Figure 2ec) shows qualitatively
that if we draw values of z" by first drawing z' ~ (0, 1) and then adding (x' -x") for independent draws of x', x", the z" vectors fall outside the annulus of typical standard Gaussian noise. Similarly, Figure d) shows that the constructed noise vectors z" have higher autocorrelation than expected. As
a result, Figure b visually shows that the generated samples have little overlap with the expected samples from typical draws of z'.
This suggests empirically that when training 2-rectified flow, intersections are rare (i.e. E[x[x; (1 (1- + tz'] = x'), which in turn implies that the optimal 2-rectified flow trajectories are nearly straight. Hence, additional rounds of Reflow are unnecessary, while also degrading sample quality. This intuition allows us to focus on better training techniques for 2-rectified flow rather than training 3- or 4-rectified flow. It also leads us to several improved techniques, discussed in Sec. [
Edge cases: Note that if ||x' x"Il2 is small, I-rectified flow could map z" to some point on MM. However, it does not alter the conclusion because the average of x' and x" is close x' anyway, so E|x|x, = (1 f)x' + iz'] ~ x'. Similarly, if is close to 1, 1 (x' x") R O, so 1-rectified flow can map x" to Mx. If the 1-rectified flow is I--Lipschitz, ||x 12 I L||z' z z"112. Therefore, the expectation TE [x] again will not deviate much from x'.
4 Improved Training Techniques for Reflow
The observation in Sec. 3 suggests that the optimal 2-rectified flow is nearly straight. Therefore, if the one-step generative performance of the 2-rectified flow model is not as good as expected, it is likely due to suboptimal training. In this section, we show that the few-step generative performance of the 2-rectified flow can be significantly improved by applying several new training techniques.
4.1 Timestep distribution
As in diffusion models, rectified flows are trained on randomly sampled timesteps f, and the distribu- tion from which t is sampled is an important design choice. Ideally, we want to focus the training effort on timesteps that are more challenging rather than wasting computational resources tasks. One common approach to focus on the tasks where the training loss is high [Shrivastava et al.) 2016). However, the training error of rectified flows is not reliable measure of difficulty because different timesteps have different non-zero lower bounds. To understand this, let us decompose the training error into two terms:
easy
((0,t) = B Blllx (x t)1
1 B|| E +2(0,t) t).
(4)
Lower bound
The first term does not depend on f and thus can- not be reduced. The second term represents the ac- tual minimizable training error, but its value cannot be directly observed because the first term is usu- ally unknown. Fortunately, because of the finding in Sec. 3 we expect that the first term is nearly zero when training 2-rectified flow, so we can use C(0,t) for designing the timestep distribution.
Loss (scaled)
0.030
0.000
0.055
0.100
0.055
Figure 3 shows that the training loss of 2-rectified flow is large at each end of the interval & E [0, 1] and
0.0
0.
3/2
Figure 3: Training loss of the vanilla 2- rectified flow on CIFAR-10 measured on 5, 000 samples after 200, 000 iterations. The shaded area represents the 1 standard devi- ation of the loss. The dashed curve is our U-shaped timestep distribution, scaled by a constant factor for visualization.
small in the middle. We thus propose to use a U- shaped timestep distribution for Pa. Specifically, we define pe (v(u) xx exp(au) + exp(-au) on u E (0. 1). We find that a = 4 works well in practice (Table IT. Compared to the uniform timestep distribution (con- fig B), the U-shaped distribution (config C) improves
5

---

Table 1: Effects of the improved training techniques. The baseline (config A) is the 2-rectified flow with the uniform timestep distribution and the squared l metric [Liu et al. 2022]. Config B is the improved baseline with EDM initialization (.3) and increased batch size (128 - 512 on CIFAR-10) FII (the lower the better) is computed using 50, 000 synthetic samples and the entire training set. We train the models for 800, 000 iterations on CIFAR-10 and 1, 000, 000 iterations on AFHQ and FFHQ and report the best FID for each setting.
[4.3]
CIFAR-10
AFHQ 64 x 64
FFHQ x 64
Base Liu ct al.
2022] (A)
12.21
(A) EDM ini + larger batch size (B)
7.14
3.61
12.39
4.16
8.84
4.79
(B) Our (C)
5.17
3.37
9.03
3.61
6.81
4,66
(C) Huber (D)
5,24
3.34
8.20
3.55
7.06
4.79
(C) LPIPS- PIPPS-Huberdee (E)
3.42
2.95
4.13
3.15
5.21
4.26
(C) + LPIPS-HHube--- (F)
3.38
2.76
4.11
3.12
5.65
4.41
(F) Incorporating real data (G)
3.07
2,40
NFE
1
2
2
2
the FID of 2-rectified flow from 7.14 to 5.17 (a 28% improvement) CIFAR-10, 12.39 to 9.03 (27%) on AFHQ. and 8.84 to 6.81 (23%) on FFHQ in the one-step setting.
For I-rectified flow training, Pr was chosen to be the uniform distribution [Liu et al. 2022 2023 or logit-normal distribution [Esser al. 2024 which puts more emphasis on the middle of the interval. When training 1-rectified flow, model learns to simply output the dataset average when = 1 and the noise average (i.e., zero) when t 0. The meaningful part of the training thus happens in the middle of the interval. In contrast, from Eq. 3 we can see that 2-rectified flow learns to directly predict the data from the noise at and the noise from the data at O, which are nontrivial tasks. Therefore, the U-shaped timestep distribution is more suitable for 2-rectified flow.
4.2 Loss function
Previously, the squared & distance was used the training metric for rectified flow to obtain the
MMSE estimator [ xz) However, as we have shown in Sec. 3 that [[[[Xx Xt = (1 +)xx + Iz'] R x',
we can generalize Eq. (2) or equivalently Eq. to any premetric m (i.e. m(a, b) = 0 + a a b):
min E x-x Bamp, [m(z x, vo())
(5)
Note that without the intuition in Sec. 3 only the squared & distance would have been a valid premetric, any other premetric makes the model deviate from the intended optimum (the posterior expectation E[x[x,]> Although the choice of m does not affect the optimum, it does affect the training dynamics and thus the obtained model. Other than the squared lz distance, we consider the following premetrics:
Pseudo-Huber [Charbonnier al. 1997 Song and Dhariwal 2023]: Miub x, vo((x,, t))
viz-x (V((,, f)l12 c² c, where c 0.00054d with d being data dimensionality.
LPIPS-Huber: Mip-hub (z x, vo(x,, ()) (1 ()
(x,, + LLIPS((x
ve(x,, t)), where LPIPS(...) ) is the learned perceptual image patch similarity [Zhang et al. 2018
LPIPS-Huber- mip-hub-! (z x, vo(x,, ()) (1 l)) iz x, Vo(x,, ()) + L X
t vo(x,,t)).
The Rxcudo-Huber loss is less sensitive to the outliers than the squared lg loss, which can potentially reduce the gradient variance [Song and Dhariwal 2023] and make training easier. In our initial experiments, found that the Sgcudo-Huber loss tends to work better than the squared & loss with a small batch size (e.g. 128 on CIFAR-10) 10). When the batch size is sufficiently large, it performs on par with the squared &g loss CIFAR- 10 and DFHQ-64 and outperforms it on AFHQ-64 664, as shown in T As it is less sensitive to the batch size, we choose to use the Eseudo-Huber loss in the following experiments.
We also explore the LPIPS, which forces the model to focus on reducing the perceptual distance between the generated data and the ground truth. Since LPIPS is not premetric as two different

---

Table The converted time and scale for the variance preserving (VP) and variance exploding (VE)
diffusion models. Here, a<(t) = exp( f (19.9s + 0.1)ds) following Song et al. 2020b), and the
perturbation kernel of the VE diffusion is Mxx, I2I) following Karras et al. 2022
tvp
tVE
SvP
5VE
-0.05 + 0.0025 19.9 In
a(br)) T-t
9.95
V 122-22222
T-
points could have zero LPIPS if they are perceptually similar, we use it in combination with the Pseudo-Huber loss with the weighting where the task is more challenging. Note that in huh the gradient vanishes when is close to zero. To compsensate, we experiment with a where we multiply LPIPS by T. Compared to config D, the improvement) on CIFAR- 10, 8.20 to 4.11 (50%) on AFHQ, and 7.06 to 5.21 (26%) on FFHQ in the one-step setting, seen in Table I
thereby relying more LPIPS when is close to 1
loss improves the FID of 2-rectified flow from 5.24 to 3.38 (a 35%
4.3 Initialization with pre-trained diffusion models
Training 1-rectified flow from scratch is computationally expensive. Recently, Pokle al. 2023 showed that pre-trained diffusion models can be used to approximate E[x[x; 2/ in Eq. 1 by adjusting the signal-to-noise ratio. The following proposition the special cases of Lemma 2 of
Pokle et al. [2023] restated with extended proof and a minor fix. We provide the constants and proof in Appendix.
C.I
Proposition 1 Let (x((xx), t) be the posterior distribution of the perturbation kernel N(11 t)x, II). Also, let p (x xx, t) and y (x(xxxe,, t) be the posterior distributions of (a(())x (1 II and Nxx 12), each. Then,
f [x = Zz, 1)x dx = I (x|xx= SypZe, LVP) dx / 1E =Syphe, LVE) dx,
(6)
where Svp and sve are the scaling factors and typ and tvE are the converted times for the VP and VE diffusion models, respectively.
We have explicitly computed the time and scale conversion factors for the VP and VE diffusion models in Table 2 See Appendix derivation.
Proposition allows us to initialize the Reflow with the pre-trained diffusion models such as EDM Karras et al. 2022 or DDPM (Ho al. 2020 and use Table to adjust the time and scaling factors.
Starting from the vanilla 2-rectified flow setup Liu al. 2022 (config A), initialize I-rectified flow with the pre-trained EDM (VE). We also increase the batch size from 128 to 512 on CIFAR-10 compared to Liu al. [2022). Overall, these improve the FID of 2-rectified flow from 12.21 7.14 (a 42% improvement) in the one-step setting on CIFAR-10 (config B).
4.4 Incorporating real data
Training 2-rectified flow does not require real data (i.e.. it can be data-free), but if is available. To see the effects of incorporating real data, we integrate the generative ODE of 1-rectified flow backward from t 0 to = 1 using an NFE of 128 to collect 50, 000 pairs of (real data, synthetic noise) on CIFAR-10 10. For quick validation, we take the pre-trained 2-rectified flow model (config F) and fine-tune it using the (real data, synthetic noise) pairs for 5, 000 iterations with a learning rate of 1e-5. This improves the FID of 2-rectified flow from 3.38 to 3.07 in the one-step setting CIFAR-10 10 (config G).
real data
In this fine-tuning setting, we also explored using (synthetic data, real noise) pair with a probability of p. but we found that not incorporating (synthetic data, real noise) pairs at all (i.e., 0) performs
7

---

Table 3: Unconditional generation CIFAR-10.
Table 4: Class-conditional generation Ima-
geNet 64 X 64.
METROD
NFE (JD (s
Diffesion models
(L)
(f)
20206
Diffusion models
DDPM
al 204y
1000
3.17
9:16
LS0M
M021
147
2.10
DDIM [Son
2020a
50
13.7
0.65
0.56
18.3
0.60
0.49
Distilled dnlisioo mimiees
DPM solver
Lu et al. 2122
10
7.93
Knowledge Distilatton
Jihmmnn Lahhman 2021
3.42
DFNO (IP)) Zhene 2 NUM4
3.78
DEIS
Zhang and Chen 2022
10
6.66
TRACT Berthelor
3.78
3.32
DDPM
Ho et al 2020
250
11.0
0.67
0.58
PD Salimans and Ho. 2022
9.12
ADM
Wal and Nichol 2121
250
2,07
0.74
0.63
Score distillation
Distilled dilfusion modees
DMD Yin et al 00.5
3.77
DINO (LPIPS) Zheng et 2022
7.83
0.61
GANs
TRACT Berthefo
7.43
RigGAN Bruck al 7018
14.7
9.22
4,97
StyleGAN2
HU000
8.32
9.21
BOOT Ga 2023
16.3
0.68
0.36
StyloGAN2-4 ADA Karma
2020h
2.92
9.83
PD Salumans and Ho 2022
15.39
0.59
0.62
Consistency models
2
8.95
0.63
0.65
CD (LPIPS) Song et al 2023
3.55
9:48
6.77
0.66
0.65
2,93
9:75
Score distillation
(LPIPS)
Song
biff-Instrett Lao et
2024
5.37
5.83
8.85
DMD Vin
2.62
ICT
Song and Dhariwal
2023
2.83
9.51
GANs
BigGAN-deep
2018
0.79
ICT deep
Song and Dhariwal 20235
2,51
9.76
Consistency models
2.24
9.89
CTM
Kim et al.
2023
5.19
CD (LPIPS)
Song et
2023
6.20
0.68
0.63
CTM
Kim al
AIPEET
GAN
1.98
4.70
0.69
0.64
Nows
4.32
0.70
0.64
I-rectified Bow (+distil)
2023
6,18
9:05
CT (LPIPS) Song et
2023
13.0
0.71
0.47
2.rectified
11.1
0.69
0.56
110
3.36
9.24
iCT Song and Dhariwal
2023
4,02
0,70
0.63
edistill Liu 0022
4.85
3.20
0.73
0.63
3-rectifieT BOW
2022
iCT-deep Song and Dhariwal
2013
3.25
0,72
0.63
104
3.96
9.01
2.77
0.74
0.62
CTM GAN Kim al.
2023
1.92
0.57
2-rectified Nowss (OITS
3.07
0.57
2,40
Rectified Hows
E-rectibed flow+ (ours)
4.31
3.64
The red rows correspond to the top-5 baselines for the I-NFF setting, and the blue correspond top baselines the 2-NFE setting. The lowest FID scores for I-NFE and 2-NFE are boldfaced.
the best. We expect that training from scratch will further improve the performance with different values of p, and leave it to future work. A similar idea is also explored in Anonymous 2024
Experiments
We call these combined improvements to Reflow training 2-rectified flow++ and evaluate it on four datasets: CIFAR- 10 Krizhevsky et al. 22009. AFHQ Choi al. 2020). FFHQ Karras al. 20199)
and IeageNet Deng el al. [2009]. We compare our improved Reflow to up to 20 recent baselines, in the families of diffusion models, distilled diffusion models, score distillation, GANs, consistency models, and rectified flows. The details of our experimental setup are included in Appendix
5.1 Unconditional and class-conditional image generation
In Tables and ImageNet 64 X 64. We observe two main messages:
compare 2-rectified flow++ with the tlate-of-the-ar methods CIFAR-10 and
On both datasets, 2-rectified flow++ (ours) outperforms or is competitive with SOTA baselines in the 1-2 NFE regime. On CIFAR-10 (Table B). our 2-rectified flow achieves an FID of 3.07 one step, surpassing existing distillation methods such consistency distillation (CD) [Song et al. 2023). progressive distillation (PD) [Salimans and Ho 2022]. diffusion model sampling with neural operator (DSNO) (Zeng et al. 2022]. and T Annistivee Closure Time-distillation (TRACT) [Berthelot et al. 2023). On ImageNet x ( our model surpasses the distillation methods such as CD, PD, DFNO, TRACT, and BOOT in one-step generation. We also close the gap with iCT (4.01 vs

---

44
Euler
Edder
(a) CIFAR-10
(b) AFHQ 64 x 64
(c) FFHQ 64
Figure 4: Effects of ODE Solver and new update rule.
4.31), the state-of-the-art consistency model, even with half the batch size. Note that on ImageNet, our model does not use real data during training, while consistency training (CT) requires real data. We believe could further reduce the gap by using config G in Tab. Uncurated samples of our model are provided in Appendix. G
2-rectified flow++ reduces the FID of 2-rectffied flows by up to 75%. Compared to vanilla rectified flows [Luu et al. 2022]. our one-step FID on CIFAR-10 is lower than that of the previous 2-rectified flow by 9.14 (a reduction of 75%). and of the 3-rectified flow by 5.08 (see also Table T for ablations on other datasets). In addition, it outperforms the previous 2-rectified flow with 110 NFEs using only one step and also surpasses 2-rectified flow + distillation, which requires additional distillation stage.
5.2 Reflow can be computationally more efficient than other distillation methods
At first glance, Reflow seems computationally expensive compared to CD and CT requires generating synthetic pairs before training. How- ever, CD requires (1 for student, for teacher, and 2 for Heun's solver) forward passes for each training iteration, and CT requires (1 for student and for teacher) forward passes, while Reflow requires only 1. For example, in ImageNet experiment setting, the total number of forward passes for Reflow 395M + 1433.6M 6M = 1828.6 6M (395M for generating pairs and 1,433.6M for training), while the total numbers of forward passes for CD and CT would be 433.6-6-- 4 5, 734.4M and 1, 433.6 comparison. Moreover, generating pairs is a one- time cost since we can reuse the pairs for multiple training runs.
Table 6: Comparison of the number of forward passes. Reflow uses 395M forward passes for gen- erating pairs and 1, 433.6M for training.
Method
Per iteration
Total
Rel. total cost
Reflow CD CT
1
1828, 6M
x1
4
5734.4M,
x3.1
2
2867.2M
x1.5
2, 867. .2M under the same setting. See Table 6 the
In terms of the storage cost, the synthetic images for ImageNet 64 x 64 require 42 GB. For noise, we only store the states of the random number generator, which is negligible.
While these results should be further validated for larger datasets, our results suggest that the fact that Reflow requires generating synthetic pairs does not necessarily make it less computationally efficient than other distillation methods.
5.3 Effects of samplers
Unlike distillation methods, rectified flow is neural ODE, and its outputs approach the true solution of the ODE as NEE increases (i.e., precision grows). Figure that with the standard Euler solver, FID decreases as NFE increases on datasets. Moreover, Heun's second-order solver further improves the trade-off curve between FID and NFE. This suggests that there may be further room for improvement by using more advanced samplers. We provide some preliminary ideas towards this goal in Appendix D

---

0209
EDM
0192
0.89
0009
0009
(a) Reconstruction loss
(b) Distribution of ||z||
(c) Inversion/recoontrruiioon results
Figure 5: Inversion results on CIFAR-10. (a) Reconstruction enror between real and reconstructed data is measured by the mean squared eπror (MSE), where the x-axis represents NFEs used for inversion and reconstruction (e.g. 2 means for inversion and for reconstruction). (b) Distribution of || of the inverted noises proxy for Gaussianity (NFE 8). The green histogram represents the distribution of true noise, which Chi-squared with 32 32 3072 degrees of freedom. (c)
Inversion and reconstruction results using (8 + 8) NFEs. With only 8 NFEs, EDM fails to produce realistic noise, and also the reconstructed samples are blurry.
Lion- Tiger
Tiger-oLion
(a) Interpolation
(b) Image translation
Figure 6: Applications of few-step inversion. (a) Interpolation between two real images. (b) Image- to-image translation. The total NFEs used are 6 (4 for inversion and 2 for generation).
5.4 Inversion
Unlike distillation methods, rectified flows are neural ODEs, thus they allow for inversion from data to noise by simply integrating the ODE the backward direction. diffusion models, inversion has
been used for various applications such as image editing [Hertz et al. [022 Kim et al. 2022 Wallace
jet al., 2023 (Su et al., [022] Hong et al. [023] and watermarking Wen et al. 2023 but usually requires many NFEs. Figure (a) demonstrates that our 2-rectified flow++ achieves significantly
lower reconstruction error than EDM. Notably, the reconstruction error of 2-rectified flow++ with
only 2 NFEs is lower than that of EDM with 16 NFEs. In (b), compare the quality of the inverted noise, where we find that the noise vectors of 2-rectified flow are more Gaussian-lik than those of EDM, in the sense that their norm is closer to that of typical Gaussian noise. These are also shown visually in (c). In Figure we show two applications inversion: interpolating between two real images (a) and image-to-mmage translation (b). Notably, the total NFE used is only 6 (4 for inversion and 2 for generation), which is significantly lower than what is typically required in diffusion models (2 100) [Hong al. 2023
Conclusion
In this work, we propose several improved training techniques for rectified flows, including the U- shaped timestep distribution and LPIPS-Huber loss. We show that by combining these improvements, 2-rectified flows++ outperforms the state-of-ue of-the-art distillation methods in the 1-2 NFE regime on CIFAR-10 and ImageNet 64 64 and closes the gap with iCT, the state-of-the-ar consistency model. 2-rectified flows++ have limitations though- they still do not outperform the best consistency models (ICT), and their training is slower (by about 15% per iteration on ImageNet) than previous rectified flows because of the LPIPS loss. Despite these shortcomings, the training techniques we propose can easily and significantly boost the performance of rectified flows in the low NFE setting, without harming performance at the higher NFE setting.
10

---

Acknowledgments
This work made possible in part by the National Science Foundation under grant CCF-2338772, CNS-2325477, as well as generous support from Google, the Sloan Foundation, Intel, and Bosch. This work used Bridges-2 GPU the Pittsburgh Supercomputing Center through allocation CIS240037 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by National Science Foundation grants 2138259, 2138286, 2138307, 2137603, and 2138296 Boerner et al. 2023
References
Michael Albergo and Eric Vanden- Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint (tti:22209/5577/ 2022.
Anonymous. Balanced conic rectified flow. In Submitted to The Thirteenth International Conference on Learning Representations, 2024. httptttt///ooeeeeeiiee wtt.://ooooor..ddtttt
under review.
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint 23Xi:23333,0228,, 2023.
Timothy J Boerner, Stephen Deems, Thomas R Furlani, Shelley Knuth, and John Towns. Access:
Advancing innovation: Nsf's advanced cyberinfrastructure coordination ecosystem: Services & support. In Practice and Experience in Advanced Research Computing, pages 173-176. 2023.
Andrew Brock, Jeff Donahue, and Karen Simonyan, Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1096, 2018.
Pierre Charbonnier, Laure Blenc-Féraud, Gilles Aubert, and Michel Barlaud. Deterministic edge- preserving regularization in computed imaging. IEEE Transactions on image processing, 6(2): 298-311, 1997.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8188-8197, 2020.
Guillaume Couairon, Jakob Verbeek, Holger Schwenk, and Matthieu Cord. Diffedit: Diffusion-basee semantic image editing with mask guidance. arXiv preprint arXiv:2210.1 /427, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-255 Ieee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34, 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.03206, 2024.
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free distil- lation of denoising diffusion models with bootstrapping. In ICML 2023 Workshop on Structured Probabilistic Inference { Generative Modeling, 2023.
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or, Prompt- to-prompt image editing with cross attention control. arXiv preprint arXiv:2208.01 220.00200.,,6666, 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update nule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
11

---

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6857 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. arXiv preprint arXiv:220..03558 2022.
Scongmin Hong, Kyeonghyun Lee, Suh Yoon Jeon, Hyewon Bae, and Se Young Chun. On exact inversion of dpm-solvers. arXiv preprint arXiv:2311. 2023.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4410, 2019.
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in Neural Information Processing Systems, 33:12104- 10000111100 14, 2020a.
Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8110-8119, 2020b.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv preprint arXiv:2206 00364, 2022.
Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Ue- saka, Yutong He. Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279, 2023.
Gwanghyun Kim, Taesung Kwon, and Jong Chul Ye. Diffusionclip: Text-guided diffusion models for robust image manipulation. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pages 2426-2435, 2022.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Toronto, ON, Canada, 2009.
Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ode-based generative models. arXiv preprint vrXiv:2311.12000 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arX1w:2210.02747, 2022.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209. 14577, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion- based text-to-image generation. arXiv preprint arXiv:2309.066: 2023.
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. arXiv preprint arXiv:2206. 00927, 2022.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint SXXi::2011.0288, 2021.
Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diff- instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training, arXiv preprint 1700::1710 2017.
12

---

Ron Mokady, Amir Hertz, Kfir Aberman, Yael Pritch, and Daniel Cohen Or. Null-text inversion for editing real images using guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6038-6047, 2023.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-811 8171. PMLR, 2021.
Ashwini Pokle, Matthew J Muckley, Ricky TQ Chen, and Brian Karrer. Training free linear image inversion via flows. arXiv preprini arXiv:2310.04432, 2023.
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint WrXXi:22304.47722 2023.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint 1tti::22020005122 2022.
Neta Shaul, Juan Perez, Ricky TQ Chen, Ali Thabet, Albert Pumarola, and Yaron Lipman. Bespoke solvers for generative flow models. arXiv preprint arXiv:2310. 19075, 2023.
Abhinav Shrivastava, Abhinav Gupta, and Ross Girshick. Training region-based basion object detectors with online hard example mining. In Proceedings of the IEEE conference on computer vision and pattem recognition, pages 761-769, 2016.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference Machine Learning, pages
PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020a.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310. 14199, 2023.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint OrXiv:2011.134555 2020b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon. Dual diffusion implicit bridges for image-to-image translation. arXiv preprint arXiv:2203.08382, 2022.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 1287-11121. 2021.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compula- tion, 23(7):1661-1674, 201
Bram Wallace, Akash Gokul, and Nikhil Naik. Edict: Exact diffusion inversion via coupled transfor- mations. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 22532-22541, 2023.
Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers for diffusion models by differentiating through sample quality. In International Conference on Learning Representations, 2021.
13

---

Yuxin Wen, John Kirchenbauer, Jonas Geiping, and Tom Goldstein. Tree-ring watermarks: Fin- gerprints for diffusion images that are invisible and robust. arXiv preprint arXiv:2305.20030. 2023.
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:2311.18828, 2023.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204. 13902, 2022.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features perceptual metric. In Proceedings of the IEEE conference on computer vision and paitern recognition, pages 586-595 2018.
Hongkai Zheng. Weili Nie, Arash Vahdat, Kamyar Avizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. arXiv preprint arXiv:2211. 2022.
14

---

3000
2 2+ -x"
z-generated
2300
2+ xor'=x'-geeeene generated
Frequency
1500
1000
500
0.0
0.5
1.0
1.5 2.3 2.5
Autocorrelation (lag=3)
Reconstruction loss
(b) Distribution |
Figure 7: (a) Autocorrelation plot CIFAR-10 analogous to Figure 2d). (b) t-SNE visualization of the inception-v3 features of the samples in Figure b). They show negligible overlap.
A Equivalence of v-parameterization and x-parameterization
Given x, = (1- + and z = (x: 1x.-x]a /1, the equivalence of the x-parameterization (Eqq 22) be shown in the following way. The result borrowed from
and the v-parameterizatioon (Eq. the Appendix Lee et al. 2023 and we provide here for completeness.
Elll( x) - vo(x, 1)1 =
[ll((x x)/t vo t)(
(7)
[ll((x xx)/t- (xi -
dt
(8)
L.
18l 77 ||x xo(x, t))11) dt.
(9)
This is equivalent to Eq. 2 with w(t) = 1/t?.
B Additional Details for Figure 2
In this section, we provide additional results details for Figure 2 Figure 7(a) shows the autocorrelation histogram on CIFAR-10 while Figure d) is on FFHQ-64. Figure 7()) shows that the inception features of the samples in Figure b) rarely overlap with each other.
For the autocorrelation plots, we use 30,000 pairs of (x', x") and randomly sample z from the standard Gaussian distribution. For d-dimensional vector define the autocorrelation as:
R1(() = u Ug+1
(10)
where is the k-th element of vector u and / > 0 represents the lag.
C Initialization Results
C.1 Proof of Proposition I
Consider two perturbation kernels pe (x, (((())xx((()))))) and = o(t)
1
xx(xx (2)2(() exp(
|x: s(t)x|1)
(11)
20(1)2
1
(x(xxx = (2 (1)2)//2
exp(
1xx s'())x11))
(12)
20' (t)²
15

---

Let t'(t) be such that
e(0) s'(e)
We will show that pc(x(xx)
(x
s'(c') a(r)
xx).
We start by
o(t) o'(e)
showing:
Pe (010)) x1|x)= = (2 (t))))/22
8(())
1
exp(-
(1))
X - s((')xx 12)
(13)
s(t)
(0)
s(t)
(t'))
(2 (t)2)d/22 exp(
(x, (t)xx)]
(14)
20 20
s(t)
1
st())
exp(
||x, (()x1
(15)
(2
(t'
s(t)
s(t)
(2 (t)2)dd12 exp(
(t)² s(t) || s(txx
(16)
exp(
(t)2 1|x, - x(1)x
(17)
(P)o((())
1
(2))(1))))
exp(- 20 ||x| s())x(1)
(18)
(2 (t)))//22
(t)
(t)2)//2
(2TO' (1)))d/2P
(19)
(tt)
o'(t') (x(
(20)
says that pa (x, xx) plu (1()) xx(xx) (but not equal), which is minor fix from the
Here, Eq. 20
s(t)
original proof Pokle al. 2023 Then, we have
1
pr (x(x) (x,)
px (x)p((xx)))
(21)
stt')
ot))
1/. (x)
xx)
((1)))
Px (xx(x)x(Xx(x))
(22)
s(1)
(8(00) a(1)
xx)
Since
plz (x|
('()) s()) x,) should be integrated to one,
((())
d
(x(x))x (x)xx and
1 -ty x.)
thus two densities equal. As the posterior densities are the same, their expectations also the same.
In our case, s(t) = 1 - t, att t, and (x( x) is the perturbation kernel of either the VP or VE diffusion model. Now we have to find t' such that 1-t
s'((') t (1()")
C.2 Perturbation Kernel Instantiations
We next provide the values of the converted time and scale of the variance preserving (VP) and variance exploding (VE) diffusion models.
VE diffusion model Karras al. 2022 defines the perturbation kemnel of the VE diffusion model
as (xx = (xx,, II). Then, t' satisfies 1-
o'(e' r so
s'(t')
t' - and
d'(t') 1
<te
T
which
correspond to lye and sve in Table 2
VP diffusion model Song et al.
2020b
defines the perturbation kernel of the VP diffusion model
as (xxxx (N(((t)x, (1 It))
where at) cxp(- _ f (19.9s + 0. 0.1)ds) defined on
t [0, 1]. Then, t' satisfies 1-t
x'(e")
aft') From here, we have
('())
(1 - t)²
a(t')²
(23)
t²
1- -ta2
(1 t)²
a(l')
(24)
12 + (1 - -
16

---

where we used the fact that axt) > 0. Since att) is monotonically decreasing function for 2 0,
we can use its inverse a= to find L' = a a--1( (v (1+11-1))
(1-1)2
y = (t) exp( 1 / (19.9s + 0.1)ds) = exp(
19.9 0.05t)
(25)
2 /
4
Iy = 19.9 22 0.05t
(26)
4
19.9 12 + 0.05t + Iyy =0 0
(27)
4
Applying the quadratic formula, we have
-0.05 + Vo.05 4. 19.9 In
-0.05 + v0.0025 - 19.9 yy
1
(28)
2 19.9
9.95
Since y al) is monotonically decreasing, we can choose the positive root:
av) -0.05 + v0.0025 19.9
(29)
9.95
Now, we arrive at
-0.05 + V 0.0025 19.9 In
(1-1)
(1- 1)2
P+(1-17
t' = a= -==--(11
(V 12 (1-t)2
(30)
9.95
which corresponds lyp in Table
2
Also, we have
s'(e) o(t) a(0) I-
which is syp in Table 2
D New Update Rule
In the standard Euler solver, the update rule is 2t-Ar = 2t - v(z,, (). Alternatively, as X (2,, () of our model generates pretty good samples, xe( t) and 2j to get the next step: žę-Ár := (1 (t - Ax) t) + (t - At)z Note that when NFE < 3, the two update rules are equivalent and do not affect our results in Section 5.1 Fig.
instead use the linear interpolation between
shows that when applied to existing solvers, the new update rule improves the sampling efficinccy up
4x, achieving the best FID with < 5 NFEs.
Algorithm 2 shows the pseudocode for generating samples using the new update rule. Unlike the standard Euler update rule which only depends on the current state z, new update rule utilizes the previous state (i.e., 21) to generate the next state at-At and thus can be viewed as a form of history-dependent samplers. Obviously, incorporating the initial state only would not be the best choice. We believe that the result can be further improved, especially by using learning- based solvers Watson et al. 2021 Shaul et al., 2023]; and leave such exploration future work.
E Experimental Details
Before training 2-rectified flow, we generate data-noise pairs following the sampling regime of EDM [Karras al. 2022]. For CIFAR-10, we generate 1M pairs using 35 NFEs. For AFHQ, FFHQ. and ImageNet, we generate 5M pairs using 79 NFEs. We use Heun's is second-order solver for all cases.
In Table 3, we report the result of config G in Table T In ImageNet, we use the batch size of 2048 and train the models for 700,000 iterations using mixed-prec precision training MMciievviiius et al. 2017 with the dynamic loss scaling. We use config E in Tar ImageNet.
We provide training configurations in Table 7 For all datasets, we use Adam optimizer. We use the exponential moving average (EMA) with 0. 0,9999 decay for all datasets.
ImageNet, the training takes roughly 9 days with 64 NVIDIA V100 GPUs. On CIFAR-10 and FFHQ/AFHQ, it takes roughly 4 days with 16 and 8 V100 GPUs, respectively. For all cases, use
17

---

Algorithm 2 Generate
daf gznerata(Z1, label, model, time_achedule, N, selver, sampler, davica):
initial
label: class label
model: v_theta
acneuuchd time achedule, 8.B., 19999, 0] steps
N:
solver: 'euler' 'houn!
sampler: 'default? "new?
zl.clone()
ent
for in range [agge(eea(cinneecddlld-------]]: ecceeeddllll::-1))::
torch. ones (z..hhae [0]), devicerdovice) rima_schedulef3 [s]
t_next torch. Perchhonnes(C) on..( (2aa((22ahap)())) devicesdevice) time_schedule
t_next [0] t[0]
nodel label)
t.viev(-1111111)
solvar 'beun' and cnt <. 1: Houn correction
if sampler == 'default";
2_next 2.detachh cllone)))))))
olif samplor 'nev":
2_next = (1 t_next. F_next..tee(-111111111 x0hat. t_naxt.
ve_next
t_next, label)
vt (vt + ve_next) /
x0hat. = vt t.21,x(-1,1,1,1)
sampler 'default";
z dedeacc()..ccooe elif sampler -- 'nev';
dt
(1 (_oxx_.1
x0hat (1111,1)
Table 7: Training configurations for each dataset. We linearly ramp learning rates for all datasets.
Datasets
Batch size
Dropout
Learning rate
Warm up iter.
CIFAR-10
512
0.13
2e-4
5000 2c-4
FFHQ AFHQ
256
0.25
5000
ImageNet
2048
0.10
lc-4
2500
the NVIDIA DGX-2 cluster. To prevent zero-division error with EDM initialization, we sample from (0.00001, 0.99999) in practice. For two-step generation, we evaluate vo at = 0.99999 and t 0.8. For other NFEs, we uniformly divide the interval [0. [0.00001, 0.99999).
License
The following are licenses for each dataset we use:
CIFAR-10: Unknown
FFHQ: CC BY-NC-SA 4.0
AFHQ: CC BY-NC 4.0
ImageNet: Custom (research, non-commercial)
F
Broader Impacts
This paper proposes an advanced algorithm to generate realistic data at high speed, which could have both positive and negative impacts. For example, it could be used for generating malicious or misleading content. Therefore, such technology should be deployed and used responsibly and with caution. We believe that our work is not expected to have any more potential negative impact than other work in the field of generative modeling.
18

---

G Uncurated Synthetic Samples
We provide uncurated synthetic samples from our 2-rectified + on CIFAR-10, AFHQ, and ImageNet in Figures g 10
14. and 15]
We use our new sampler
(Sec. 5.3
to generate these images.
19

---

Figure 8: Synthetic samples from 2-rectified flow++ on CIFAR-10 with NFE = (FID=3.38).
Figure 9: Synthetic samples from 2-rectified flow+ on CIFAR-10 with NFE = 2 (FID=2.76).
20

---

Figure 10: Synthetic samples from 2-rectified flow++ CIFAR-IO with NFE 4 (FID=2.50).
Figure 11: Synthetic samples from 2-rectified flow++ on CIFAR-10 with NFE 5 (FID=2.45).
21

---

Figure 12: Synthetic samples from 2-rectified flow++ on ImageNet 64x64 with NFE = (FID=4.31).
G
Figure 13: Synthetic samples from 2-rectified flow++ on ImageNet 64x64 with NFE = (FID=3.644)
22

---

G
Figure 14: Synthetic samples from 2-rectified flow++ on ImageNet 64x64 with NFE 4 4 (FID-3.44)
G
Figure 15: Synthetic samples from 2-rectified flow++ on ImageNet 64x64 with NFE = 8 (FID=3.32).
23

---

Figure 16: Synthetic samples from 2-rectified flow++ on AFHQ 64 64 with NFE = 1 (FID=4.11).
Figure 17: Synthetic samples from 2-rectified flow++ on AFHQ 64 64 with NFE = 2 12)
24

---

Figure 18: Synthetic samples from 2-rectified flow++ on AFHQ 64 64 with NFE (FID=2.90).
Figure 19: Synthetic samples from 2-rectified flow++ on AFHQ 64 64 with NFE = 5 (FID=2.86).
25

---

Figure 20: Synthetic samples from 2-rectified flow++ FFHQ 64 with NFE = (FID=5.21).
Figure 21: Synthetic samples from 2-rectified flow++ on FFHQ 64x64 with NFE = 2 (FID=4.26).
26

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction reflect the paper's 's contributions and scope such
improved empirical performance, which is backed up by our experiments in Sec.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: The limitations are discussed in the conclusion, such as the increased training time of our method.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Our only theoretical result is Proposition ! for which we provide the proof in Appendix C.I The argument in Sec. 3is intuitive, and we do not prove theoretically.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide our experimental details in Sec. 5 App. For the new update algorithm, we provide full pseudocode in Sec. D and we will release code publicly as soon obtain approval to do so.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The datasets we used are publicly available. We will release the code publicly as soon as we obtain internal approval.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide these details in Sec.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification; The shaded area Figure Sinureaindscates the standard deviation. We only compute FID once due to cost constraints.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
27

---

Answer: [Yes]
Justification: We provide details in Sec.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /tte:://neuiipsp
Answer: [Yes]
Justification: The paper conforms with the NeurIPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We provide impact statement in Sec. Effectively, with other work
generative models, they can be used for both beneficial and harmful purposes. Our work, being focused on the mechanies of these models, does not introduce new risks, nor does mitigate existing ones.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We believe that work is not expected to have any more potential risks than other work in this field. We have discussed some of these considerations in Sec. [F]
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: For the datasets we use, we cite the original papers and describe the license information in Sec.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide necessary information reproduce new models in Sec.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing or research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
28


Reviews:
[{"review_id": "kUXqnJqv8p", "reviewer": "Reviewer_nW7F", "summary": "The method introduces an one-stage training of rectified flows, mitigating the costly process of multi-iteration training of the former model. Particularly, the authors propose a U-shaped timestep distribution for sampling and modified LPIPS-Huber loss. The method demonstrates superior FID scores for 1-NFE setting on CIFAR-10 and ImageNet $64\\times 64$, outperforming other distillation methods like consistency distillation.", "questions": "NA", "limitations": "This is included in conclusion.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Provide an analysis of reflow algorithm and shows that one-stage reflow are sufficient for nearly straight solution trajectory\n- Improve training of rectified flow via introducing U-shaped time distribution and new loss function, LPIPS-Huber.", "weaknesses": "- Have the authors tested with more than two-NFE like 100 to directly compare with 2-rectified flow? Does model still get consistent improvement?\n- What is the advantage of U-shape time distribution over lognormal distribution? As it is constrained to already hold a straight solution paths like in 2-rectified flow which cannot directly apply to 1-rectified flow. The intuition of this part does not clearly present.\n- Have the authors tested LPIPS-Huber-1/t (config F) for 1-rectified flow?\n- New sampler seems to have counter effect when many-NFE (>6) is used. Could the authors explain this behaviour?\n\nOverall, the novelty seems to be limited as the method mainly involves engineering works in using U-shape distribution and LPIPS-Huber loss (a combination of huber loss and LPIPS loss) to improve model performance under few-NFE settings. Besides, though the performance shows a notable gain compared to RectifiedFlow and ConsistencyModels, it still falls behinds ImprovedConsistencyModels and CTM. However, I appreciate the findings in paper and these proposed techniques can be useful for the community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's valuable comments.\n\n## Have the authors tested with more than two-NFE like 100 to directly compare with 2-rectified flow? Does model still get consistent improvement?\n\nWe tested up to 16 NFEs and observed consistent improvement (see Figure 4). Compared to 2-rectified flow, 2-rectified flow++ with 2 NFEs outperforms 2-rectified flow with 110 NFEs by a large margin (2.76 vs 3.36, Table 2).\n\n## What is the advantage of U-shape time distribution over lognormal distribution? The intuition of this part does not clearly present.\n\nLognormal emphasizes the middle, while U-shaped distribution emphasizes both ends of the interval. The intuition is that when training 1-rectified flow (or diffusion model), there is little to learn on both ends because when t=1, a model just learns the dataset average, and when t=0, it just predicts noise average (i.e., 0). In contrast, in 2-rectified flow, a model should learn to generate a very sharp sample at t=1, and a very accurate noise at t=0. Therefore, the most difficult parts are the both ends.\n\n## Have the authors tested LPIPS-Huber-1/t (config F) for 1-rectified flow?\n\nWe emphasize that our scope is to improve the performance of the 2-rectified flow. 1-rectified flow training is not intersection-free, so using loss functions other than the squared L2 violates the theory. Nonetheless, we tested LPIPS-Huber-1/t loss for 1-rectified flow and provide the result in the general response. We see that the result is pretty bad. These results confirm our unique insights in 2-rectified flow presented in Section 3.\n## New sampler seems to have counter effect when many-NFE (>6) is used. Could the authors explain this behaviour?\n\nThe goal here is to obtain as good FID as possible using as few NFE as possible. The new sampler's FID goes up after some NFE but **its performance with a low NFE is better than the original sampler's best FID with many more NFEs** (see Figure 4). We introduced the new sampler as a prototypical example to show that our 2-rectified flow++ can be further improved by designing a better sampler, which is another advantage over other distillation methods. It is certainly possible to design a better sampler that does not have this increasing FID behavior, and we leave this to future work.\n\n## 2-rectified flow++ falls behinds CTM and iCT.\n\nTable 2 shows that our method outperforms CTM without GAN loss on CIFAR-10 by a large margin. 2-rectified flow++ can also be trained with GAN loss and may obtain a similar gain. Combining two different generative models is an interesting but orthogonal direction to our paper’s scope.\n\nIt is true that our 2-rectified flow++ falls behind iCT (FIDs 4.32 vs 4.01 on ImageNet64), but we want to emphasize that 2-rectified flow++ does not require real data during training while iCT does. These days, many foundation models such as SDXL [2] only disclose model weights but not training data. iCT cannot be trained in those cases.\n\n\n[1] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow\n\n[2] SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis", "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": "01B3vV01ga", "reviewer": "Reviewer_eZ8c", "summary": "This paper mainly improves the training of rectified flows empirically, making it comparable to the distillation method in terms of performance with fewer steps.", "questions": "See Weaknesses.", "limitations": "The authors briefly discuss their limitations.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This article provides a comprehensive analysis of a single return, and there is a clear motivation for improvement.\n\n2. The authors mainly improved the training of the whole network from three aspects: time-step sampling, loss function, and initialization from the diffusion model.  \n\n3. The improved rectified flows method can be compared to the sota distillation method, and it also supports operations such as inversion, which may be important in applications such as image editing or translation.", "weaknesses": "1. The main weakness of this paper is that all the improvements are incremental and empirical. In addition, many of the improved techniques overlap with existing diffusion improvements, such as Pseudo-Huber loss [1] and Initialization with pre-trained diffusion models [2].\n\n2. The comparison in Section 5.2 may be somewhat unfair. Although the reflow requires a forward pass for each iteration, the reflow needs to run ODE in advance to generate noise-sample pairs, while the CD does not need to simulate, so the number of forward passes should also be considered to generate noise-sample pairs.       \n\n\n[1] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models.\n[2] Liu X, Zhang X, Ma J, et al. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's valuable comments.\n\n## The main weakness of this paper is that all the improvements are incremental and empirical.\n\nWe respectfully disagree with the reviewer's comment. Our proposed techniques improve the FID of 2-rectified flow **from 12.21 to 3.38 on CIFAR-10, 12.39 to 4.11 on AFHQ64, 8.84 to 5.21 on FFHQ64** compared to baseline, which is difficult to call incremental in any sense. Regarding our improvements being empirical, it is true that the paper is empirical at its core, motivated by the insights we provide in Sections 3 and 4. However, we disagree that this is a reason to reject. Indeed, many highly impactful papers in the area of generative modeling are empirical [3,4,5,6]. While we agree that having rigorous theoretical backing is more beneficial than not, we don’t believe it is always necessary or possible for all practical scenarios, especially those in deep learning.  \n\n## In addition, many of the improved techniques overlap with existing diffusion improvements, such as Pseudo-Huber loss and Initialization with pre-trained diffusion models\n\n**Pseudo-Huber loss**\n\nThe fact that Pseudo-Huber loss works well for consistency models does not necessarily mean that it should work well for rectified flows because of their fundamental differences in formulation. Consistency models' loss functions can be generalized to any metric function while rectified flow requires l2 squared loss (see the general response above). Otherwise, the obtained solution is not a posterior expectation and thus violates the marginal-preserving property of Liu et al. [1] (Theorem 3.3). Without our finding in Section 3, this generalization is not justified. Indeed, the 1-rectified flow’s FID gets worse if we use other losses, as we demonstrated in the supplementary results in the general response.\n\n**Pre-trained diffusion models**\n\nIt is true that the Instaflow paper used the pre-trained diffusion models for 1) paired data generation and 2) initialization, but they did not provide any theoretical justification for why we can do so. This raises the question of whether doing so yields the same results as standard 2-rectified flow training since diffusion models are trained with a curved, nonlinear interpolation (see Figure 5 in Liu et al. [1]). An illustrative example is when we use the pre-trained variance-exploding EDMs. In this case, it doesn’t make sense to directly use them as they are trained in t \\in [0, +\\infty) while rectified flows operate on [0, 1] range, so we can intuitively see that some sort of conversion is needed.\n\nIn fact, it turns out that one needs to convert the scale and time following Proposition 1 in our paper to make them compatible; the proof is provided in Appendix D. Proposition 1 makes sure that 1) the coupling generated from a pre-trained diffusion model is the same as 1-rectified flow coupling required by Reflow algorithm (Algorithm 1), and 2) 2-rectified flow is properly initialized by 1-rectified flow model as suggested by Liu et al. [1]. Table 1 in the attached pdf file shows that naively initializing with EDM without the proposed conversion leads to slower convergence.\n\n\nThis result is quite surprising, as it suggests that 1-rectified flow (with linear interpolation) and diffusion models (with nonlinear interpolation) not only belong to the same model class but are actually equivalent (i.e., they are interchangeable after training by simple time and scale conversion). We believe that this is not widely known in the literature, as significant community effort has been invested in developing new training techniques for 1-rectified flows/flow matching models, with the hope that the linear interpolation “connects data and noise on a straight line” and thus “has better theoretical properties” [2].\n\n\n## The comparison in Section 5.2 may be somewhat unfair. The number of forward passes should also be considered to generate noise-sample pairs.\n\nIt already does so. The caption of Table 5 says, \"Reflow uses 395M forward passes for generating pairs and 1, 433.6M for training.\". They sum up to 1,828.6M.\n\n[1] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow\n\n[2] Scaling rectified flow transformers for high-resolution image synthesis\n\n[3] Diffusion Models Beat GANs on Image Synthesis\n\n[4] Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks\n\n[5] Analyzing and Improving the Image Quality of StyleGAN\n\n[6] Improved Techniques for Training GANs", "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": "fVRoSovWKc", "reviewer": "Reviewer_aFYc", "summary": "The paper targets efficient training of a class of flow models called Rectified Flow (RF) trained using flow matching objective. The paper has two broad contributions: (1) justification of 2-RF (‘reflow’-ed once) being close to optimal, (2) and using those findings to improve training of 2-RF.\n\nAuthors argued that, in practical scenario, the pairs generated by optimal 1-RF model is ‘crossing-free’, i.e. the stochastic interpolating paths rarely cross each other. They provided an intuitive and empirical evidence of the same. They claimed this to be a motivating factor for two improvements of 2-RF model training — a new time-step distribution $p_t$ & the  a new distance measure for the regression task of diffusion objective.\n\nGood empirical performance is shown in terms of FID on several datasets and models. An ablation study is also done for the relevant part of the finding.", "questions": "1. Authors said “*any other premetric .. deviates from the posterior expectation*” — but is it really true ? Can you mathematically show why ?\n2. Is it possible to have trajectory crossings *at all* when samples are from $p_{xz}^2$ ? If no, can you prove it formally ?\n3. L99: “.. use a specific non-linear interpolation” — what is that exactly ? I thought the interpolation is still linear.\n4. Eq.4: How did you decompose the loss — can you show the steps ? And what is $\\bar{\\mathcal{L}}$ ?\n5. L65: “*training loss of 2-RF has zero lower bound*” — please clarify: Isn’t it true that *any* L2 loss has zero lower bound ? How does it matter whether trajectories cross or not ? Even the FM loss with independent coupling $p_{\\mathbf{xz}}^1$ has zero lower bound — no ?\n6. I think it is unclear which quantity the authors are arguing to be zero when trajectories do not cross. The term ‘curvature’ was used some times (L124) — what does that mean ? Can you write this object in mathematical terms ? Just curious, what is the equation that needs to be proved if one wants to show that no trajectories from 1-RF coupling cross each other (ideal case) ? (Related to Q2 above).", "limitations": "Some limitation are mentioned, which are reasonable.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The problem statement chosen and the arguments provided are quite credible. It is well known that, for RF to work well, one requires several reflows. Reflows are expensive, hence the contributions of the paper (if credible) can be significant.\n- A good intuitive analysis is done by the authors to justify their argument of 2-RF being near-optimal.", "weaknesses": "I like the overall outcome of the paper in terms of empirical performance. I also like the problem statement. However, some concerns remain. Following points explain the issues and points to some more questions in question section.\n\n### Major concerns\n\n- This is a major concern for me. The paper has two parts: (1) justifying that 2-RF is near optimal & providing intuitive and empirical evidence; (2) proposing two new measures for training improvement. I think, (1) & (2) both are overall correct on their own. I just don’t think that (1) is the right motivator for (2).\n- Continuing the above point, I felt that the rationale behind both ‘improvements’ are weakly connected to the empirical observations (in section 3). After all, what the authors ultimately proposed are ad-hoc and exists in the literature. (L162) “.. *focusing on the tasks where training loss is high ..*” is a known technique which authors admitted themselves. Using non-L2 loss is also not unheard of [1]. More importantly, I don’t think (not) using L2 loss has anything to do with your observation is section 3 (as argued in L195). The L2 loss has its origin in score-matching which yields theoretical benefits, but I suspect it is not necessary. Q1 is related to this.\n- Let’s talk about the observation of section 3 itself (i.e. 2-RF being optimal). Authors must be clear about whether they are making a theoretical assertion (L64-65) or just an empirical observation (they used words like “under *realistic* setting”, “*rarely* intersect each other”). If you are providing a *guarantee*, you must provide better formal proofs. To clarify, I think the argument provided in section 3 is indeed correct and seems reasonable. But if it is a guarantee, empirical evidence is not enough. Q2 is a related question.\n\n[1] “Improving Diffusion Models's Data-Corruption Resistance using Scheduled Pseudo-Huber Loss”, Kharpov et al., arXiv 2403.16728.\n\n### Presentations\n\n- The notation used in section 2 & 3 are confusing sometimes. There are three pairs of notation — $(\\mathbf{x}, \\mathbf{z})$, $(\\mathbf{x}_0, \\mathbf{z}_0)$ \\& $(\\mathbf{x}_1, \\mathbf{z}_1)$. I am confused about which is what. I would recommend authors to follow the same notation as the original RF paper by Liu et al.\n- Notations like $\\mathbf{z}\\_0 \\sim p\\_{\\mathbf{x}}$ (L86) are very confusing.\n- The factor $\\frac{t}{1-t}$ in Fig 2(a) (at the top left) should be $\\frac{1-t}{t}$, right ?\n- Fig. 2 caption says $\\mathbf{z}'' = \\mathbf{z} + (\\mathbf{x}' - \\mathbf{x}'')$ — where is the factor $\\frac{1-t}{t}$ ? Did you assume $t = 0.5$ ?\n- Again, it is hard to parse notations like $\\mathbf{x}, \\mathbf{x}', \\mathbf{x}''$ or $\\mathbf{z}, \\mathbf{z}', \\mathbf{z}''$.\n- Eq.4: Can you please denote the suffix of the $\\mathbb{E}_{??}\\left[ \\cdot \\right]$ properly ? It is hard to read otherwise.\n\n### Results/Experiments\n\nResult section is okay-ish. The following are some comments/questions.\n\n- Table.1: Need clarification: ‘Base (A)’ is a 2-RF, and the other ones are written as “(A) + <something>” — does that mean they are 3-RF (meaning 2 reflows) ?\n- Section 5.2 seems totally unnecessary. That has nothing to do with the core contributions of the paper.\n- Section 5.3 is also very much unnecessary. I even doubt its correctness. You seem to be proposing a new sampler/solver with a very weak (intuitive) motivation. Designing a solver require a lot more than that. And then “*.. detailed analysis is provided in appendix E*” — appendix E barely has any details !  Also, obvious question, why are FIDs going up (fig.4) with higher NFE ? Does that even make sense ?\n- Fig.5(b): The inverted noise norm distribution still looks quite different (higher variance) from the true noise. Just having the norm closer to the truth isn’t necessarily making a good case for your method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's valuable comments.\n## Why does the training loss of rectified flows (and FM with independent coupling) have a nonzero lower bound?\n\nWhen two interpolation trajectories cross, given the intersection, there are two possible directions to go. As the neural net we use is deterministic, if there are two possible regression targets, the loss cannot be zero. \n## How did you decompose the loss? What is $\\bar {\\mathcal L}$?\n\nAs we described in the general response above, the L2 squared loss is lower-bounded by $ \\frac{1}{t^2} \\mathbb E [|| \\mathbf x - \\mathbb E[\\mathbf x | \\mathbf x_t]||_2^2]$. $\\bar {\\mathcal L}$ is simply defined as the difference between $\\mathcal L$ and its theoretical lower bound (i.e., $\\bar {\\mathcal L}:= \\mathcal L - \\frac{1}{t^2} \\mathbb E [|| \\mathbf x - \\mathbb E[\\mathbf x | \\mathbf x_t]||_2^2]$), which is the only term we can reduce. The decomposition in Eq. 4 holds by definition.\n## “any other premetric .. deviates from the posterior expectation” — Why? \n\nSee general response. \n## Proposed techniques are ad-hoc, exist in the literature. “Focusing on the tasks where training loss is high” is a known technique.\n\n “Putting emphasis on tasks with high training loss” does not always help; knowing how to apply it requires care.\n\nE.g., one cannot naively apply this technique to 1-rectified flow training. Figure 1 in the attached PDF shows that the training loss of 1-rectified flow is also U-shaped as 2-rectified flow loss. However, the current state-of-the-art 1-rectified flow uses logit-normal distribution [4], which has the opposite shape.\n\nOn the other hand, for 2-rectified flow, we can design $p_t$ based on the loss shape as Sec. 3 tells us that the lower bound of the loss is nearly zero.\n## Sec. 3: If you are providing a guarantee, you must provide better formal proofs. \n\nWe are not making a formal statement—we do not expect our intuition to hold for every dataset/model. We found empirically that for many natural image datasets, our intuition appears to hold, allowing us to design significant training improvements.\n\nMaking section 3 formal poses several challenges. For example, we reason that if we add the difference of two synthetic data points to a noise, it is not a common noise realization that is used in the training. To make this formal, we would need to add restrictive assumptions to the data distribution, model, and training algorithm, as we are dealing with synthetic data. We also assumed that the model would not generate high-quality data if the input is very different from the noise it has seen during training, which is true in practical cases but also requires additional assumptions. As such, we ultimately decided to leave the intuition at an empirical level; we felt that adding a stylized theoretical result would not significantly strengthen the main results.\n## Q3: use a specific non-linear interpolation” — Isn’t the interpolation linear?\n\nThe variance-preserving diffusion model is a special case of rectified flows with the interpolation defined as $\\mathbf x_t = \\alpha(t) \\mathbf x + \\sqrt{1-\\alpha(t)} \\mathbf z$, where $\\alpha(t) = \\exp(-\\frac{1}{2} \\int_0^t (19.9s + 0.1)ds)$ (See Liu et al. [1]).\n\n## Q6: \n\n**What does “curvature” mean ?**\n\nWhen there is no intersection, the resulting (optimal) rectified flow's ODE trajectories are completely straight, or in other words, have zero curvature (Theorem 3.6 in Liu et al. [1]). Curvature [2,3] represents the degree to which the trajectory deviates from a straight path–one widely used definition is $\\int_0^1 \\|\\mathbf{z}_1-\\mathbf{z}_0-\\frac{\\partial}{\\partial t} \\mathbf{z}_t\\|_2^2 dt$.\n\n**What equation needs to be proved to show that no trajectories from 1-RF coupling cross each other?**\n\nTo show that there is no intersection, one needs to show that $\\mathbb E[\\mathbf x | \\mathbf x_t = (1-t)\\mathbf x' + t\\mathbf z'] = \\mathbf x'$ for all pairs of $\\mathbf x'$ and $\\mathbf z'$ and for all $t$.\n\n## Table.1:“(A) + something” — does that mean they are 3-RF (meaning 2 reflows) ?\n\nAll of our models are 2-rectified flows. “(A) + something” means we add the technique something to the base model configuration (A). We will clarify. \n\n## Sections 5.2 and 5.3 seem unnecessary. \n\nWe aim to draw attention to the potential of rectified flows. The computational cost to generate synthetic pairs is sometimes cited [5] as a downside, and we provide Section 5.2 as a counterargument. We introduced the new sampler in Section 5.3 to show that 2-rectified flow++ can be further improved by designing a better sampler–another advantage over other methods. \n\n## I doubt correctness of 5.3. Designing a solver require a lot more.\n\nWe did not mean to imply that our proposed update rule is itself an ODE solver, and will clarify. We use terms like “sampler” or “update rule” that **are applied to existing solvers like Euler and Heun**, instead of the term “solver” to refer to the new sampling algorithm. We do not even believe it converges to the true ODE solution in the limit as its behavior is very different from Euler solver’s in a large NFE setting. \n\nOur claim is that the new update rule achieves better FID with a smaller NFE (Fig. 4).\n\n## Why are FIDs going up (fig.4) with higher NFE ?\n\nThe new sampler's FID goes up after some NFEs because it is not a proper solver, but its performance at low NFEs is better than the original sampler's best FID with many more NFEs (Figure 4). \n## Fig.5(b): The inverted noise norm distribution still looks quite different from the true noise... Not necessarily a good case for your method.\n\nThe inverted noise norm distribution is not the same as the true noise distribution but closer to it than the baseline. Also, the norm distribution is meant to complement the visual comparison in Figure 5(c), where we see that the inverted noises of the baseline exhibit a strong spatial correlation. \n## Fig. 2 caption: where is the factor $\\frac{1-t}{t}$ ? Did you assume $t=0.5$?\n\nYes.", "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": "jcIfSFglp3", "reviewer": "Reviewer_Wooo", "summary": "By retraining with Rectified flows, they straighten the ODE, allowing sampling for small number of steps. They propose using rectified flow as replacement to complex distillation methods such as consistency models. They reflow only once to straighten the path which makes it as efficient as distillation. They propose some propose using a different distribution on t and a LPIPS loss. They obtain better performance than the original rectified flow and competitive results with distillation.\n\nThe authors provide geometric intuition to why one reflow should be enough. They provide well detailed ablation on each aspects of their improvements on multiple datasets.\n\nThey combine Pseudo-Huber (differentiable Huber loss to get less sensitivity to outliers while being differentiable) with LPIPS (better distance metric than L2 for image quality). They take a weighted combination of both losses. If one was not generating images, we could envision using pseudo-huber with L2, so the method could still be useful in that case. The focus here is on images though.\n\nEdit: I apologize, I added some extra text to the summary that was not related. I just removed it.", "questions": "The only thing I would like added is the difference in memory from the first flow and then the reflow only because of the use of the LPIPS. Basically, I would like to know if LPIPS massively improve the memory cost over L2 and if so it would be important to mention it as a limitation. The authors already mentioned the 15% speed reduction (which is quite minor), it would be great to know about memory.", "limitations": "The authors have well addressed the limitations. I just asked that they talk about the memory cost of LPIPS too.", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The experiments are on multiple datasets, focusing on the 1-2 NFE case which is arguably the most important problem in diffusion-flow models to solve. The results are quite impressive. Reflow seems more efficient that distillation and its especially much more clean and natural as a solution. The experiments and discussion is quite thorough.", "weaknesses": "Honestly this is a near perfect in my view. Maybe, the only thing I could think of is the focus being only on images, so maybe not fully generalizable accross domains.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review. On ImageNet 64x64, the additional memory overhead of using LPIPS is less than 5%. We expect this would be even smaller on more large-scale settings since the generative model’s size would be relatively huge compared to the feature extractor such as AlexNet or VGG.", "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:
# General response

We thank the reviewers for their valuable comments. Here, we provide an additional background to help clarify some of the points raised in the reviews. We also have fixed some typos and clarified notations. The supplementary PDF file is attached.

**Question: Clarify notation**

In our paper, we use $\mathbf x, \mathbf z$ as the random variables (data and noise, respectively), $\mathbf x_t$ as a function of them, and $\mathbf x',\mathbf x'’, \mathbf z',\mathbf z'’$ as the specific values of them. $\mathbf z_t$ is a dynamical system governed by the rectified flow ODE. $\mathbf z_0 \sim p_{\mathbf x}$ means that the initial value of $\mathbf z_t$ at t=0 is sampled from $p_{\mathbf x}$.

We will make the notations clear in the revised version.

**Question: What does your intuitive argument in Section 3 have to do with the proposed loss functions? Alternative losses are widely used in many generative models.**

We would like to summarize the logic below. 

**Minimizer of L2 squared loss**

For a random variable $\mathbf x$, $\arg \min_{\hat {\boldsymbol \mu}} \mathbb E_{\mathbf x}[||\mathbf x - \hat {\boldsymbol \mu}||^2]$ = $\mathbb E_{\mathbf x}[\mathbf x]$, as $\nabla_{\hat {\boldsymbol \mu}} \mathbb E_{\mathbf x}[||\mathbf x - \hat {\boldsymbol \mu}||^2] = 2(\hat {\boldsymbol \mu} - \mathbb E_{\mathbf x}[\mathbf x]) = 0$ at $\hat {\boldsymbol \mu} = \mathbb E_{\mathbf x}[\mathbf x]$. This minimum mean squared error (MMSE) estimator minimizes the l2 squared loss, but it is generally not a minimizer of other loss functions. For example, the minimizer of the l1 loss is the median. See e.g., [2], p. 176.

**Lower bound.**  We can see that the l2 squared loss cannot be zero and is lower bounded by $\mathbb E[||\mathbf x - \mathbb E_{\mathbf x}[\mathbf x]||^2]$ unless $\mathbf x$ is not random.

**Posterior expectation and marginal-preserving property**

The goal of the rectified flow training is to obtain the vector field $\mathbb E[\mathbf z - \mathbf x |\mathbf x_t =  \mathbf z_t]$ assuming the linear interpolation is used, because this vector field is shown to generate the same marginal distribution as the interpolation (Liu et al. [1], Theorem 3.3). This is called the “marginal preserving property”, and it is central to the efficacy of rectified flows. This is why Liu et al. [1] used the L2 squared loss for training. If we use a different loss function, the obtained solution is not a posterior expectation and the marginal-preserving property (Theorem 3.3) is inapplicable.

**Our finding**

Since we found that the intersection is generally not a problem in 2-rectified flow training (Section 3), our regression target is no longer a random variable. Now $\mathbb E[\mathbf x | \mathbf x_t = (1-t)\mathbf x' + t \mathbf z'] = \mathbf x'$, and the loss can actually become zero: 

$$\mathbb E_{\mathbf x, \mathbf z, t}[\frac{1}{t^2} || \mathbf x - \mathbf x_\theta((1-t)\mathbf x + t\mathbf z, t)||_2^2] \geq \mathbb E || \mathbf x - \mathbb E[\mathbf x | \mathbf x_t]||_2^2] = \mathbb E [\frac{1}{t^2} || \mathbf x - \mathbf x||_2^2] =  0$$

, where the optimum is achieved at $\mathbf x_\theta((1-t)\mathbf x' + t \mathbf z', t) = \mathbf x'$. Since this is also the unique optimum of any premetric $m$ such that $m(\mathbf{a}, \mathbf{b})=0 \iff \mathbf{a}=\mathbf{b}$, we can generalize the rectified flow loss to any premetric.

**What if we use different losses for 1-rectified flow training?**

Upon the reviewers’ questions, we provide additional experiment results where we compare the FID of 1-rectified flow models trained with different losses.

| Method        | FID  |
| ------------- | ---- |
| **L2 squared**| **2.66** |
| LPIPS         | 47.24 |
| Pseudo-Huber  | 2.92 |

*Table: FID of 1-rectified flow trained with different losses on CIFAR-10. During sampling, the RK45 solver is used following Liu et al. [1]*


This result is expected–Unlike 2-rectified flow, 1-rectified flow training is not intersection-free, so only the L2 squared loss is valid.

**Typos**

Section 3: $p^1(\mathbf x) = \int p^1(\mathbf x, \mathbf z) d\mathbf z$ -> $p^2(\mathbf x) = \int p^2(\mathbf x, \mathbf z) d\mathbf z$

Figure 2(a) (top left): $\frac{t}{1-t}\mathbf x' - \mathbf x''$ -> $\frac{1-t}{t}(\mathbf x' - \mathbf x'')$

Section 4: log-normal -> logit-normal

[1] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

[2] https://probml.github.io/pml-book/book1.html

Meta-review: The paper presents an improved training approach for rectified flows, demonstrating competitive performance with state-of-the-art distillation methods, particularly in low NFE settings. Overall, the reviewers recognized the paper’s empirical contributions as significant and practically relevant. Reviewer Wooo commended the results, noting that “Reflow seems more efficient than distillation and is especially much more clean and natural as a solution.” Reviewer nW7F observed that improvements are consistent even when scaling to higher NFEs, emphasizing that “the proposed techniques can be useful for the community.” While some concerns were expressed about the incremental nature of the work and overlap with existing methods—such as those highlighted by Reviewer eZ8c, who pointed out that “many of the improved techniques overlap with existing diffusion improvements”—the consensus is that the paper’s contributions substantially advance the training of rectified flows by reducing complexity and enhancing efficiency. The presentation and clarity of the paper were also praised, with Reviewer Wooo describing the empirical studies as “thorough” and the presentation as “excellent.” Given the paper’s strong empirical performance, practical impact, and clear exposition, I recommend its acceptance.
Decision: Accept (poster)