On the Relation between Rectified Flows and Optimal Transport

Many thanks for your reading of our paper, and your review, comments and suggestions! We reply to all of your critiques separately.

While the proposed "smoothed rectification" method is a valuable idea, its analysis feels preliminary. The theoretical guarantees are shown for the specific case where the source distribution is Gaussian, which may not cover all use cases. More importantly, while the authors prove that their method maintains a good convergence rate for the loss, they do not investigate if it actually helps in converging to the true optimal transport plan, especially in the challenging cases (like disconnected supports) that they identify.

Please find our answer to this comment together with the answer to your last comment at the end of our reply.

While the counterexample in Section 4.1 based on disconnected supports is technically sound and arguably the paper's most critical contribution, its practical significance is not immediately obvious from the text. The authors should explicitly state why this scenario is not just a theoretical curiosity but a common and fundamental challenge in practice.

One key motivation and application of rectified flows is distribution translation, i.e., constructing a mapping that transports one probability distribution to another. In this setting, connecting rectified flows to optimal transport is particularly appealing, as it yields a mapping that is unique (under absolute continuity), stable, and interpretable (minimizing the expected squared distance between inputs and outputs).

However, our counterexample demonstrates that this connection breaks down when the distributions are disconnected. Since disconnected datasets are common in practice, this indicates that rectified flows cannot reliably compute optimal transport mappings in general.

Another major motivation for rectified flows is the fact that the flow paths tend to become "straighter", improving computational efficiency during sampling by simplifying the ODE trajectories. This aspect is largely independent of optimal transport, and our counterexample does not affect this advantage.

In the next revision, we will extend the comments in introduction, Section 4.1 and Conclusions based on this discussion.

No experiments demonstrate the proposed smoothed rectification procedure.

In the next revision, we will extend Appendix B by some numerical experiments in the setting of Remark 9 (see our reply to the next comment, a brief overview over these experiments is given at the end of the rebuttal).

Does noise injection also help the iterative process escape these non-optimal fixed points? For instance, in your disconnected support counterexample (Fig. 1), if one were to run the smoothed rectification procedure starting from the independent coupling, would it be more likely to converge to the globally optimal transport plan (Fig. 1a) compared to the standard procedure? My evaluation of the paper's practical contribution would be further strengthened if the authors could provide either (a) an empirical study, perhaps by extending the experiments in Appendix B, to show that smoothing helps find the correct OT plan, or (b) a theoretical argument or conjecture on how noise might provide a regularizing effect that favors the true optimum.

From a theoretical side, the definition of the smoothed rectification ($X_0^{(i)}=\sqrt{1-c_i}Z_0^{(i)}+\sqrt{c_i}W^{(i)}$ for standard Gaussian noise $W^{(i)}$) ensures not only rectifiability, but also that $\mathrm{supp}(X_t^{(i)})=\mathbb R^{d}$ (for $t<1$). Consequently, the smoothed rectification indeed overcomes the problem with disconnected supports which is the main idea for the counterexamples from Section 4.1. However, while the smoothed rectification helps for the counterexamples from Section 4.1, we are currently unable to provide a proof that it generally converges to the optimal coupling (for appropriate $c_i$).

We have now also tested the smoothed rectification for the counterexample from Remark 9 (we do not use the other example from Section 4.1, because the smoothed rectification assumes that $\mu_0$ is Gaussian). As initialization, we use the non-optimal fixed point from eqt (12). The results (see below) suggest that smoothed rectification does converge to the optimal coupling. Furthermore, the convergence speed significantly decreases as the noise level approaches zero. This is expected since a small change in the coupling results only in the small change in the rectification.

We will present these findings and the corresponding discussion in Appendix B of the next revision. A brief overview of the experiments and results is given below (for the paper we will also generate figures in the style of Figure 6/7/8 visualizing the resulting couplings).

---

Additional Experiment with Smoothed Rectification

We consider the example of Remark 9 and apply the smoothed rectification with parameter $c_k=\frac{c_0}{k}$ for $c_0\in\{0,0.01,0.05,0.1,0.2\}$, where $c_0=0$ resembles the case of the standard rectification $\mathcal R_p$. The initial coupling $(X_0^{(0)},X_1^{(0)})$ is set to the non-optimal fixed point of $\mathcal R_p$ from Remark 9 (defined in eqt (12), visualization in Figure 1c). We report the transport distance $(\mathbb E[\\|X_0-X_1\\|^2])^{1/2}$ versus the number of steps of the smoothed rectification procedure in the table below. A lower transport distance indicates that the coupling is closer to otpimal transport. The optimal coupling admits a transport distance of $2.66$ (evaluated based on $20000$ samples using the POT package). Since all experiments are initialized with the same coupling, the first step coincides over all runs, independent of the noise level. Due to numerical errors in the marginals, the transport distance is sometimes slightly smaller than the anlytical optimum. Overall, we observe that the smoothed rectification indeed escapes the non-optimal fixed points from Section 4.1. However, the convergence becomes slower if the noise level approaches zero.

