[论文解读] Rectified Flow: A Marginal Preserving Approach to Optimal Transport

Qiang Liu|arXiv (Cornell University)|Sep 29, 2022

Stochastic Gradient Optimization Techniques被引用 10

一句话总结

本文提出了一种新颖的边际保持流方法，用于最优传输（OT），通过迭代最小化用户指定的凸成本函数 $c$，同时保持输入分布 $ olimits_0$ 和 $ olimits_1$ 不变。通过将修正流（rectified flow）适配至单一成本函数，该方法在神经ODE上通过无约束回归构建单调的内部优化过程，确保在无需外部约束强制的情况下收敛至最优耦合。

ABSTRACT

We present a flow-based approach to the optimal transport (OT) problem between two continuous distributions $π_0,π_1$ on $\mathbb{R}^d$, of minimizing a transport cost $\mathbb{E}[c(X_1-X_0)]$ in the set of couplings $(X_0,X_1)$ whose marginal distributions on $X_0,X_1$ equals $π_0,π_1$, respectively, where $c$ is a cost function. Our method iteratively constructs a sequence of neural ordinary differentiable equations (ODE), each learned by solving a simple unconstrained regression problem, which monotonically reduce the transport cost while automatically preserving the marginal constraints. This yields a monotonic interior approach that traverses inside the set of valid couplings to decrease the transport cost, which distinguishes itself from most existing approaches that enforce the coupling constraints from the outside. The main idea of the method draws from rectified flow, a recent approach that simultaneously decreases the whole family of transport costs induced by convex functions $c$ (and is hence multi-objective in nature), but is not tailored to minimize a specific transport cost. Our method is a single-object variant of rectified flow that guarantees to solve the OT problem for a fixed, user-specified convex cost function $c$.

研究动机与目标

• 为解决在连续最优传输中强制执行无穷维边际约束的问题，而不依赖对偶性或正则化。
• 开发一种单调的内部优化方法，使其在耦合的可行集内移动以降低传输成本。
• 将原始为多目标设计的修正流方法，适配至具有固定、用户指定成本函数 $c$ 的单目标OT问题。
• 在高维连续设置中实现数据驱动的最优传输，通过一系列简单且无约束的回归问题完成。

提出的方法

• 该方法构建一系列神经常微分方程（ODE），在时间 $t \in [0,1]$ 内演化耦合 $(X_0, X_1)$，同时保持边际分布律 $\mathrm{Law}(X_0) = \pi_0$ 和 $\mathrm{Law}(X_1) = \pi_1$。
• 在每一步中，通过最小二乘回归学习速度场 $v^X_t$，以预测沿线性插值路径 $X_t = tX_1 + (1-t)X_0$ 的位移 $X_1 - X_0$。
• 由此产生的ODE流 $\mathrm{d}Z_t = v^X_t(Z_t)\mathrm{d}t$ 确保在整个变换过程中边际分布保持不变。
• 该方法是修正流的一种变体，其目标是特定成本函数 $c$，通过 $c$-修正流过程将多目标问题转化为单目标问题。
• 优化过程以迭代方式进行，每一步求解一个无约束回归问题，从而实现 $c$-传输成本的稳定且单调下降。
• 理论分析表明，该方法满足OT问题的KKT条件，最优映射表现为凸共轭函数的梯度。

实验结果

研究问题

• RQ1 能否构建一种单调的内部方法用于连续最优传输，以避免依赖外部约束强制？
• RQ2 如何将修正流方法适配以最小化特定的、用户定义的凸成本函数 $c$，而非同时最小化所有凸成本？
• RQ3 迭代 $c$-修正流过程的收敛行为和最优性间隙如何？
• RQ4 在递归应用 $c$-修正流时，如何减轻误差累积？
• RQ5 修正流在何种意义上求解了最优传输的多目标变体？在何种情况下，具有固定 $c$ 的单目标OT更优？

主要发现

• 该方法通过一系列无约束回归问题实现 $c$-传输成本的单调降低，确保收敛至最优耦合。
• 通过ODE流自动保持边际约束，无需使用拉格朗日对偶或熵正则化。
• 理论分析证实，该方法满足OT问题的KKT条件，最优映射为 $v_t = \nabla c^*(\nabla \lambda_t)$，且 $\dot{\lambda}_t + c^*(\nabla \lambda_t) = 0$。
• 该方法为极小化极大法和基于正则化的方法提供了一种实用替代方案，后者常因不稳定性和约束强制不足而表现不佳。
• 通过重加权或使用置换匹配或差异最小化将最终 $Z_1^k$ 与目标 $ olimits_1$ 对齐，可减轻迭代应用中误差的累积。
• 该方法实现了高维连续设置下的数据驱动最优传输，且可通过线性时间调整步骤保持边际保真度。