2-Rectifications are Enough for Straight Flows: A Theoretical Insight into Wasserstein Convergence

Authors: Saptarshi Roy, Vansh Bansal, Purnamrita Sarkar, Alessandro Rinaldo
Primary area: generative models
Keywords: generative model, optimal transport, rectified flow, wasserstein distance

Abstract:
Diffusion models have emerged as a powerful tool for image generation and denoising. Typically, generative models learn a trajectory between the starting noise distribution and the target data distribution. Recently Liu et al. (2023b) designed a novel alternative generative model Rectified Flow(RF), which aims to learn straight flow trajectories from noise to data using a sequence of convex optimization problems with close ties to optimal transport. If the trajectory is curved, one must use many Euler discretization steps or novel strategies, such as exponential integrators, to achieve a satisfactory generation quality. In contrast, RF has been shown to theoretically straighten the trajectory through successive rectifications, reducing the number of function evaluations (NFEs) while sampling. It has also been shown empirically that RF may improve the straightness in two rectifications if one can solve the underlying optimization problem within a sufficiently small error. In this paper, we make two key theoretical contributions: 1) we provide the first theoretical analysis of the
Wasserstein distance between the sampling distribution of RF and the target distribution. Our error rate is characterized by the number of discretization steps and a new formulation of straightness stronger than that in the original work. 2) under a mild regularity assumption, we show that for a rectified flow from a Gaussian to any general target distribution with finite first moment (e.g. mixture of Gaussians), two rectifications are sufficient to achieve a straight flow, which is in line with the previous empirical findings. Additionally, we also present empirical results on both simulated and real datasets to validate our theoretical findings. The codes are available at https://github.com/bansal-vansh/rectified-flow.

Decision: accept (poster)
Reviewer scores: 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730566943964

thread_type:
official_review

<thread_object>
creator: reviewer_m3xn

<content>
summary:
The paper provides theoretical analysis on the distribution generated by rectified flow and elaborate on the case of Gaussian mixtures.

strengths:
1. The paper is clearly written and well-organized.
2. The theoretical framework is both novel and broadly applicable to the study of rectified flow, particularly in demonstrating the relationship between RF’s performance, the straightness of the flow, and the discretization accuracy.
3. The authors present experiments on both synthetic and real datasets. The figures, particularly those illustrating the sampling path and straightness, effectively convey the key findings.

weaknesses:
1. Assumption 3.1 regarding the estimation error and the Lipschitz property of the drift function, while common in previous diffusion model research, may not be essential for distribution estimation. These assumptions impose additional requirements on the RF sampling path. For the estimation error assumption, could the authors prove rather than assume a small estimation error when learning the drift function of a target distribution with some regularity assumptions using neural networks? Additionally, could the Lipschitz condition be derived using assumptions based solely on the target distribution rather than on the entire sample path?
2. The tightness and effectiveness of the theoretical results on the 2-Wasserstein bound (Theorems 3.2 and 3.5) are uncertain, as the error is exponentially dependent on the Lipschitz constant $\hat{L}$. Moreover, in theorem 3.5, the error bound even tends to infinity when $\hat{L} \rightarrow 0$, which seems counterintuitive. It would be helpful if the authors could conduct additional experiments to perform a quantitative sensitivity analysis on both $\hat{L}$ and $T$ to demonstrate the tightness of the results.
3. Minor typos: 
(Line 220) "on-side" Lipschitzness ->one-side Lipschitzness 
(Line 323) Under "he" condition -> Under the condition

questions:
1. Regarding the definition of straightness in Definition 2.1, why straigtness is equivalent to the latter expectation form?
2. Regarding the straightness of RF for Gaussian mixtures, is it possible to extend the theory to multivariate Gaussian mixtures with multiple modes? (more than 2 modes). What are the difficulty for proving the straightness?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730577133254

thread_type:
official_review

<thread_object>
creator: reviewer_ptxd

<content>
summary:
In the context of flow-based generative models, the authors provide a theoretical analysis of Rectified Flow, a generative model introduced by Liu et al (2023). Since the generative model has an associated ODE that has to be discretized to be solved in practice, how straight the solution paths of the ODE is a relevant question, where straighter paths can be solved with fewer discretization steps. Rectified Flow provides an iterative method that successively straigthens the paths. The authors provide a theoretical analysis showing that for generating Gaussian Mixtures, two iterations are enough to get straight paths. The authors also analyze the Wasserstein distance between the target distribution and the generated distribution after one step of Rectified Flow, showing that it can be upper bounded in terms of a measure of straightness.

