[論文レビュー] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

Xingchao Liu, Chengyue Gong+1 more|arXiv (Cornell University)|Sep 7, 2022

Model Reduction and Neural Networks被引用数 75

ひとこと要約

Rectified flowは、直線経路に従うことで2つの経験的分布間を輸送するODEを学習し、単純な非線形最小二乗法目的で達成される。再帰的整流によりほぼ直線の流れが得られ、生成と転送タスクのための高速な1ステップ推論を可能にする。

ABSTRACT

We present rectified flow, a surprisingly simple approach to learning (neural) ordinary differential equation (ODE) models to transport between two empirically observed distributions π_0 and π_1, hence providing a unified solution to generative modeling and domain transfer, among various other tasks involving distribution transport. The idea of rectified flow is to learn the ODE to follow the straight paths connecting the points drawn from π_0 and π_1 as much as possible. This is achieved by solving a straightforward nonlinear least squares optimization problem, which can be easily scaled to large models without introducing extra parameters beyond standard supervised learning. The straight paths are special and preferred because they are the shortest paths between two points, and can be simulated exactly without time discretization and hence yield computationally efficient models. We show that the procedure of learning a rectified flow from data, called rectification, turns an arbitrary coupling of π_0 and π_1 to a new deterministic coupling with provably non-increasing convex transport costs. In addition, recursively applying rectification allows us to obtain a sequence of flows with increasingly straight paths, which can be simulated accurately with coarse time discretization in the inference phase. In empirical studies, we show that rectified flow performs superbly on image generation, image-to-image translation, and domain adaptation. In particular, on image generation and translation, our method yields nearly straight flows that give high quality results even with a single Euler discretization step.

研究の動機と目的

• 輸送マップフレームワークの下で生成モデルとドメイン転送を統合する。
• 分布間の直線的で最短経路の流れを好む、スケーラブルなODEベースのアプローチを導入する。
• 輸送を導く速度場を学習するための単純な最小二乗目的関数を開発する。
• 理論的性質を示す：非交差流、周辺分布の保存、削減された凸輸送コスト。
• 画像生成、画像間翻訳、ドメイン適応における実証的有効性を示す。

提案手法

• Rectified flowを、結合データからサンプリングされた X_1 - X_0 に沿ってほぼ直線方向を辿る ODE dZ_t = v(Z_t, t) dt として pi_0 を pi_1 に輸送するものとして定義する。
• X_1 - X_0 と v(tX_1 + (1 - t)X_0, t) の不一致を [0,1] の t について最小化する非線形最小二乗目的関数を解くことで速度場 v を学習する。
• X_t の周辺法則を保存しつつ、交差する軌跡を再配線して非交差流を得ることで因果性を保証する。
• 学習した v を用いた pi_0 および pi_1 からの前向き・後向きサンプリングを使用する。より直線的な経路を得るための任意の reflow を適用する。
• 高速な1ステップ推論のために、k-rectified flow を模倣するニューラルネットワーク T の蒸留を任意で行う。
• 線形補間を一般の時間次微分可能な X_t に置換した非線形拡張を議論し、確率流ODEとDDIMとの関連を述べる。

実験結果

リサーチクエスチョン

• RQ1 最小限の計算オーバーヘッドで、2つの経験的分布間で質量を移動させる単純なODEベースの輸送モデルを学習できるか。
• RQ2 Rectified flows は凹凸のない、決定論的な結合を生み出し、凸コストに対する輸送コストを低減するか。
• RQ3 再帰的整流によってますます直線的な軌道を生み出し、正確な単一步推論または数ステップ推論を可能にするか。
• RQ4 生成と転送タスクにおいて、rectified flowsは確率流ODEsやノイズ除去拡散モデルとどう関連し、どのように異なるか。
• RQ5 Rectified flows は画像生成、画像間翻訳、ドメイン適応に有効か。

主な発見

• Rectified flow は、敵対的訓練や尤度に基づく推論を必要とせず、経験的分布間の輸送を学習するための単純な最適化ベースの方法を提供する。
• 得られた流れは非交差で、線形補間過程の周辺法を常に保存する。
• 整流は、すべての凸コスト c に対して入力結合よりも低い凸輸送コストをもつ結合を生み出す。
• 再帰的整流はますます直線的な流れを生み出し、粗い時間離散化でも正確なシミュレーションを可能にし、潜在的には1ステップ推論を実現する。
• 経験的に、Rectified flow は高品質な画像生成と翻訳成果を達成し、報告された最先端の1ステップ CIFAR10 FIDスコア4.85 および recall 0.51 を含む。