In the paper, we will also include figures (similar to Fig 6/7/8) for this experiments.

We would like to thank the reviewer for having read carefuly our paper and their efforts in evaluating it. We provide answers to the specific critiques separately.

To improve clarity, Definitions should be included. For example, the rectified coupling is defined on line 110 with a reference to equation (1). I think the paper would have been easier to follow if such important definitions were in Definition sections such as theorems and propositions.

Thank you for the suggestion! We will revise the paper and put important definitions (like the rectification) into a definition environment.

Isn't there a mistake in the definition (12) of $X_1$ in Remark 9? Shouldn't the second case also contain instead $X_0$ of $X_1$?

You are absolutely right! Thank you for pointing out this typo.

To make sure that I understand the smoothed rectification in Theorem 15, you assume $\mu_0$ is gaussian, and since rectification preserves the marginals, $X_0$ will remain gaussian, but it is the conditional distribution of $X_0$ (given $X_1$) that must be absolutely continuous with a smooth and positive density, which is why you add an independent Gaussian noise?

Correct! Adding independent Gaussian noise (and rescaling $Z_0^{(i)}$) in the definition to $X_0^{(i)}$ ensures that the conditional distribution $P_{X_0|X_1=x_1}$ has a smooth positive density such that we can apply Theorem 11. At the same time $X_0$ remains a Gaussian distribution since the rectification preserves the marginals and a sum of two Gaussians remains Gaussian.

Counterexamples are provided which mean that stronger assumptions are needed for rectification to converge to an optimal transport, but these stronger assumptions are not investigated. This is of course minor as I think the paper is rich enough in results.

Well, we can at least state sufficient assumptions such that Theorem 5.6 from [1] is correct (where the additional assumption is marked italic):

Assume that $(X_0,X_1)$ is rectifiable and let $v_t=\nabla \varphi_t\in\mathrm{argmin}_{w_t=\nabla \psi_t}\mathcal L(w_t|X_0,X_1)$ fulfill that $\varphi_t\in C^{2,1}(\mathbb R^d\times[0,1])$. Moreover, suppose that $\mathrm{supp}(X_t)=\mathbb R^d$ for $X_t=(1-t)X_0+tX_1$. Then the following are equivalent:

• $\mathcal R_p(X_0,X_1)=(X_0,X_1)$,
• $\mathcal L(v_t|X_0,X_1)=0$,
• $(X_0,X_1)$ is an optimal coupling.

However, we stress that within the iterative rectification $(X_0^{(i+1)},X_1^{(i+1)})=\mathcal R(X_0^{(i)},X_1^{(i)})$ it is usually intractable to check, whether these assumptions remain fulfilled for all $i$ based on the initial coupling $(X_0^{(0)},X_1^{(0)})$. Moreover, we showed in Corollary 14 that the fact that the loss decays does not necessary imply convergence of the (couplings) to a fixed point of $\mathcal R$. How (or if at all) the decay of the loss relates to convergence of the couplings to a fixed point remains an open question. In the next revision, we will include the statement of Theorem 5.6 from [1] with the corrected assumptions (also considering the comments of reviewer 3EWA) and extend the explanations on this point in the conclusions.

Thank you for having read accurately our paper and for your comments! We reply to all of these separately.

While the main contribution presented in the paper is probably already known to theorists in this area [...]

We are not aware of any comparable results from the literature. In addition, the slightly incorrect results which we discuss are cited in plenty of papers, without any discussion of their validity. We really feel that pointing this out is very useful to the community.

The structure of the paper is a bit confusing. For instance, what is the purpose of Section 3.1? This section appears abruptly from the previous discussion of rectified flows with the gradient potential conditions, and appears to have little connection with the rest of the presentation. From inspection, it appears that Section 3.1 is only used in the proofs of later results. The layout disrupts from the presentation of the main argument of the paper and makes the paper harder to follow.

Some relations between rectified flows and optimal transport are already true for the case without the gradient potential condition. Section 3.1 gathers these foundational results, including

• Rectified flows and optimal transport share some invariance properties (Theorem 1 (ii) and (iii))
• The observation that, in both the Gaussian setting (Theorem 2 (ii)) and the one-dimensional case (Proposition 3), the optimal coupling is achieved after a single rectification step.

As a side effect, these results are quite useful to understand the basic behavior of rectified flows and are consequently used in the proofs of subsequent results.

The current structure aims to provide a general overview of the literature in Section 2, followed by our contributions in the subsequent sections. However, we recognize that a revised layout could enhance accessibility from a pedagogical perspective. In the next revision, we consider to restructure Sections 2 and 3 as follows:

• Section 2 will present the general background, excluding the discussion on the gradient potential condition (currently lines 88–119), and incorporate the results presently found in Section 3.1.
• Section 3 will begin with the background on the gradient potential condition (lines 120–137), followed by the results currently contained in Section 3.2.

Line 193: "$v^{A,b}_t$ has a potential if and only if the Jacobian is symmetric." - what is the justification of this? I'm not sure how this immediately follows?