strengths:
* The authors do a good job introducing the background and putting their results into context with what was previously proved, with a clear presentation for their results.
* The 2-Wasserstein bound provided in Theorem 3.5 and Corollary 3.6 are useful result since they gives an explicit dependence on the straightness parameters they defined and furthermore provides a guide on how to discretize to achieve a certain error if we have access to these straightness parameters. The intuition for these parameters and comparison with another definition of straightness parameter is also helpful.
* Theorem 4.4 is a clean and simple illustration of the effect rectifying flows and provides a stepping stone towards understanding rectifying flows in more general data distributions.

Overall, the authors provide an useful theoretical foundations both for understanding the exact effect of rectifying flows when the data is coming from the Gaussian Mixture, and also in bounding $W_2$ between target and data coming from the rectified flow. The new definitions of straightness introduced seem natural and could prove to be useful in future research.

weaknesses:
It is stated and proved that 1-Rectified flows are straight for Gaussian Mixtures. It seems to me that a very important question is not considered: whether there is any example of initial laws for $X_0$ and $X_1$ whose 1-Rectified flow is not straight. Equivalently, we can ask whether there is any example of laws for $X_0$ and $X_1$ whose 1-Rectified coupling is different from the 2-Rectified coupling. I may be missing something, but I tried to find for these counterexamples and did not succeed. I will raise my score if this question is addressed.

It is clear that 1-Rectified flows are always straight in 1 dimension, since any coupling coming from an ODE is monotonically increasing and deterministic so that by Lemma 4.2 is straight. However, say in Proposition 4.5, the intuitive argument given is that the flows along each coordinate decouples. Hence it seems to me that a much simpler proof may just be to show that the coordinates decouple and then just use Lemma 4.2. A similar reasoning may apply to the proof of Theorem 4.4.

The experiments do not seem to illustrate the upper bound on the 2-Wasserstein distance much. Indeed, the fact that the $W_2$ decays as $T$ grows is to be expected even if we did not know about the bound. It would be interesting if we saw in the experiments a clear evidence that $W_2$ decays with $\gamma_{2,T}.$ However, we see in Figure (3)(b) that $\gamma_{2,T}$ for the 2 component Gaussian Mixture is almost identical to the $\gamma_{2,T}$ for the 3 component, but the $W_2$ in Figure (3)(a) is decreasing much faster for the 2 component than for the 3 component. Of course, this does not contradict the bound since it is just an upper bound, but it is not illustrating it well either.

The plots in Figure (3)(a) and (c) are also not very informative since they just show zero for all $T\geq 5.$ It may be better to plot them in log scale. In particular, the bound suggests that there should be a $1/T^2$ decay for the $W_2,$ but since it is not in log scale it is quite hard to tell whether this is True in the empirical data or not.

Smaller comments: 
* In between line 216 and 217 it is said that the bound on the 2-Wasserstein is $O(\epsilon^2),$ but that is the bound on the 2-Wasserstein squared.
* The plots on Figure 1 are not displayed correctly, the x label is not showing.
* In line 206 it says $\hat \rho _\text{data}$ it should say $\rho$
* In line 837, there is an incomplete line that is likely missing an equation.
* The proof sketch for Theorem 4.4 is a bit confusing. I.e., the presentation could be improved e.g. it says '(WLOG call this coordiante 1). 
 Denote this coordinate by $c_1,$' and 'isoperimetric mixture' without defining it.

questions:
* In between liens 217 and 218, it is said that the requirement on $b(t)$ of $\sup_{t\in[0,1]}b(t)\leq \epsilon^2$ is much more stringent than that of Assumption 3.1(a). Is there actually an application/relevant setting where Assumption 3.1(a) is fulfilled but $\sup_{t\in[0,1]}b(t)\leq \epsilon^2$ is not?
* The points chosen for the experiment on the Gaussian Mixture mentioned in Example 1 are a bit unmotivated to me. Why did you chose them this way? 
* See the first paragraph of weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730687353143

thread_type:
official_review

<thread_object>
creator: reviewer_apm1

