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

(P4) 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?-- We now have a fairly general result for multivariate distribution, and we have added it as Theorem 4.5 (or Theorem 4.3). Note that, in the new result, 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 with multiple modes. See discussion in Lines 381-384 and Lines 388-393. The main difficulty in the proof arose in showing the invertibility of a certain multivariate map, for which we used the Inverse function theorem. Also, we had to take care of the existence and uniqueness of the ODE under a more general setting.

Thank you for your questions and constructive comments. We address your concerns point by point below.

(P1) 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? -- Thank you for your constructive comments and the interesting questions. We think , one can prove results similar to estimation error bound by imposing assumptions on the neural network. In particular, one can also obtain sample complexity results (under finite VC dimension or similar assumptions) that will guarantee good estimation error. However, in this paper we solely wanted to investigate the efficacy of Rectified flow. However, we do agree that a theory for estimation error is pertinent in such studies and defer it to future research.

We think that it is possible to impose a condition on the target distribution so that the velocity function is Lipschitz. For example, if the pdf $\xi_1$ of the target distribution is fairly smooth with bounded derivatives, then it can be shown that the resulting velocity is Lipschitz. However, one can ensure the lipschitzness of the estimated velocity $\hat v$ by restricting the class of neural networks with bounded weights.

(P2) 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.-- We do agree that the upper bound result has room for improvement and it may not be tight. However, we want to point out that this the very first result that connects $W_2$ distance with the straightness of rectified flow. We mainly wanted to highlight the effect of straightness on the generation quality of rectified flow. We will also add some numerical plots to highlight the effect of straightness in the revised version.

Also, thank you for pointing out the issue regarding $\hat L \to 0$. A refined analysis shows that the bound can be improved to

The refined bound now circumvents the aforementioned issue. Please refer to the updated Theorem 3.5.

We also provide the plot of $\log(W_2)^2$ vs $L$ in this link. It shows that $W_2$ is increasing with $L$. However, the plot also suggests that there is room for improvement in the upper bound in Theorem 3.5. We defer this to future work as mentioned in Section 6.

(P3) Regarding the definition of straightness in Definition 2.1, why straightness is equivalent to the latter expectation form?-- Straight flow implies that the velocity $\mathbb{E}(X_1 - X_0 \mid X_t)$ along the path is not changing with time $t$. Let's say it is a constant $c$ independent of $t$. $X_1$ and $X_0$ are the two end-points then we need $X1 = X_0 + c$. In this case velocity is $\mathbb{E}(X1- X_0 \mid X_t)$, and that is why it follows that $\mathbb{E}(X_1 - X_0 \mid X_t) = c = X_1 - X_0$ almost surely. Moreover, if $\mathbb{E}(X_1 - X_0 \mid X_t) = c = X_1 - X_0$, then velocity is $t$ independent, and hence the trajectory is straight. Also, we rectified a minor error regarding the non-intersecting trajectories in Definition 2.1. Please refer to the updated version.

Thank you for your questions and constructive comments. We address your concerns point by point below.

