Title: Improving the Training of Rectified Flows
Venue: NeurIPS 2024
Accepted: yes
Keywords: generative modeling

rectified flow

diffusion model
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.

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 LPIPSHuber 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 × 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 .

1 Introduction

Diffusion models [Sohl-Dickstein et al., 2015, Ho et al., 2020, Song and Ermon, 2019, Song et al., 2020b] have shown great promise in image [Ramesh et al., 2022] and video [Ho et al., 2022] generation. They generate data by simulating a 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 of 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 et al., 2022]).

In knowledge distllation-based methods [Luhman and Luhman, 2021, Salimans and Ho, 2022, Song et al., 2023, Zheng et al., 2022] a student model is trained to directly predict the solution of the PF-ODE. These models are currently state-of-the-art in the low number of function evaluations (NFEs) regime (e.g. 1-4).

Another promising direction is simulation-free flow models such as 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 a special case. Rectified flows can learn smooth

Giulia Fanti

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 to 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 et al., 2022] rectified flows have several advantages. First, they can be generalized to map two arbitrary distributions to one another, while distillation methods are limited to a 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., 2022, Kim et al., 2022, Wallace et al., 2023, Couairon et al., 2022, Mokady et al., 2023, Su et al., 2022, Hong et al., 2023] and watermarking [Wen et al., 2023]. 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-based 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-the-art techniques for training rectified flows use the Reflow algorithm to improve low NFE performance [Liu et al., 2022, 2023]. 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 is 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 ℓ 2 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 as EDM [Karras et al., 2022], 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 64 × 64 [Deng et al., 2009]), our improved rectified flow outperforms the state-of-the-art distillation 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-image translation.

2 Background

2.1 Rectified Flow

Rectified flow (see also flow matching [Lipman et al., 2022] and stochastic interpolant [Albergo and Vanden-Eijnden, 2022]) is a generative model that smoothly transitions between two distributions p x and p z by solving ODEs [Liu et al., 2022]. For x ∼ p x and z ∼ p z , we define the interpolation between x and z as x t = (1 -t ) x + t z for t ∈ [0 , 1] . Liu et al. [2022] showed that for z 0 ∼ p x , the

<latexi sh

1_b

64="

d

q

mX

J

v7oyT

2

R

3QELM

>A

C

c

ZD

S

N

FI

P

W

9

U

8

Y

u

gr

G

H

fj

+

K

w

/

5pz

n

0

V

B

<latexi sh

1_b

64="Wv+gCX

n

cAZV

k

J

>

B

H

7T

F

3

q

UY

oQr

G

I9

p

E

2

P

8

0D

S

d

5

R

z

u

m

f

K

/

y

<latexi sh

1_b

64="knB

G3mO

Kgv7

8+

of

Q

r

>A

c

VDLS

M

F

U

2

q

w

jP

0

C

J

Z

5I

H

u/Rp

d

E

N

X

9W

Y

T

z

y

<latexi sh

1_b

64="oC

2OZz80yY

RX

gPp7

kFM

>A

B9H

c

VDLS

T

W+q

d

Iv

r

j

K

U

J

3N

G

u/

w

f

E

5

m

Q

n

<latexi sh

1_b

64="yMWzg c/

Ep

QC

O

>A

B+H

V

7T

qj

UY

J

o

r

G

X

2

P

8

0D

kS

dn

5

R

u

m

f

K

<latexi sh

1_b

64="oC

2OZz80yY

RX

gPp7

kFM

>A

B9H

c

VDLS

T

W+q

d

Iv

r

j

K

U

J

3N

G

u/

w

f

E

5

m

Q

n

<latexi sh

1_b

64="

d

q

mX

J

v7oyT

2

R

3QELM

>A

C

c

ZD

S

N

FI

P

W

9

U

8

Y

u

gr

G

H

fj

+

K

w

/

5pz

n

0

V

B

<latexi sh

1_b

64="yMWzg c/

Ep

QC

O

>A

B+H

V

7T

qj

UY

J

o

r

G

X

2

P

8

0D

kS

dn

5

R

u

m

f

K

<latexi sh

1_b

64="knB

G3mO

Kgv7

8+

of

Q

r

>A

c

VDLS

M

F

U

2

q

w

jP

0