<content>
summary:
This submission presents an error analysis of rectified flow (RF) in terms of the Wasserstein distance. (1) It establishes an error bound for the Wasserstein distance between the target distribution and the sampling distribution generated by RF, based on the Lipschitz continuity of the vector fields and the average curvature of the flow curve, which measures its straightness. (2) For Gaussian mixtures, the authors demonstrate that 2-rectifications are sufficient to achieve a straightened flow.

strengths:
The theoretical analysis is clear and intuitive. The choice of Wasserstein distance naturally aligns with the formulation of least squares (least action principle). The authors use that to control the Wasserstein distance between the target and sampling distributions based on the squared distance between the velocity fields. Additionally, this squared distance is further bounded by the average curvature through a Poincaré-type inequality. The authors also provide a rigorous justification for achieving strict straightness in 2-rectifications for Gaussian mixtures, a phenomenon that has also been observed numerically in more complex distributions.

weaknesses:
The benefit of straightness in the rectified flow is not demonstrated. In the conclusion, the authors claim that "straight flows are desirable because they require fewer Euler discretization steps." However, there is no quantification of this gain or a clear link to straightness provided in the analysis.

Although the straightness of two-step rectified flow is proved for Gaussian mixtures, it remains unclear how this result can be extended to more general distributions.

questions:
1. Could the authors provide a more detailed explanation or quantification of the computational advantages of achieving straightness in rectified flows? For instance, how does straightness impact the number of discretization steps required in practice?

2. How robust is the theoretical result for two-step rectified flow straightness when moving beyond Gaussian mixtures? Are there any indicators or preliminary results that suggest potential extensions to more general distributions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731896662819

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised version uploaded and A new general straightness result is added

comment:
We thank all the reviewers for their suggestions and constructive comments. We have revised the paper and uploaded the new version. More importantly, In Section 4, we have added two new general results (Theorem 4.3 and Theorem 4.5) that shows that **1-RF yields a straight coupling for a large class of distributions (which extends beyond Gaussian mixtures).** We would like to emphasize the generality of this result, which, to the best of our knowledge, has not been previously reported in the literature. An informal version of Theorem 4.5 is stated below.

*_Theorem 4.5 (informal version)_: Assume $X_0 \sim N(0, I_d)$ and the target data $X_1 \sim \rho_1$ (independent of $X_0$) with finite 1st moment, i.e., $\mathbb{E} \Vert X_1\Vert_2< \infty$. Also, assume that the ODE (3) is non-explosive (essentially solutions $Z_t$ do not diverge to infinity). Then, the resulting rectified coupling $(Z_0,Z_1):= \texttt{Rect}(X_0, X_1)$ under the aforementioned conditions is  a straight coupling*

Note that, in the theorem, target distribution does not need to have a pdf, it only needs to have a finite 1st moment. Also, the non-explosive condition is fairly mild and is enjoyed by a large class of distributions including a general mixture of Gaussians. Therefore, the current results provide a general straightness result for a large class of rectified flow. We have also added recommended simulation experiments that corroborate our theoretical claims.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732570087007

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised(minor) version uploaded: Slight change in paper title and minor typo corrections

comment:
We want to sincerely thank all the reviewers again for handling our submission. We wanted to point out that we have made a slight modification of the paper title to underscore the new straightness results (Theorem 4.3 and Theorem 4.5) mentioned in the rebuttal. The current title reads "2-RECTIFICATIONS ARE ENOUGH FOR STRAIGHT FLOWS: A THEORETICAL INSIGHT INTO WASSERSTEIN CONVERGENCE." We could not change the title in the Openreview directly as we could not find the option. However, we will change it during the camera-ready submission period if accepted. In addition, we have also addressed a few typographical errors and minor issues that were pointed out by the reviewers.

Regards,
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper is on the topic of rectified flow. It makes two theoretical contributions to the study of rectified flow. First, it provides a theoretical analysis to quantify the Wasserstein error of a trained rectified flow model. Second, it shows, under proper assumptions, 2 rectifications are sufficient to achieve a straight flow. All reviewers agree the paper is well written and present the theoretical contributions clearly. Two reviewers point out that the experiments are insufficient to illustrate the theoretical results. One reviewer questions the tightness of the theoretical bound. The authors provide satisfactory answers to these questions. Overall, this is a nice contribution to the rectified flow framework.