Rectified Flow: A Marginal Preserving Approach to Optimal Transport

Introduction

The Monge–Kantorovich (MK) optimal transport (OT) problem concerns finding an optimal coupling between two distributions ${\pi}_{0},{\pi}_{1}$:

We present a different approach tocontinuous OTthat re-frames (1) intoa sequence of simple unconstrained nonlinear least squares optimization problems,whichmonotonically reduce the transport cost of a couplingwhile automatically preserving the marginal constraints.Different from theminimax and regularization approaches that enforce the constraints from outside,our method is an interior approachwhich starts from a valid coupling (typically the naive independent coupling), and traverses inside the constraint set to decrease the transport cost.Such an interior approach is non-trivial and has not been realized before, because there exists no obvious unconstrained parameterization of the set of couplings of ${\pi}_{0}$ and ${\pi}_{1}$. Our method is made possible by leveragingrectified flow ,a recent approach to constructing (non-optimal) transport maps for generative modeling and domain transfer.What makes rectified flow specialis that it provides a simple procedurethat turns a given coupling into a new one thatobeys the same marginal laws, while yielding no worse transport cost w.r.t. all convex functions $c$ simultaneously.Despite this attractive property,as pointed out in ,rectified flow can not be used to optimize any fixed cost $c$,as it is essentially a special multi-objective optimization procedure that targets no specific cost.Our method is a variant of rectified flow that targets a user-specified cost function $c$ and hence yields a new approach to the OT problem (1).

Rectified flow

We provide a high-level overview ofthe rectified flow of and the main results of this work.For a given coupling $(X_{0},X_{1})$ of ${\pi}_{0}$ and ${\pi}_{1}$,the rectified flowinduced by $(X_{0},X_{1})$is the time-differentiable process ${\boldsymbol{Z}}=\{Z_{t}\colon t\in\}$ over an artificial notion of time $t\in$,that solves the following ordinary differential equation (ODE):

and $X_{t}$ is the linear interpolation between $X_{0}$ and $X_{1}$.Eq (3) isa least squares regression problem of predicting the line direction of $(X_{1}-X_{0})$ from every space-time point $(X_{t},t)$ on the linear interpolation path, yielding a solution of

c𝑐c-Rectified flow

In this work,we modify the $\mathtt{Rectify}$ procedure so that it can be used tosolve (1) given a user-specified cost function $c$.We show that this can be done easily by properly restricting the optimization domain of $v$ and modifying the loss functionin (3).The case of quadratic loss $c(x)=\frac{1}{2}\left\lVert x\right\rVert^{2}$is particularly simple,for which we simplyneed to restrict the $v$to be a gradient field $v_{t}=\nabla f_{t}$ in the optimization of (3).For more general convex $c$,we need to restrict $v$ to have a form of $v_{t}(x)=\nabla c^{*}(\nabla f_{t}(x))$, with $f$ minimizing the following loss function:

Notation

Outline

The rest of the work is organized as follows. Section 2introduces the background of optimal transport. Section 3reviews rectified flow of from an optimization-based view.Section 4 characterizesthe if and only if conditionfor two differentiable stochastic processes to have equal marginal laws.Section 5 introduces the main $c$-rectified flow method and establishes its theoretical properties.

Background of Optimal Transport

This section introduces the background of optimal transport (OT),including both the static and dynamic formulations.Of special importance is the dynamic formulation,which is closely related to the rectified flow approach.The readers can find systematic introductions to OT ina collection of excellent textbooks.

The optimal transport problem wasfirst formulated byGaspard Monge in 1781 when he studiedthe problem of how toredistribute mass, e.g., a pile of soil, with minimal effort.Monge’s problem can be formulated as

As the left side of (7)only depends on $(X_{0},X_{1})$ and the right side only on $(\mu,\nu)$, one can show that $(X_{0},X_{1})$ is $c$-optimal and $(\mu,\nu)$ solves (6) iff$\mu(X_{0})+\nu(X_{1})=c(X_{1}-X_{0})$holds with probability one, which provides a basic optimality criterion. Many existingOT algorithms are developed by exploiting the primal dual relation of (1) and(6)(see e.g., ), but have the drawback of yielding minimax problems that are challenging to solve in practice. If $c$ is strictly convex,the optimal transport map of (5)is unique (almost surely) and yields a form of