C

J

Z

5I

H

u/Rp

d

E

N

X

9W

Y

T

z

y

<latexi sh

1_b

64="oC

2OZz80yY

RX

gPp7

kFM

>A

B9H

c

VDLS

T

W+q

d

Iv

r

j

K

U

J

3N

G

u/

w

f

E

5

m

Q

n

<latexi sh

1_b

64="

DzAQ+

0NEruBF

U

M

I

>

H

c

VC7T

w

L

yW

R

S8

o

9

G

q

PZ

J

/

5

v

gO

n

3

f

2

X

m

k

K

<latexi sh

1_b

64="oC

2OZz80yY

RX

gPp7

kFM

>A

B9H

c

VDLS

T

W+q

d

Iv

r

j

K

U

J

3N

G

u/

w

f

E

5

m

Q

n

E

[

x

|

x

t

]

p

0

x

p

1

x

p

z

p

z

p

z

p

z

p

1

x

p

2

x

(a) Linear interpolation

induced by

𝑝

𝒙𝒛

#

(c) Linear interpolation

induced by

𝑝

𝒙𝒛

$

(b) Generative process

of 1-rectified flow

(d) Generative process

of 2-rectified flow

x

t

x

t

E

[

x

|

x

t

]

following ODE yields the same marginal distribution as x t for any t :

Since x 1 = z , Eq. (1) transports p x to p z . We can also transport p z to p x by drawing z 1 from p z and solving the ODE backwards from t = 1 to t = 0 . During training, we estimate the conditional expectation E [ x | x t = z t ] with a vector-valued neural network x θ trained on the squared ℓ 2 loss:

where p xz is the joint distribution of x and z , x θ is parameterized by θ , and ω ( t ) is a weighting function. p t is chosen to be the uniform distribution on [0 , 1] in Liu et al. [2022, 2023]. In the optimum of Eq. (2), x θ becomes the conditional expectation as it is a minimum mean squared error (MMSE) estimator, which is then plugged into the ODE (1) to generate samples. Instead of predicting the conditional expectation directly, Liu et al. [2022] choose to parameterize the velocity v t with a neural network v θ and train it on

which is equivalent to Eq. (2) with ω ( t ) = 1 t 2 . See Appendix. A.

In this paper, we consider the Gaussian marginal case, i.e., p z = N ( 0 , I ) . In this case, if we define x and z as independent random variables (i.e., p xz ( x , z ) = p x ( x ) p z ( z ) ) and use a specific nonlinear interpolation instead of the linear interpolation for x t , Eq. (2) becomes the weighted denoising objective of the diffusion model [Vincent, 2011], and Eq. (1) becomes the probability flow ODE (PF-ODE) [Song et al., 2020b].

2.2 Reflow

Algorithm 1 Reflow Procedure

1: First iteration: 2: θ 1 = arg min θ E x , z ∼ p 0 xz E t ∼ p t [ ω ( t ) ∥ x -x θ ( x t , t ) ∥ 2 2 ] ▷ Train 1-rectified flow 3: for k = 1 to K -1 do 4: T k ( z ) = z + ∫ 0 1 1 t ( z t -x θ k ( z t , t )) dt with z 1 = z 5: p k xz ( x , z ) = p z ( z ) δ ( x -T k ( z ) ) ▷ Generate synthetic pairs for next coupling 6: θ k +1 = arg min θ E x , z ∼ p k xz E t ∼ p t [ ω ( t ) ∥ x -x θ ( x t , t ) ∥ 2 2 ] ▷ Train ( k +1) -rectified flow 7: end for

The independent coupling p xz ( x , z ) = p x ( x ) p z ( z ) is known to lead to curved ODE trajectories, which require a large number of function evaluations (NFE) to generate high-quality samples [Pooladian et al., 2023, Lee et al., 2023]. Reflow [Liu et al., 2022] is a recursive training algorithm to find a better coupling that yields straighter ODE trajectories. Starting from the independent coupling p 0 xz ( x , z ) = p x ( x ) p z ( z ) , the Reflow algorithm generates p k +1 xz ( x , z ) from p k xz ( x , z ) by first

<latexi sh

1_b

64="v

0Q

2/jVR

BXL

+

CNUJpH

Kg

>A

c

D