(P1) 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 $X_0$ for $X_1$ and 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... -- Thank you for this thought-provoking question. Just to be, clear, we show that for the mixture of Gaussian, 1-rectification generates a straight coupling, or equivalently 2-rectification is a straight flow. This is also equivalent to saying that 1-rectified coupling is the same as the 2-rectified coupling as you have already mentioned. We do agree that it is important to understand when straightness is violated. First, we would like to draw your attention to a fairly general result (Theorem 4.5 and Theorem 4.3) that we recently proved. An informal version is displayed 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):= `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, to come up with a counter-example, we have to look outside the aforementioned examples.

We first want to draw your attention to the following example: $X_0\sim \delta_0$ and $X_1 \sim 0.5 \delta_1 + 0.5 \delta_2$, where $\delta_n$ denotes Dirac measure on $n$. In this case, there is no straight coupling. To see this, assume there is a straight coupling. Then it can be represented by a deterministic map $\psi$ such that $X_1 \overset{d}{=} \psi(X_0)$. This is contradiction as $\psi(X_0) \sim \delta_{\psi(0)}$. In fact, $(X_0, X_1)$ can not be rectified, as rectification always leads to a deterministic coupling which again raises contradiction by the previous logic. However, under the knowledge of the current general theorem, we suspect that if $(X_0, X_1)$ is rectifiable for some initial choices of distributions, then the 1-rectified flow will result in a straight coupling under very mild regularity conditions. In fact, we conjecture that 1-rectified flow (under certain regularity conditions) will result in the optimal coupling induced by the Monge map (if it exists), as it is known that RF iteratively solves the optimal transport problem [1]. Although, this is an interesting research direction we defer it to future research.

(P2) 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.-- We would like to point out that all 1-rectified flows in 1-dimension are not straight as shown in the previous example. If the velocity $v(x,t)$ is Lipschitz, then one can argue that the resulting coupling is monotonic as the solutions to the ODE is unique due to Picard-Lideloff theorem.

Next, we would like to point out that the decoupling of the coordinates phenomenon is not a general phenomenon. For example, even in Theorem 4.4, the coordinates do not decouple if the target distribution is $\pi_1 N(\mu_1, \Sigma_1) + \pi_2 N(\mu_2, \Sigma_2)$ for some general choices of $\{\pi_i, \mu_i, \Sigma_i\}_{i \in [2]}$. In the current paper, $\Sigma_i = I_d$ helps us to decouple the co-ordinates. Also, for Proposition 4.4, symmetric choices of the means are crucial for the decoupling argument which may not hold for general choices of the mean vectors.

To this end, we want to Theorem 4.5 which we mentioned in the previous point. Note that, 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 aforementioned theorem covers a large class of rectified flows.

[1] Liu, Q. (2022). Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577.

Dear Reviewer ptxD,

Thank you so much for your thoughtful remarks and thorough review. We have added a new general straightness result (Theorem 4.3 and 4.5) which shows that 2-RF yields straight flow from Gaussian to a fairly large class of distributions including a general mixture of Gaussians. In the first rebuttal, we also provided a simple example where 2-RF is not straight, but as it turns out, in that case, the Monge map does not exist. Under the light of Theorem 4.5 (and Theorem 4.3), we believe it is hard to develop an example where a Monge map exists but the flow is not straight. In the discussion of the updated manuscript, we pose an open question "Does the existence of Monge map imply straightness of 2-RF?". We discussed this in our previous comments and Section 6 of the paper. Hopefully, the discussion will help clarify some of the concepts. Given the discussion period is coming to an end, we were wondering if you have any further questions that we could address so that we can improve the paper.

P1. The new general result addresses my concern.

P2. I do not mean that for any target distribution we have decoupling. I am just saying that a way to argue in Theorem 4.4, given that the problem does decouple for the given GM target distribution, would be to show that in each dimension we are straight. Can this idea simplify the proof of Theorem 4.4? (When I say that "1-Rectified flows are always straight in 1 dimension" I am assuming that the 1-Rectified flow exists. The counterexample you provide can not be rectified so it does not apply.)

P3-P6: My concerns are addressed.

I raise my score to 8 since this new general result appears to me as very important for the field of rectified flows.

Thank you very much for raising the score and your valuable comments. We are glad that you found our results important. Regarding your concern related to (P2), you are right that coordinate wise straightness would imply overall straightness. Therefore, showing coordinate wise straightness is enough to establish the straightness under the decoupling phenomenon. However, for general $\mu_1, \mu_2 \in \mathbb{R}^d$, the problem do not necessarily decouple. We have to change the coordinate system via a judicial rotation to allow us to decouple the coordinates. After that step, we essentially show that the problem is coordinatewise straight. The application of rotation is crucial in this proof.

Also, thank you for pointing out the typo in lines 364-365. We will fix this in the updated version.

Thank you for your questions and constructive comments. We address your concerns point by point below.

(P1) 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. -- We would like to point out that Theorem 3.5 (and Corollary 3.6) shows the dependency of $W_2$ distance on the straightness parameter $\gamma_{2,T}$. The discussion after Theorem 3.5 (line 304 - 315) clearly points out that $T = \Omega(\sqrt{\gamma_{2,T}/\epsilon})$ is enough to achieve desired error bound on $W_2$ distance. This shows that more straight flow requires less number of discretizations (i.e., small $T$) steps resulting in faster generation time.

(P2) 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.-- The benefit of rectified flow in terms of generation quality has been numerically demonstrated in past works like [1, 2]. In this, paper we provide a theoretical underpinning to this fact by connecting $W_2$ distance with the straightness parameter $\gamma_{2,T}$. [1, 2] also theoretically showed that the straightness parameter $S(\mathcal{Z})$ decreases (in some sense) with successive rectification. Therefore, we expect that $\gamma_{2,T}$ decreases for higher rectification. In Section 4, we actually theoretically show that $\gamma_{2,T} = 0$ (i.e., straight flow) after 2-rectification for some curated examples. In fact, while revising, we proved a more general result, i.e., under very mild assumptions on the target distribution, the parameter $\gamma_{2,T} = 0$ after 2-rectification (Theorem 4.5 and Theorem 4.3). An informal version of the Theorem 4.5 is shown 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):= `Rect`(X_0, X_1)$ under the aforementioned conditions is a straight coupling

This shows that we do not need any discretization after 2-rectification.

(P3) 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? -- Numerical experiments in [1,2] have already demonstrated the benefit of rectified flow and we also discuss those in the 2nd paragraph of Page 1 (line 44-51). Essentially, rectification aims to straighten a potentially curved flow between two distributions, thus reducing the transport cost. This allows the rectified flow to generate higher-quality data within fewer steps. We also discuss about the dependency of discretization steps $T$ with the straightness parameter $\gamma_{2,T}$ in the discussion after Theorem 3.5 (Line 304-316).

(P4) 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?-- Thank you for your question. Yes, the result actually holds in more generality and we proved this while revising the paper. Please refer to Theorem 4.5 (and Theorem 4.3) of the updated version.

Note that, 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.

[1] X. Liu, X. Zhang, J. Ma, J. Peng, et al., “Instaflow: One step is enough for high-quality diffusion-based text-to-image generation,” in The Twelfth International Conference on Learning Representations, 2024.

[2] X. Liu, C. Gong, and Q. Liu, “Flow straight and fast: Learning to generate and transfer data with rectified flow,” in The Eleventh International Conference on Learning Representations, 2023.