where $c^{*}$ is the convex conjugate function of $c$, and $\nu$ is an optimal solution of (6), which is $c$-convex in that $\nu(x)=\sup_{y}\left\{-c(y-x)+\mu(y)\right\}$ with $\mu$ the associated solution.In the canonical case of quadratic cost $c(x)=\frac{1}{2}\left\lVert x\right\rVert^{2}$, we can write $T(x)=\nabla\phi(x)$, where $\phi(x)\coloneqq\frac{1}{2}\left\lVert x\right\rVert^{2}+\nu(x)$ is a convex function.

Dynamic formulations

Both the MK and Monge problems can be equivalently framed in dynamic waysas finding continuous-time processes that transfer ${\pi}_{0}$ to ${\pi}_{1}$.Let $\{x_{t}\colon t\in\}$ be a smooth path connecting $x_{0}$ and $x_{1}$, whose time derivative is denoted as $\dot{x}_{t}$.For convex $c$, by Jensen’s inequality, we can represent the cost $c(x_{1}-x_{0})$ in an integral form:

where the infimum is attained when $x_{t}$ is the linear interpolation (geodesic) path: $x_{t}=tx_{1}+(1-t)x_{0}$.Hence, the MK optimal transport problem (1) is equivalent to

Hence, we can rewrite (9) into an optimization problem on $(v,\varrho)$, yielding the celebrated Benamou-Brenier formula :

Rectified Flow: An Optimization-Based View

We introduce rectified flow of from an optimization-based perspective:we show that rectified flow can be viewed as the solutionof a special constrained dynamic optimization problem,which allows us to gain more understanding of rectified flow and motivates the development of $c$-rectified flow.Following ,for a time-differentiable stochastic process ${\boldsymbol{X}}=\{X_{t}\colon t\in\}$, its expected velocity field $v^{{\boldsymbol{X}}}$ is defined as

where $\dot{X}_{t}$ denotes the time derivative of $X_{t}$.Obviously, $v^{{\boldsymbol{X}}}$ is the solution of

Precisely, Equation (13) should be interpreted by its weak and integral form:

Moreover, the rectified flow of a coupling $(X_{0},X_{1})$is defined as the rectified flow of ${\boldsymbol{X}}$ when ${\boldsymbol{X}}$ is the linear interpolation of $(X_{0},X_{1})$.

A stochastic process ${\boldsymbol{X}}$ is called rectifiable if $v^{\boldsymbol{X}}$ exists and is locally bounded,and Equation (15)has an unique solution.A coupling $(X_{0},X_{1})$ is called rectifiable if its linear interpolation process ${\boldsymbol{X}}$, following $X_{t}=tX_{1}+(1-t)X_{0}$, is rectifiable.In this case, we call ${\boldsymbol{Z}}=\mathtt{Rectflow}({\boldsymbol{X}})$ the rectified flow of $(X_{0},X_{1})$, and write it (with an abuse of notation) as ${\boldsymbol{Z}}=\mathtt{Rectflow}((X_{0},X_{1}))$.The corresponding$(Z_{0},Z_{1})$ is called the rectified coupling of $(X_{0},X_{1})$, denoted as $(Z_{0},Z_{1})=\mathtt{Rectify}((X_{0},X_{1}))$.

Assume that ${\boldsymbol{X}}$ is rectifiable. We have

Hence,rectified flowturns a rectifiable stochastic process into a flow while preserving the marginal laws.

We show thatthe rectified flow ${\boldsymbol{Z}}$ of ${\boldsymbol{X}}$achieves the minimum of the path-wise $c$-transport cost in the set of time-differentiable stochastic processes whose expected velocity field equals $v^{{\boldsymbol{X}}}$.This explains that the property of non-increasing convex transport costs of rectified flow/coupling.

