[論文レビュー] Rectified Flow: A Marginal Preserving Approach to Optimal Transport

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

Stochastic Gradient Optimization Techniques被引用数 10

ひとこと要約

この論文は、入力分布 $\pi_0$ および $\pi_1$ を保存する一方で、ユーザーが指定した凸コスト関数 $c$ を反復的に最小化する、新しいマージナル保存型フローサンプル法を導入する。整列フローを特定のコスト関数に適応させることで、ニューラル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$ をもつ単一目的の最適輸送に適応させること。
• 高次元連続設定におけるデータ駆動型最適輸送を実現するため、単純で制約なしの回帰問題の系列を用いること。

提案手法

• 本手法は、時間 $t \in [0,1]$ の間に結合 $(X_0, X_1)$ を進化させる一連のニューラル常微分方程式（ODE）を構築し、マージナル法則 $\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$ をターゲット $\pi_1$ に順列付けまたは不一致最小化により一致させることで軽減可能である。
• 本手法は、マージナル忠実度を維持するための線形時間調整ステップを備えた高次元連続設定におけるデータ駆動型OTを可能にする。