Here we refer to the fact that a vector field admits a potential (locally or in a "simple" topology) if and only if the Jacobian is symmetric (such vector fields are often referred as "conservative"). This is a non-trivial, yet well-known result, contained in many text books: One direction immediately follows from the fact that the Jacobian of the vector field is the Hessian of the potential and consequently symmetric (known as Schwarz's theorem). The other direction is known as Poincaré's lemma (and equivalent to the property that the circulation of such field along a closed loop vanishes). In the next revision, we will add an additional sentence that we use the above equivalence and add a reference to some text book.

Line 244: "The statement of...". The full statement of the corrected proposition and proof would be valuable to include in this paper to resolve any ambiguity, instead of leaving this offhand remark halfway in the paper

We agree that presenting the corrected statement could be beneficial. In the next revision, we will include the following version of Thm 5.6 from [1] for $c(x-y)=\\|x-y\\|^2$ (where the additonal assumption is marked italic):

Assume that $(X_0,X_1)$ is rectifiable and let $v_t=\nabla \varphi_t\in\mathrm{argmin}_{w_t=\nabla \psi_t}\mathcal L(w_t|X_0,X_1)$ fulfill that $\varphi_t\in C^{2,1}(\mathbb R^d\times[0,1])$. Moreover, suppose that $\mathrm{supp}(X_t)=\mathbb R^d$ for $X_t=(1-t)X_0+tX_1$. Then the following are equivalent:

• $\mathcal R_p(X_0,X_1)=(X_0,X_1)$,
• $\mathcal L(v_t|X_0,X_1)=0$,
• $(X_0,X_1)$ is an optimal coupling.

The proof is mainly the same as in [1]. But we will add an appendix, which provides a more detailed explanation on this point.

For Theorem 2 and Theorem 4, covariance matrices are referred to as both $X_0$ and $X_{00}$. Why is this inconsistent?

This is a typo, thank you for pointing this out. We will correct it within the next revision.

In line 184, I do not understand how the independent coupling of $(X_0,X_1)$ follows a GMM distribution as $X_1$ follows a discrete measure. Hence, I do not see how Theorem 4 can be applied to this case.

Here we use the definition of the multivariate normal distribution which includes singular covariance matrices. More precisely, we say that a $d$-dimensional random variable $X$ follows normal distribution if and only if all one-dimensional projections $a^T X$ follow a normal distribution with the notation that $\delta_x=\mathcal N(x,0)$ for $x\in\mathbb{R}$. It can be shown that using this definition, Gaussian distributions can still be identified by their mean and covariance matrix (which is still positivie semi-definite but not necessarily invertible) and that it preserves all computation rules of Gaussian distributions which are used within the statement and proofs of the Theorem 2 and 4. Then, the joint distribution of independent $X_0\sim\mathcal N(0,\mathrm{Id})$ and $X_1\sim\sum_{k=1}^K\frac1K\delta_{m^k}$ follows the GMM $\sum_{k=1}^K\frac1K\mathcal N\left(\left(\begin{array}{c}0\\\\m^k\end{array}\right),\left(\begin{array}{cc}\mathrm{Id}&0\\\\0&0\end{array}\right)\right)$ In the next revision, we will add a remark about this and add a reference to some text book introducing Gaussian distributions and their computation rule for singular covariance matrices.

The notation in Remark 5 is confusing. I assume that here $\varphi^a_t$ is referring to the transformation by adding the constant $a$. However, in Theorem 1, the notation for this transformation uses the letter $b$ in $v_t^b$. I don't see any reason for changing up the notation here, apart from to reduce clarity.

This is again a typo, which we will correct in the next revision. We are sorry for the confusion!

To what extend can the counter-examples/proofs/intuition from the paper be generalised to rectified flows under the setting of a generic cost function, as presented in [1]?

The counterexample regarding the non-rectifiable transport plan from Section 4.2 works for any cost function fulfilling the assumptions from [1]. The main argument here is that the distribution $\mu_t$ of $X_t=(1-t)X_0+tX_1$ becomes singular for $t=\frac12$. Therefore, for $v_t$ fulfilling the continuity equation $\partial_t \mu_t+\mathrm{div}(v_t\mu_t)=0$ (which is true for the minimizer of the loss function) and for any solution of $\dot Z_t=v_t(Z_t)$ with $Z_0=X_0$, it holds that $Z_{1/2}=0$ almost surely. In other words, almost-every ODE path collapses to a single point at time $\frac12$ which contradicts the uniqueness of solutions.

The counterexample from Section 4.1 might not work "as-is" for other cost funcitions (already because it might be possible to construct cost functions such that both couplings from eqt (11) admit the same cost). However, it should be straight-forward to construct analogous examples for other cost functions (which are convex and have a $C^1$ convex conjugate) by the same idea: Construct two locally optimal transport plans which are sepatated by some area where the coupling has zero-mass. Then, glue them together such that the global plan is not optimal.

Within the next revision, we will add a remark about other cost functions in the Conclusions section.

Reference:
[1] Qiang Liu (2022) Rectified Flow: A Marginal Preserving Approach to Optimal Transport