The rectified flow ${\boldsymbol{Z}}=\mathtt{Rectflow}({\boldsymbol{X}}_{t})$ in (15) attains the minimum of

which yields a proof of Theorem 3.2 of that the rectified coupling $(Z_{0},Z_{1})$ yields no larger convex transport costs than $(X_{0},X_{1})$.

A primal-dual relation

Let us generalize the least squares loss $L_{{\boldsymbol{X}}}(v)$ in (12) to aa Bregman divergence based loss:

For any differentiable convex function $c$,and rectifiable process ${\boldsymbol{X}}$, we have

andthe optima above are achieved when $v=v^{{{\boldsymbol{X}}}}$ and ${\boldsymbol{Y}}=\mathtt{Rectflow}({\boldsymbol{X}})$.

Straight couplings

Take ${\pi}_{0}={\pi}_{1}=\mathcal{N}(0,I)$.Hence, for $c(x)=\left\lVert x\right\rVert^{p}$ with $p>0$,the $c$-optimal mapping is the trivial identity coupling $(X_{0},X_{0})$ with $X_{0}\sim{\pi}_{0}$.However,consider the coupling $(X_{0},AX_{0})$,where $A$ is anon-identity and non-reflecting rotation matrix (namely $A^{\top}A=I$, $\det(A)=1$, $A\neq I$ and $A$ does not have $\lambda=-1$ as an eigenvalue).Then $(X_{0},AX_{0})$ is a straight coupling of ${\pi}_{0}$ and ${\pi}_{1}$,but it is not $c$-optimal for allsecond order differentiable convex function $c$ whose Hessian matrix is invertible everywhere. See Appendix for the proof.It is the rotation transformthat makes $(X_{0},AX_{0})$ sub-optimal,which is removed in the proposed $c$-rectified flow in Section 5 via a Helmholtz like decomposition.

Differentiable Processes with Equivalent Marginal Laws

The marginal preserving property of rectified flow is due to the property of $v^{{\boldsymbol{Z}}}=v^{{\boldsymbol{X}}}$ by construction.However, we show in this section that$v^{{\boldsymbol{X}}}=v^{{\boldsymbol{Z}}}$ is only a sufficient condition:two differentiable processes ${\boldsymbol{X}}$ and ${\boldsymbol{Z}}$can have the same marginal lawseven if $r\coloneqq v^{{\boldsymbol{X}}}-v^{{\boldsymbol{Z}}}\neq 0$.This is because $r$, as illustrated in Example 3.5,can be a rotation-only vector field (in a generalized sense shown below)that introduces rotation components into the dynamics without modifying the marginal distributions.Therefore, the constraint of $v^{{\boldsymbol{Y}}}=v^{{\boldsymbol{X}}}$ inthe optimization problem (16) may be too restrictive.A natural relaxation of (16) would be

which yields a dynamic OT problem with a continuum ofmarginal constraints.In Section 5, we show thatthe solution of (17) yields our $c$-rectified flow thatsolve the OT problemat the fixed point.Solving (17) allows us to remove the rotational components of $v^{{\boldsymbol{X}}}$, which is whatwhat renders rectified flow non-optimal.In this section,we first characterize the necessary and sufficient conditionfor having equivalent marginal laws.

which gives $\nabla\cdot(r_{t}\varrho_{t})=0$. This says that $r_{t}\varrho_{t}$ is a rotation-only (or divergence-free) vector field in the classical sense.

which is the definition of${\boldsymbol{Y}}$-marginal-preserving following (18).∎

c𝑐c-Rectified Flow