S

M

F

3

Zdu

k

Y

G

o

9

y

I

O

5

7

r

n

f

q

m

w

z

8

P

ET

W

<latexi sh

1_b

64="X30HuBYG2o/EZAWFDM

VRj

>

8

cd

LSg

U

vqk

w

K

m

T

Ny

r

I

Q

J

C

+

n

z

P

9

O

f

p

5

7

<latexi sh

1_b

64="qvwG

D9X

2

z

UkVL

7KC

/

0

>A

B8

c

SgM

F

T

W+

d

I

o

y

r

Rj

H

Jp

3N

EO

Q

Z

P

n

f

u

m

Y

5

<latexi sh

1_b

64="

3dXOjEGMIkU

/ZTn

PV7A

>

B9H

c

C

wFL0pr

J

YW

q

oQ

mv

R

+

g2

z

u

f

K

8

5

S

D

N

y

<latexi sh

1_b

64="d7jAM

I8kN

G/Tf

>

B

X

c

VDLSg

F

W+q

y

r2

R

K9

H

U

p

E

Q

3

0

Z

w

P

n

C

Y

5

m

<latexi sh

1_b

64="MIOw

f

8

K

WRm3H

C

BD

Q

>A

X

cdV

LSg

F

rPU

kZ

Ey

N

z0

7o

u

2q

n

/

9

+j

5

v

G

p

Y

J

T

<latexi sh

1_b

64="X30HuBYG2o/EZAWFDM

VRj

>

8

cd

LSg

U

vqk

w

K

m

T

Ny

r

I

Q

J

C

+

n

z

P

9

O

f

p

5

7

<latexi sh

1_b

64="+

H8

9Vd

BKAy

2oX0Cc rZTnY

>

3

DLS

N

F

p

f

u

vU

E

7

m

M

g

R

G

J

/

P

wO

j

5

W

IQ

k

z

q

<latexi sh

1_b

64="NvVOC

X

nG

YU2mzEd

A9

p

0

>

B/3

c

DLS

F

r

f

8

u

qJ+

oW

gRK

H

M

T

Q

7

y

w

ZP

k

j

I

5

x

′

M

x

z

′

x

′′

x

′

-

x

′′

Generative process

(1-Rectified Flow)

‖

‖

(a)

(c)

(d)

(b)

Generated from

𝐳 + 𝐱

!

-𝐱′′

Generated from

𝐳

z

′′

z

′′

=

z

′

+ 1

-

t

t

(

x

′

-

x

′′

)

09

2500

z + (x'-x")

50

I|zll2

2000

Z

Frequency

40

1500

30

20

10

500

60

70

Norm

08

06

100

0

0.00 0.36 0.72 1.08 1.44 1.80

Autocorrelation(lag=3)

generating synthetic ( x , z ) pairs from p k xz , then training rectified flow on the generated synthetic pairs (Figure 1 ( 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. 1.

Convergence: Liu et al. [2022] show that Reflow trajectories are straight in the limit as K →∞ . 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 be applied at 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].

3 Applying Reflow Once is 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 1-rectified flow-generated pairs do not intersect, or equivalently, E [ x | x t = (1 -t ) x ′ + t z ′ ] = x ′ for all pairs ( x ′ , z ′ ) [Liu et al., 2022].

To begin with, consider the manifold M x of the synthetic distribution p 1 ( x ) = ∫ p 1 ( x , z ) d z . 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 two pairs ( x ′ , z ′ ) and ( x ′′ , z ′′ ) intersect if ∃ t ∈ [0 , 1] s.t. (1 -t ) x ′ + t z ′ = (1 -t ) x ′′ + t z ′′ . For example, in Figure 2(a), we observe that the two trajectories intersect at an intermediate t .

For an intersection to exist at t it must hold that 1) 1-rectified flow maps z ′′ to x ′′ , and 2) z ′′ = z ′ + 1 -t t ( x ′ -x ′′ ) by basic geometry. However, note that for realistic data distributions and if