We introduce $c$-rectified flow,a $c$-dependent variant of rectified flow that guarantees tominimize the $c$-transport costwhen applied recursively. This section is organized as follows:Section 5.1defines and discusses the $c$-rectified flow of a differentiable stochastic process ${\boldsymbol{X}}$, which we show yields the solution of theinfinite-marginal OT problem (17).Section 5.2considers the $c$-rectified flow of a coupling $(X_{0},X_{1})$, which we show is non-increasing on the $c$-transport cost.Section 5.3proves that the fixed points of $c\text{-}\mathtt{Rectify}$ are $c$-optimal.Section 5.4interprets $c$-rectified flowas an alternating direction descent methodfor the dynamic OT problem (8),and a majorize-minimization (MM) algorithm for the static OT problem (1).Section 5.5 discussesa key lemma relating $c$-optimal couplings and its associated displacement interpolation with Hamilton-Jacobi equation.

Note that we have ${\mathsf{m}_{c}}\!\left(x;y\right)\geq 0$ for $\forall x,y$ following the definition of the conjugate $c^{*}$ (or the Fenchel-Young inequality).Losses of form ${\mathsf{m}_{c}}\!\left(x;y\right)$is equivalent to the so called matching lossproposed forlearning generalized linear models . Compared with the original rectified flow,the difference of $c$-rectified flow is i) restricting the velocity field to a form of $g_{t}=\nabla c^{*}\circ\nabla f_{t}$, and ii) replacing the quadratic objective function to the matching loss.These two changes combined yield a Helmholtz like decomposition of $v^{{\boldsymbol{X}}}$ as we show below, allowing us to remove the “rotation-only” component of $v^{{\boldsymbol{X}}}$ and obtain $c$-optimal couplings at fixed points.

We can equivalently write (23) using Bergman divergence associated with $c$, that is,

Then it is easy to see that ${\mathsf{m}_{c}}\!\left(x;y\right)={\mathsf{b}_{c}}\!\left(x;\nabla c^{*}(y)\right)$, by using the fact that $\nabla c(\nabla c^{*}(y))=y$ and $c^{*}(y)=y^{\top}\nabla c^{*}(y)-c(\nabla c^{*}(y))$.Hence, ${\mathsf{m}_{c}}$ and ${\mathsf{b}_{c}}$ are equivalent up to the monotonic transform $\nabla c^{*}$ on $y$.The minimum ${\mathsf{b}_{c}}\!\left(x;y\right)=0$ is achieved when $y=x$,while ${\mathsf{m}_{c}}\!\left(x;y\right)=0$ is achieved when $\nabla c^{*}(y)=x$.Therefore, (23) is equivalent to

Moreover,the generalized Pythagorean theorem of Bregman divergence (e.g., ) gives

which can be viewed as projecting the expected velocity $v^{{\boldsymbol{X}}}_{t}$ to the set of functions of form $g_{t}=\nabla c^{*}\circ\nabla f_{t}$,w.r.t. the Bregman divergence.This yields an orthogonal decomposition of $v_{t}^{{\boldsymbol{X}}}$:

where $r^{{{\boldsymbol{X}}},c}_{t}=v^{{{\boldsymbol{X}}},c}_{t}-\nabla c^{*}\circ\nabla f^{{{\boldsymbol{X}}},c}_{t}$ is the residual term.The key result below shows that $r^{{{\boldsymbol{X}}},c}$ is ${{\boldsymbol{X}}}$-marginal-preserving, which ensures thatthe $c$-rectified flow preserves the marginals of ${{\boldsymbol{X}}}$.

We say that ${\boldsymbol{X}}$ is $c$-rectifiable if $v^{{\boldsymbol{X}}}$ exists,the minimum of (23) exists and is attained by a locally bounded function $f^{{\boldsymbol{X}},c}$,and the solution of Equation (22) exists and is unique.

Taking $g_{s}=h$ if $s<t$ and $g_{s}=0$ if $s>t$ yields that $r^{{{\boldsymbol{X}}},c}(x)=\nabla c^{*}(\nabla f_{s}^{{{\boldsymbol{X}}},c}(X_{s}))-v^{{\boldsymbol{X}}}(X_{s})$ is ${\boldsymbol{X}}$-marginal-preserving following (18).ii) Note that ${\boldsymbol{Z}}$ is rectifiable if ${\boldsymbol{X}}$ is $c$-rectifiable. Applying Lemma 4.2yields the result. ∎

For the quadratic cost $c(x)=c^{*}(x)=\frac{1}{2}\left\lVert x\right\rVert^{2}$, the $\nabla c^{*}$ is the identity mapping, and (27) reduces to the Helmholtz decomposition,which represents a velocity field into the sum of a gradient field and a divergence-free field. Hence,(27) yields a generalization of Helmholtz decomposition, in which a monotonic transform $\nabla c^{*}$ is applied on the gradient field component.We call (27) aBregman Helmholtz decomposition.

Remark: score matching

In some special cases, $v^{{\boldsymbol{X}}}$ may already be a gradient field,and hence the rectified flow and $c$-rectified flow coincide for $c(x)=\frac{1}{2}\left\lVert x\right\rVert^{2}$. One example of this is when $X_{t}=\alpha_{t}X_{1}+\beta_{t}\xi$for some time-differentiable functions $\alpha_{t}$ and $\beta_{t}$, and $\xi\sim\mathcal{N}(0,I)$, satisfying $\alpha_{1}=1,\beta_{1}=0$, and $X_{0}=\alpha_{0}X_{1}+\beta_{0}\xi$. In this case, one can show that

c𝑐c-Rectified flow solves Problem (17)

We are ready to show that the $c$-rectified flow solves the optimization problem in (17).Further, (23) forms a dual problem of(17).

Under the conditions in Theorem 5.2, we havei) ${\boldsymbol{Z}}=c\text{-}\mathtt{Rectify}({\boldsymbol{X}})$ attains the minimum of (17).ii) Problem(17) and (23) has a strong duality:

As the optima above are achieved by $f^{{{\boldsymbol{X}}},c}$ and ${\boldsymbol{Z}}$, we have$L_{{{\boldsymbol{X}}},c}(f^{{{\boldsymbol{X}}},c})=F_{c}({\boldsymbol{X}})-F_{c}({\boldsymbol{Z}}).$

Moreover, if we take${\boldsymbol{Y}}={\boldsymbol{Z}}$ and $f=f^{{{\boldsymbol{X}}},c}$,then the inequality in $\overset{(1)}{\leq}$ is tight because $\dot{Z}_{t}=\nabla c^{*}(\nabla f_{t}(Y_{t}))$ holds $t$-almost surely.Therefore, $R_{{{\boldsymbol{X}}},c}({\boldsymbol{Z}})=L_{{{\boldsymbol{X}}},c}(f^{{{\boldsymbol{X}}},c})\geq R^{{{\boldsymbol{X}}},c}(Y)$, which suggests that ${\boldsymbol{Z}}$ attains the maximum of $R_{{{\boldsymbol{X}}},c}$ (under the marginal constraints) and the strong duality holds.∎

Similar to the case of rectified flow,the $c$-rectified flow/coupling of a coupling $(X_{0},X_{1})$is defined as the $c$-rectified flow/coupling of its linear interpolation process.In the following, we show that the $c$-rectified coupling of a coupling yields no larger$c$-transport cost.

Let ${\boldsymbol{X}}$ be the linear interpolation ofcoupling $(X_{0},X_{1})$ in that $X_{t}=tX_{1}+(1-t)X_{0},\forall t\in$. We say that $(X_{0},X_{1})$ is $c$-rectifiable if ${\boldsymbol{X}}$ is $c$-rectifiable, and call ${\boldsymbol{Z}}=c\text{-}\mathtt{Rectflow}({\boldsymbol{X}})$ the $c$-rectified flow of $(X_{0},X_{1})$.We call the induced $(Z_{0},Z_{1})$ the $c$-rectified coupling of $(X_{0},X_{1})$, denoted as $(Z_{0},Z_{1})=c\text{-}\mathtt{Rectify}((X_{0},X_{1}))$.

which establishes that $(Z_{0},Z_{1})$ yields no larger transport cost than $(X_{0},X_{1})$.