1-rectified flow is sufficiently well-trained, z ′′ = z ′ + 1 -t t ( x ′ -x ′′ ) is not a common noise realization (e.g., it is likely to have nonzero autocorrelation or a norm that is too large to be on a Gaussian annulus), as shown visually in Figure 2(a). As 1-rectified flow is almost entirely trained on common Gaussian noise inputs, it cannot generally map an atypical z ′′ to M x . Figure 2(c) shows qualitatively that if we draw values of z ′′ by first drawing z ′ ∼ N (0 , I ) 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 2(d) shows that the constructed noise vectors z ′′ have higher autocorrelation than expected. As a result, Figure 2(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 t = (1 -t ) x ′ + t z ′ ] ≈ 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. 4.

Edge cases: Note that if || x ′ -x ′′ || 2 is small, 1-rectified flow could map z ′′ to some point on M x . However, it does not alter the conclusion because the average of x ′ and x ′′ is close to x ′ anyway, so E [ x | x t = (1 -t ) x ′ + t z ′ ] ≈ x ′ . Similarly, if t is close to 1, 1 -t t ( x ′ -x ′′ ) ≈ 0 , so 1-rectified flow can map z ′′ to M x . If the 1-rectified flow is L -Lipschitz, || x ′ -x ′′ || 2 ≤ L || z ′ -z ′′ || 2 . Therefore, the expectation E [ x | x t ] 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 t , and the distribution 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 on easy tasks. One common approach is to focus on the tasks where the training loss is high [Shrivastava et al., 2016]. However, the training error of rectified flows is not a 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:

The first term does not depend on θ and thus cannot be reduced. The second term represents the actual minimizable training error, but its value cannot be directly observed because the first term is usually 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 L ( θ , t ) for designing the timestep distribution.

Figure 3 shows that the training loss of 2-rectified flow is large at each end of the interval t ∈ [0 , 1] and small in the middle. We thus propose to use a Ushaped timestep distribution for p t . Specifically, we define p t ( u ) ∝ exp( au ) + exp( -au ) on u ∈ [0 , 1] . We find that a = 4 works well in practice (Table 1). Compared to the uniform timestep distribution (config B), the U-shaped distribution (config C) improves

the FID of 2-rectified flow from 7 . 14 to 5 . 17 (a 28% improvement) on 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 1-rectified flow training, p t was chosen to be the uniform distribution [Liu et al., 2022, 2023] or logit-normal distribution [Esser et al., 2024] which puts more emphasis on the middle of the interval. When training 1-rectified flow, a model learns to simply output the dataset average when t = 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 t = 1 and the noise from the data at t = 0 , which are nontrivial tasks. Therefore, the U-shaped timestep distribution is more suitable for 2-rectified flow.

4.2 Loss function

Previously, the squared ℓ 2 distance was used as the training metric for rectified flow to obtain the MMSEestimator E [ x | x t ] . However, as we have shown in Sec. 3 that E [ x | x t = (1 -t ) x ′ + t z ′ ] ≈ x ′ , we can generalize Eq. (2) or equivalently Eq. (3) to any premetric m (i.e. m ( a , b ) = 0 ⇔ a = b ):

Note that without the intuition in Sec. 3, only the squared ℓ 2 distance would have been a valid premetric , as any other premetric makes the model deviate from the intended optimum (the posterior expectation E [ x | x t ] ). 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 ℓ 2 distance, we consider the following premetrics:

Pseudo-Huber [Charbonnier et al., 1997, Song and Dhariwal, 2023]: m hub ( z -x , v θ (( x t , t )) = √ || z -x -v θ ( x t , t ) || 2 2 + c 2 -c , where c = 0 . 00054 d with d being data dimensionality.

LPIPS-Huber1 t : m lp-hub1 t ( z -x , v θ ( x t , t )) = (1 -t ) m hub ( z -x , v θ ( x t , t )) + 1 t LPIPS ( x , x t -t · v θ ( x t , t )) ,

LPIPS-Huber: m lp-hub ( z -x , v θ ( x t , t )) = (1 -t ) m hub ( z -x , v θ ( x t , t )) + LPIPS ( x , x t -t · v θ ( x t , t )) , where LPIPS ( · , · ) is the learned perceptual image patch similarity [Zhang et al., 2018].

The Pseudo-Huber loss is less sensitive to the outliers than the squared ℓ 2 loss, which can potentially reduce the gradient variance [Song and Dhariwal, 2023] and make training easier. In our initial experiments, we found that the Pseudo-Huber loss tends to work better than the squared ℓ 2 loss with a small batch size (e.g. 128 on CIFAR-10). When the batch size is sufficiently large, it performs on par with the squared ℓ 2 loss on CIFAR-10 and FFHQ-64 and outperforms it on AFHQ-64, as shown in Table 1. As it is less sensitive to the batch size, we choose to use the Pseudo-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 a premetric as two different

points could have zero LPIPS if they are perceptually similar, we use it in combination with the Pseudo-Huber loss with the weighting 1 -t , thereby relying more on LPIPS when t is close to 1 where the task is more challenging. Note that in m lp-hub, the gradient vanishes when t is close to zero. To compsensate, we experiment with m lp-hub1 t where we multiply LPIPS by 1 t . Compared to config D, the LPIPS-Huber loss improves the FID of 2-rectified flow from 5 . 24 to 3 . 38 (a 35% 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, as seen in Table 1.

4.3 Initialization with pre-trained diffusion models

Training 1-rectified flow from scratch is computationally expensive. Recently, Pokle et al. [2023] showed that pre-trained diffusion models can be used to approximate E [ x | x t = z t ] in Eq. (1) by adjusting the signal-to-noise ratio. The following proposition is 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.1.

Proposition 1 Let p RE ( x | x t , t ) be the posterior distribution of the perturbation kernel N ((1 -t ) x , t 2 I ) . Also, let p VP ( x | x t , t ) and p VE ( x | x t , t ) be the posterior distributions of N ( α ( t ) x , (1 -α ( t )) 2 I ) and N ( x , t 2 I ) , each. Then,

where s VP and s VE are the scaling factors and t VP and t VE 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 C for derivation.

Proposition 1 allows us to initialize the Reflow with the pre-trained diffusion models such as EDM [Karras et al., 2022] or DDPM [Ho et al., 2020] and use Table 2 to adjust the time and scaling factors.

Starting from the vanilla 2-rectified flow setup [Liu et al., 2022] (config A), we initialize 1-rectified flow with the pre-trained EDM (VE). We also increase the batch size from 128 to 512 on CIFAR-10 compared to Liu et al. [2022]. Overall, these improve the FID of 2-rectified flow from 12 . 21 to 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 we can use real data if it is available. To see the effects of incorporating real data, we integrate the generative ODE of 1-rectified flow backward from t = 0 to t = 1 using an NFE of 128 to collect 50 , 000 pairs of (real data, synthetic noise) on CIFAR-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 on CIFAR-10 (config G).

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., p = 0 ) performs

2

3.64

The red rows correspond to the top-5 baselines for the 1-NFE setting, and the blue rows correspond to the top 5 baselines for the 2-NFE setting. The lowest FID scores for 1-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].

5 Experiments

We call these combined improvements to Reflow training 2-rectified flow++ and evaluate it on four datasets: CIFAR-10 Krizhevsky et al. [2009], AFHQ Choi et al. [2020], FFHQ Karras et al. [2019], and ImageNet Deng et 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 E.

5.1 Unconditional and class-conditional image generation

In Tables 3 and 4, we compare 2-rectified flow++ with the state-of-the-art methods on CIFAR-10 and ImageNet 64 × 64 . We observe two main messages:

On both datasets, 2-rectified flow++ (ours) outperforms or is competitive with SOTA baselines in the 1-2 NFE regime. On CIFAR-10 (Table 3), our 2-rectified flow achieves an FID of 3 . 07 in one step, surpassing existing distillation methods such as consistency distillation (CD) [Song et al., 2023], progressive distillation (PD) [Salimans and Ho, 2022], diffusion model sampling with neural operator (DSNO) [Zheng et al., 2022], and TRAnsitive Closure Time-distillation (TRACT) [Berthelot et al., 2023]. On ImageNet 64 × 64 (Table 4), 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

(a) CIFAR-10

(b) AFHQ

64

×

64

(c) FFHQ

64

×

64

3.4

Euler

4.0

Euler

5.6

Euler

3.2

Heun

Euler + new sampler

3.8

Heun

Euler + new sampler

5.4

Heun

Euler+ newsampler

3.0

Heun + new sampler

3.6

Heun + new sampler

5.2

Heun+newsampler

FID

2.8

FID