Compared with the regular $\mathtt{Rectify}$ mapping,the key difference here is thatthe monotonicity of $c\text{-}\mathtt{Rectify}$ only holds for the specific $c$ that it employees, rather than all convex cost functions. More importantly,as we show below,recursively applying $c\text{-}\mathtt{Rectify}$yields optimal couplings w.r.t. $c$,a key property that the regular rectified flow misses.

3 Fixed Points ofc𝑐c-𝚁𝚎𝚌𝚝𝚒𝚏𝚢𝚁𝚎𝚌𝚝𝚒𝚏𝚢\mathtt{Rectify} arec𝑐c-Optimal

Knowing that $L_{{{\boldsymbol{X}}},c}(f^{{{\boldsymbol{X}}},c})$ is an indication of $c$-optimality, we show below that it is guaranteed to converge to zero with recursive $\mathtt{Rectify}$ updates.

Let ${\boldsymbol{Z}}^{k}$ be the $k$-th $c$-rectified flow of $(X_{0},X_{1})$, satisfying ${\boldsymbol{Z}}^{k+1}=c\text{-}\mathtt{Rectflow}((Z_{0}^{k},Z_{1}^{k}))$ and $(Z_{0}^{0},Z_{1}^{0})=(X_{0},X_{1})$.Assume each $(Z_{0}^{k},Z_{1}^{k})$ is $c$-rectifiable for $k=0,\ldots,K$. Then

Applying (28) to $(Z_{0}^{k},Z_{1}^{k})$ and $(Z_{0}^{k+1},Z_{1}^{k+1})$ yields

4 c𝑐c-Rectified Flow as Optimization Algorithms

In this section, we draw more understanding on how iterative $c$-rectified flowing solves the static and dynamic OT problems. We first show that $c$-rectified flow can be viewed as an alternative direction descent on the dynamic OT problem(8), and then that $c$-rectified coupling as a majorize-minimization (MM) algorithm on the statistic OT problem (1).The results in this section are framed in terms of a general path-wise loss function $F_{c}({\boldsymbol{Y}})$,and hence provide a useful starting point for deriving $c$-rectified flow like approaches tomore general optimization problems with coupling constraints.

The mapping ${\boldsymbol{Z}}^{k+1}=c\text{-}\mathtt{Rectflow}({\boldsymbol{Z}}^{k})$ can beinterpreted as an alternative direction descent procedure for the dynamic OT problem (8):

c𝑐c-Rectified flow as an MM algorithm

In this case,the MM update guarantees that $F(X^{k})$ is monotonically non-increasing:

and the minimum is attained by $(X_{0},X_{1})=(Y_{0},Y_{1})$,where $\Pi_{0,1}$ denotes the set of couplings of ${\pi}_{0}$ and ${\pi}_{1}$.ii) $c\text{-}\mathtt{Rectify}$ yields the MM update related $F^{+}$ in that

i) For any coupling $(X_{0},X_{1})$ and $(Y_{0},Y_{1})$, we have

where the inequality holds because remove the constraint ${\boldsymbol{Y}}\in\mathcal{M}_{X}$. In addition, it is obvious that the inequality above becomes equality when $(X_{0},X_{1})=(Y_{0},Y_{1})$.ii) Note that

whose minimum of the right side is attained by ${\boldsymbol{Y}}=c\text{-}\mathtt{Rectflow}((X_{0},X_{1}))$ following Theorem 5.3. Hence, the minimum of the left side is attained by $(Y_{0},Y_{1})=c\text{-}\mathtt{Rectify}((X_{0},X_{1}))$.∎

5 Hamilton-Jacobi Equation and Optimal Transport

On the other hand, define $h_{t}(x)=\partial_{t}f_{t}(x)+c^{*}(\nabla f_{t}(x))$. Then we have

Hence, $(X_{0},X_{1})$ is a $c$-optimal coupling.∎

where we assume that $\lambda_{t}v_{r}\rho_{t}$ decays to zero sufficiently fast at infinity.We have

At the saddle points,the functional derivations of $\mathcal{L}$ equal zero, yielding

Assume $\varrho_{t}$ is positive everywhere and note that $\nabla c^{*}(\nabla c(x))=x$, we have $v_{t}=\nabla c^{*}(\nabla\lambda_{t})$,and hence $\nabla\lambda_{t}^{\top}v_{t}-c(v_{t})=c^{*}(\nabla\lambda_{t})$.Plugging it back to $\frac{\delta\mathcal{L}}{\delta\rho_{t}}=0$ yields that $\dot{\lambda}_{t}+c^{*}(\nabla\lambda_{t})=0$.Overall, the (formal) KKT condition of (10) is

Discussion and Open Questions

For machine learning (ML) tasks such as generative models and domain transfer, the transport cost is not necessarily the direct object of interest.In these cases,as suggested in, rectified flow might be preferred because it is simpler and does not require to specify a particular cost $c$.Question: for such ML tasks, when would it be preferred to use OT with a specific $c$, and how to choose $c$ optimally?

In practice, recursively applying the ($c$-)rectification accumulates errorsbecause the training optimization for the drift field and the simulation of the ODEcan not be conducted perfectly.Howto avoid the error accumulationat each step?Assume $\{x_{1,i}\}_{i}\sim{\pi}_{1}$, and $\{z_{0,i}^{k},z_{1,i}^{k}\}_{i}$ is obtained by solving the ODE of the $k$-th $c$-rectified flow starting from $z_{0,i}^{k}\sim{\pi}_{0}$.As we increase $k$, $\{z_{0,i}^{k}\}_{i}$ may yield increasingly bad approximation of ${\pi}_{1}$ due to the error accumulation. One way to fix this is to adjust $\{z_{1,i}^{k}\}$ to make it closer to $\{x_{1,i}^{k}\}_{i}$ at each step. This can be done by reweighting/transporting $\{z_{1,i}^{k}\}_{i}$ towards $\{x_{1,i}^{k}\}_{i}$ by minimizing certain discrepancy measure,or replacing each $z_{1,i}^{k}$ with $x_{\sigma(i)}^{k}$ where $\sigma$ is a permutation that yields a one-to-one matching between $\{z_{1}^{(i)}\}$ and $\{x_{1}^{(i)}\}_{i}$.The key and challenging part is to do the adjustment in a good and fast way, ideally with a (near) linear time complexity.

With or without the adjustment step,build a complete theoretical analysis on the statistical error of the method.

In what precise sense is rectified flow solving a multi-objective variant of optimal transport?

References

Appendix A Proofs

i)The fact that $A^{\top}A=I$ and ${\pi}_{0}={\pi}_{1}=\mathcal{N}(0,I)$ensures that $AX_{0}\sim{\pi}_{1}$ and hence$(X_{0},AX_{0})$ is a coupling of ${\pi}_{0}$ and ${\pi}_{1}$.Let $X_{t}=tAX_{0}+(1-t)X_{0}$ be the linear interpolation of the coupling. Related, we have $\dot{X}_{t}=AX_{0}-X_{0}$. Canceling $X_{0}$ yields that

where $c^{*}$ is the convex conjugate of $c$.Equation (33) is equivalent to $\nabla c(Ax-x)=\nabla\phi(x)$, which means that $\nabla\phi$ is continuously differentiable.Taking gradient on both sides of (33) gives

where $H_{x},B_{x}$ are both symmetric and $H_{x}$ is positive definite,and henceThen $H_{x}B_{x}$ is a diagonalizable (all its eigenvalues are real) by Lemma A.1.However, $A-I$ is not diagonalizable because $A$ must have complex eigenvalues as a non-reflecting, non-identity rotation matrix. Hence, (34) can not hold.

Assume that $A,B$ are two real symmetric matrices and $A$ is positive definite. Then $AB$ is diagonalizable (on the real domain), that is, there exists an invertible matrix $P$, such that $P^{-1}ABP$ is a diagonal matrix.

This is a standard result in linear algebra.Because $A$ is positive definite, there exists an invertible symmetric matrix $C$, such that $CC=A.$ Then, $AB=CCB$, and it is similar to $CBC^{-1}$, which is symmetric and hence diagonalizable.∎