3.4

5.0

4.8

3.2

4.6

2.6

3.0

4.4

2.4

2

4

6

8

NFE

10

12

14

16

2.8

2

4

6

8

NFE

10

12

14

16

4.2

N

4

6

8

NFE

10

12

14

16

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 we could further reduce the gap by using config G in Tab. 1. Uncurated samples of our model are provided in Appendix. G.

2-rectified flow++ reduces the FID of 2-rectified flows by up to 75%. Compared to vanilla rectified flows [Liu 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 1 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 an 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 as it requires generating synthetic pairs before training. However, CD requires 4 (1 for student, 1 for teacher, and 2 for Heun's solver) forward passes for each training iteration, and CT requires 2 (1 for student and 1 for teacher) forward passes, while Reflow requires only 1. For example, in our ImageNet experiment setting, the total number of forward passes for Reflow is 395 M+ 1433 . 6 M= 1828 . 6 M(395Mfor generating pairs

and 1,433.6M for training), while the total numbers of forward passes for CD and CT would be 1 , 433 . 6 · 4 = 5 , 734 . 4 Mand 1 , 433 . 6 · 2 = 2 , 867 . 2 Munder the same setting. See Table 6 for the comparison. Moreover, generating pairs is a one-time cost since we can reuse the pairs for multiple training runs.

In terms of the storage cost, the synthetic images for ImageNet 64 × 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 a neural ODE, and its outputs approach the true solution of the ODE as NFE increases (i.e., precision grows). Figure 4 shows that with the standard Euler solver, FID decreases as NFE increases on all 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.

(a) Reconstruction loss

(b) Distribution of

∥

z

∥

2

2

Ours

EDM

Data

Noise

Reconstructed

(c) Inversion/reconstruction results

0.225

(MSE)

0.200

0.175

Ours

EDM

Ours

EDM

True noise

0.150

error

0.125

Recon

0.100

0.075

0.050

0.025

0.000

2

NFE

8

16

500

1000

1500

IIz|12

2000

2500

3000

(a) Interpolation

Lion→Tiger

Tiger→Lion

(b) Image translation

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 in the backward direction. In diffusion models, inversion has been used for various applications such as image editing [Hertz et al., 2022, Kim et al., 2022, Wallace et al., 2023, Su et al., 2022, Hong et al., 2023] and watermarking [Wen et al., 2023], but it usually requires many NFEs. Figure 5 (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), we compare the quality of the inverted noise, where we find that the noise vectors of 2-rectified flow are more Gaussian-like 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 6, we show two applications of inversion: interpolating between two real images (a) and image-to-image 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 ( ≥ 100 ) [Hong et al., 2023].

6 Conclusion

In this work, we propose several improved training techniques for rectified flows, including the Ushaped timestep distribution and LPIPS-Huber loss. We show that by combining these improvements, 2-rectified flows++ outperforms the state-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-art 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.

Acknowledgments

This work was 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 at 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 S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571 , 2022.

Anonymous. Balanced conic rectified flow. In Submitted to The Thirteenth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=ctSjIlYN74 . 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 arXiv:2303.04248 , 2023.

Timothy J Boerner, Stephen Deems, Thomas R Furlani, Shelley L 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:1809.11096 , 2018.

Pierre Charbonnier, Laure Blanc-Féraud, Gilles Aubert, and Michel Barlaud. Deterministic edgepreserving 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 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-based semantic image editing with mask guidance. arXiv preprint arXiv:2210.11427 , 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 on 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 distillation 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. Promptto-prompt image editing with cross attention control. arXiv preprint arXiv:2208.01626 , 2022.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458 , 2022.

Seongmin Hong, Kyeonghyun Lee, Suh Yoon Jeon, Hyewon Bae, and Se Young Chun. On exact inversion of dpm-solvers. arXiv preprint arXiv:2311.18387 , 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-12114, 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, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusionbased generative models. arXiv preprint arXiv:2206.00364 , 2022.

Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, 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 on Computer Vision and Pattern 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 arXiv:2301.12003 , 2023.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv: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.06380 , 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 arXiv:2101.02388 , 2021.

Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diffinstruct: 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 arXiv:1710.03740 , 2017.

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-8171. PMLR, 2021.

Ashwini Pokle, Matthew J Muckley, Ricky TQ Chen, and Brian Karrer. Training-free linear image inversion via flows. arXiv preprint 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 arXiv:2304.14772 , 2023.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical textconditional 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 arXiv:2202.00512 , 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 object detectors with online hard example mining. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 761-769, 2016.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. 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.14189 , 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 P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 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:11287-11302, 2021.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Bram Wallace, Akash Gokul, and Nikhil Naik. Edict: Exact diffusion inversion via coupled transformations. In Proceedings of the IEEE/CVF Conference on 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.

Yuxin Wen, John Kirchenbauer, Jonas Geiping, and Tom Goldstein. Tree-ring watermarks: Fingerprints 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 as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 586-595, 2018.

Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. arXiv preprint arXiv:2211.13449 , 2022.

3000

2500

Z

equency

2000

1500

1000

500

0

0.00.5

1.0

1.5

2.0

2.5

Autocorrelation (lag=3)

(a) Reconstruction loss

(b) Distribution of ∥ z ∥ 2 2

Z-generated

z + x' - x"-generated

A Equivalence of v -parameterization and x -parameterization

Given x t = (1 -t ) x + t z and z -x = ( x t -x ) /t , the equivalence of the x -parameterization (Eq. (2)) and the v -parameterization (Eq. (3)) can be shown in the following way. The result is borrowed from the Appendix of Lee et al. [2023], and we provide here for completeness.

This is equivalent to Eq. (2) with ω ( t ) = 1 /t 2 .

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 2(d) is on FFHQ-64. Figure 7(b) shows that the inception features of the samples in Figure 2(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 a d -dimensional vector u , we define the autocorrelation as:

where u k is the k -th element of a vector u and l > 0 represents the lag.

C Initialization Results

C.1 Proof of Proposition 1

Consider two perturbation kernels p t ( x t | x ) = N ( s ( t ) x , σ ( t ) 2 I ) and p ′ t ( x t | x ) = N ( s ′ ( t ) x , σ ′ ( t ) 2 I ) :

Let t ′ ( t ) be such that s ( t ) σ ( t ) = s ′ ( t ′ ) σ ′ ( t ′ ) . We will show that p t ( x | x t ) = p ′ t ′ ( x | s ′ ( t ′ ) s ( t ) x t ) . We start by showing:

Here, Eq. (20) says that p t ( x t | x ) ∝ p ′ t ′ ( s ′ ( t ′ ) s ( t ) x t | x ) (but not equal), which is a minor fix from the original proof [Pokle et al., 2023]. Then, we have

Since p ′ t ′ ( x | s ′ ( t ′ ) s ( t ) x t ) should be integrated to one, ( σ ( t ) σ ′ ( t ′ ) ) d 1 p ′ t ′ ( s ′ ( t ′ ) s ( t ) x t ) = ∫ p x ( x ) p t ( x t | x ) d x and thus two densities are equal. As the posterior densities are the same, their expectations are also the same.

In our case, s ( t ) = 1 -t , σ ( t ) = t , and p ′ t ( x t | x ) is the perturbation kernel of either the VP or VE diffusion model. Now we have to find t ′ such that 1 -t t = s ′ ( t ′ ) σ ′ ( t ′ ) .

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 et al. [2022] defines the perturbation kernel of the VE diffusion model as p ′ t ( x t | x ) = N ( x , t 2 I ) . Then, t ′ satisfies 1 -t t = s ′ ( t ′ ) σ ′ ( t ′ ) = 1 t ′ , so t ′ = t 1 -t , and s ′ ( t ′ ) s ( t ) = 1 1 -t , which correspond to t VE and s VE in Table 2.

VP diffusion model Song et al. [2020b] defines the perturbation kernel of the VP diffusion model as p ′ t ( x t | x ) = N ( α ( t ) x , (1 -α ( t ) 2 ) I ) , where α ( t ) := exp( -1 2 ∫ t 0 (19 . 9 s + 0 . 1) ds ) defined on t ∈ [0 , 1] . Then, t ′ satisfies 1 -t t = s ′ ( t ′ ) σ ′ ( t ′ ) = α ( t ′ ) √ 1 -α ( t ′ ) 2 . From here, we have

where we used the fact that α ( t ) > 0 . Since α ( t ) is a monotonically decreasing function for t ≥ 0 , we can use its inverse α -1 to find t ′ = α -1 ( √ (1 -t ) 2 t 2 +(1 -t ) 2 ) .

Applying the quadratic formula, we have

Since y = α ( t ) is monotonically decreasing, we can choose the positive root:

Now, we arrive at

which corresponds to t VP in Table 2. Also, we have s ′ ( t ′ ) s ( t ) = α ( t ′ ) 1 -t , which is s VP in Table 2.

D New Update Rule

In the standard Euler solver, the update rule is z t -∆ t := z t -v ( z t , t )∆ t . Alternatively, as x θ ( z t , t ) of our model generates pretty good samples, we can instead use the linear interpolation between x θ ( z t , t ) and z 1 to get the next step: z t -∆ t := (1 -( t -∆ t )) x θ ( z t , t ) + ( t -∆ t ) z 1 . Note that when NFE < 3 , the two update rules are equivalent and do not affect our results in Section 5.1. Fig. 4 shows that when applied to existing solvers, the new update rule improves the sampling efficiency up to 4 × , 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 t , our new update rule utilizes the previous state (i.e., z 1 ) to generate the next state z t -∆ t 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 to future work.

E Experimental Details

Before training 2-rectified flow, we generate data-noise pairs following the sampling regime of EDM [Karras et 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 second-order solver for all cases. In Table 3, we report the result of config G in Table 1. In ImageNet, we use the batch size of 2048 and train the models for 700,000 iterations using mixed-precision training [Micikevicius et al., 2017] with the dynamic loss scaling. We use config E in Table 1 for ImageNet.

We provide training configurations in Table 7. For all datasets, we use Adam optimizer. We use the exponential moving average (EMA) with 0.9999 decay rate for all datasets.

On 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, we use

Algorithm 2 Generate

the NVIDIA DGX-2 cluster. To prevent zero-division error with EDM initialization, we sample t from [0 . 00001 , 0 . 99999] in practice. For a two-step generation, we evaluate v θ at t = 0 . 99999 and t = 0 . 8 . For other NFEs, we uniformly divide the interval [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.

G Uncurated Synthetic Samples

We provide uncurated synthetic samples from our 2-rectified flow++ on CIFAR-10, AFHQ, and ImageNet in Figures 8, 9, 10, 11, 20, 21, 16, 17, 18, 19, 12, 13, 14, and 15. We use our new sampler (Sec. 5.3) to generate these images.

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 contributions and scope such as improved empirical performance, which is backed up by our experiments in Sec. 5.

2. Limitations

Question: Does the paper discuss the limitations of 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 a complete (and correct) proof?

Answer: [Yes]

Justification: Our only theoretical result is Proposition 1, for which we provide the proof in Appendix C.1. The argument in Sec. 3 is intuitive, and we do not prove it theoretically.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide our experimental details in Sec. 5 and App. E. For the new update algorithm, we provide full pseudocode in Sec. D, and we will release code publicly as soon as we obtain approval to do so.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide these details in Sec. E.

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 in Figure 3 indicates the standard deviation. We only compute FID once due to cost constraints.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide details in Sec. E.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 an impact statement in Sec. F. Effectively, as with other work on generative models, they can be used for both beneficial and harmful purposes. Our work, being focused on the mechanics of these models, does not introduce new risks, nor does it mitigate existing ones.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We believe that our 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.

12. 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. E.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide necessary information to reproduce our new models in Sec. E.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBy 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nHonestly 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.\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nThe authors have well addressed the limitations. I just asked that they talk about the memory cost of LPIPS too.", "TITLE": "Official Review", "RECOMMENDATION": 10, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nSome limitation are mentioned, which are reasonable.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-23"}

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

STRENGTHS:
1. This article provides a comprehensive analysis of a single return, and there is a clear motivation for improvement.

2. The authors mainly improved the training of the whole network from three aspects: time-step sampling, loss function, and initialization from the diffusion model.  

3. 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].

2. 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.       


[1] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models.
[2] Liu X, Zhang X, Ma J, et al. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.

QUESTIONS:
See Weaknesses.

LIMITATIONS:
The authors briefly discuss their limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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- 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.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nThis is included in conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}