Flow Maps: 拡散モデルの積分を学習する

拡散モデルからのサンプリングは反復的なプロセスです。各ステップで、ディノイザー（去噪器）が入力空間内のある経路上の接線方向を見積もります。我々はこの方向に沿って小さなステップを繰り返し踏み出すことで経路を進み、本質的に雑音レベル全体にわたる積分を計算することになります。これにより、単純な雑音分布からのサンプルが目標分布からのサンプルへと徐々に変換され、それらをつなぐ経路が描かれていくのです。

我々は、この積分を直接予測するニューラルネットワークを学習させることはできるのでしょうか？サンプリング速度を向上させるためにです。「はい、可能です」という歓迎の言葉を添えつつ、「フローマップ（Flow Maps）」の世界へようこそ！

拡散モデルが登場して以降、人々はそれらのサンプリングをより高速かつ低コストにする方法を模索してきました。約 2 年前、私は「拡散蒸留（diffusion distillation）」についてブログ記事を書きました。これは、高品質なサンプルを得るために必要なステップ数を削減するための主要なツールの一つです。蒸留手法の基本原理に変化はないものの、以来多くの新しいバリエーションが登場しました。

本記事では、フローマップに深く掘り下げてみたいと思います。拡散モデルは経路間の関係を、経路上の各点における接線方向を予測することによって記述するのに対し、フローマップはその同じ経路上的任意の点から任意の他の点を予測することができます。これらは高速なサンプリングのために使用されますが、より効率的な報酬ベースの学習や改善されたサンプリングのガイド可能性（steerability）などを可能にする他の特性も備えています。最近、研究対象として非常に注目されています。

フローマップを定義することは比較的容易ですが、構築および訓練する方法はさまざま存在します。さらに、拡散モデル自体と同様に、文献には異なる形式主義や用語が溢れており、すべてがどのように結びついているかを理解しようとする際には混乱を招くことになります。私は、主に Boffi らによって提案された分類学に基づいて、できるだけ状況を明確にしていきます。

フローマップは拡散モデルのアイデアの上に構築されており、通常通りこれらに対する一定の知識を前提とします。ベクトル微積分に慣れていることが訓練方法の理解にも役立ちますが、それが難しいという方には、本記事のもう他の部分も興味深い内容であるはずです。文脈のために、以前のいくつかのブログ記事（例えば「拡散に関する視座」など）を参照・再確認することを検討してください。また、Chieh-Hsin Lai らは最近、数学的厳密さと直感的説明を組み合わせた包括的な拡散モデルに関する独著を発表しており、リファッシャーとしても入門点としても大変おすすめです。

以下に目次を示します。クリックして特定のセクションへ飛び越すことができます。

• 雑音からデータへの経路の描画
• 一貫性の三つの概念
• バックプロパゲーションすべきか、するべきでないか？
• スケッチからのフローマップ訓練
• 実践におけるフローマップ
• アプリケーションと拡張
• 代替戦略
• 総括的な考察
• 謝辞
• リファレンス

雑音からデータへの経路を描画する

フローマップを理解するための鍵は、拡散モデルを「雑音とデータの間の双射（bijection）」として捉える視点です。つまり、各分布のサンプル対をつなぐ一意の経路が存在し、それらは互いに交わらないように定義されているのです。したがって、まずは拡散サンプリングアルゴリズムに詳しく見て、そこからフローマップへと発展していきましょう。

拡散モデルからのサンプリング

現在、拡散モデルには多数の異なるサンプリングアルゴリズムがありますが、それらはいずれも「確率的（stochastic）」か「決定論的（deterministic）」かの二つのカテゴリーに分類されます。決定論的サンプリングの神話は以前書いたことがありますのでここで触れる必要はありませんが、フローマップの発展にとって基礎的であり、概説する価値があります。

要点は以下の通りです：ノイズ観測値 $\mathbf{x}_t$ が与えられた場合、ディノイザーモデルを使ってクリーンな元のデータの期待値 $\hat{\mathbf{x}}_0 = \mathbb{E}\left[ \mathbf{x}_0 \mid \mathbf{x}_t \right]$ を予測できるのであれば、二つの異なる反復的な生成手順を構築できます。

確率的な方法は最も直感的です：各反復で、現在のノイズ観測値 $\mathbf{x}t$ 条件下でのわずかにノイズの少ない例の条件分布 $p(\mathbf{x}{t-1} \mid \mathbf{x}_t)$ からサンプリングし、汚染プロセスを一度に一つずつ逆転させます。都合の良いことに、ディノイザーモデルによる予測 $\hat{\mathbf{x}}_0$ を用いてこの分布の近似を構築できます。ステップ $t$ と $t-1$ の間のノイズレベルの区間が小さいほど、近似はより正確になります。多くの反復の後、ノイズが消え、$t=0$ でクリーンなデータ分布からのサンプルが得られます。これは、簡単に言えば元来の DDPM アルゴリズムの仕組みです。拡散モデルの常微分方程式 (SDE) 形式主義に基づくサンプリングアルゴリズムも、入力空間で同様の確率的経路を生み出します。

決定論的な手順は、非常に初期を除き、いずれかの時点でランダムなサンプルを引き出すことを含みません：現在のノイズ観測値 $\mathbf{x}_t$ とディノイザーからの予測 $\hat{\mathbf{x}}0$ が与えられた場合、決定論的な更新規則により $\mathbf{x}{t-1}$ が得られ、これを再帰的に適用して最終的に $\mathbf{x}_0$ に到達します。手順の各ステップが決定論的であるため、どこにもランダム性はありません：与えられた始点 $\mathbf{x}_t$ から出発すれば、必ず特定のエンドポイント $\mathbf{x}_0$ しか到達しません。この更新規則は確率的枠組み（つまり DDIM）で導出することも、常微分方程式 (ODE) 形式主義を用いても行うことができます。

Flow Matching でデフォルトで使用されるサンプリングアルゴリズムも、決定論的な手順の別の例です。ここでは、ニューラルネットワークは通常、クリーン入力 $\mathbb{E}\left[ \mathbf{x}_0 \mid \mathbf{x}_t \right]$ ではなく速度 $\mathbf{v}_t = \mathbb{E}\left[\mathbf{x}_T - \mathbf{x}_0 \mid \mathbf{x}_t \right]$ を予測するようにパラメータ化されます（ここで $t=T$ は最大のノイズレベルに対応し、純粋なガウス雑音です）。ただし、$\mathbf{v}_t$, $\hat{\mathbf{x}}_0$, および $\mathbf{x}_t$ の間に線形関係があるため、これは単に同じ基礎アルゴリズムのバリエーションに過ぎません（詳細は以前のブログ記事における異なる拡散モデル出力パラメータ化についての議論も参照）。

これらすべてのアルゴリズムが共通しているのは、各時間ステップ $t$ でのノイズ例の周辺分布 $p(\mathbf{x}_t)$ が保持されることです：決定論的か確率的かのサンプリングアルゴリズムを選択した与否に関わらず、$\mathbf{x}_t$ の分布は変わりません！もちろん、これは決定論的なケースではデルタ分布に収束する（全確質量が単一のオプション上に集中する）条件分布 $p(\mathbf{x}_t \mid \mathbf{x}_T)$ については完全に異なります。周辺分布の保持という性質は、それぞれデータ側と雑音側の特殊なケース $p(\mathbf{x}_0)$ および $p(\mathbf{x}_T)$ にも当てはまります。しかし、個々の例を特定するのではなく分布全体を見ると、サンプリングプロセスによって入力空間で描かれる経路は非常に異なるように見えます。

以下にサンプリング過程の視覚化を示します：左側に確率的、右側に決定論的な場合です。私は両方のケース（1 次元例：上部と 2 次元例：下部）を表示することにしました。提供する洞察は補完的であるためです。どちらのケースにおいても、ターゲット分布は 2 つのガウスの混合体です。サンプリングを開始すると、我々は単一のガウス分布からのサンプルでスタートし、分布が徐々にターゲットの混合体へと変換されていきます。経路上を単なるサンプルが辿るパスは、確率的なケースではかなりギザギザで不規則ですが、決定論的なケースでは滑らかで緩やかに曲がったものになります。まったく異なる微視的な行動が、全く同じ巨視的な行動を生み出します！

推計航行：ディノイザーによる経路の追跡

決定論的サンプリングアルゴリズムが存在することから重要な帰結が導かれます：個々の雑音サンプルとデータサンプルの間には、必ず一意の双射的なマッピングが存在するということです。各雑音サンプルは単一の特定のデータサンプルに関連付けられ、逆もまた然りです。雑音サンプルから始めて、対応するデータサンプルへと導く入力空間内を辿る経路に従うことができます。これを行うには、各点における経路の接線方向に従うだけで十分であり、これはディノイザーによって予測されます。注意すべきことは、同じ接線方向を使用して、逆方向（データから雑音へ）も経路に沿って誘導できるということです。

以下の図は、雑音分布からのサンプル $\mathbf{x}_T$、対応するデータサンプル $\mathbf{x}_0$、それらを結ぶ入力空間内の経路、そして経路上の中間点 $\mathbf{x}_t$ を示しています。また、この時点におけるディノイザー予測 $\hat{\mathbf{x}}_0$（これは経路に対する接線方向に対応する）も示されています。私が以前「ガイダンスの幾何学」や「蒸留」について書いた記事を読んでいる場合、おそらくこのような図は馴染みがあるはずです。以前の記事には、2 次元直感から結論を導く際の危険性についての警告が含まれており、再度言及すべきです：高次元オブジェクトを 2 次元で表現する際には慎重な注意が必要です！

ディノイザー予測を使用してこれらの経路を Traverse することは「記憶lessness」（メモリーレス）です：ディノイザーへの入力となるのは、入力空間における現在の位置と現在の雑音レベルのみであり、これに基づいて移動方向 $\hat{\mathbf{x}}_0 = f(\mathbf{x}_t, t)$ を予測します。また、これは「近見性（myopic）」でもあります：ディノイザーは最終的な目的地 $\mathbf{x}_0$ を先に見ることはできず、単に次にどこへ向かうかを示すだけです。他の情報の使用もできません：過去に訪問した位置や予測された方向、始点または終点は無視され、サンプリングプロセスにおける現在の位置しか考慮されず、それ以外は何もない。この経路の性格付けは、推計航行を連想させます。

したがって、このような方法で接続されている特定の雑音とデータサンプル対間の経路は一意である必要があります：もしいくつかの経路が同じデータサンプルへと導くのであれば、これらの経路が分岐する点では複数の有効な接線方向が存在することになります。同様の理由から、異なるサンプル対間の経路は決して交差することはできません：交差点に曖昧さをもたらすからです。ディノイザーは、現在の位置しか知らず、どの経路上にいるか分からないため、複数の交差する経路を区別することは不可能です。これは以下の図に示されています。

技術的には、この議論は経路が $(\mathbf{x}_t, t)$-空間において交わらないことを示すだけですが、理論的には異なる時間ステップ $t$ で入力空間内の同一点に到達する 2 つの経路が存在する場合、$\mathbf{x}_t$-空間で交差することは可能です。実際には、このエッジケースを無視できます：十分に異なる 2 つの時間ステップにおけるノイズ中間サンプルの分布 $p(\mathbf{x}_t)$ は実質的に重複しないためです。事実、一部の最近の論文では、現在の雑音レベルをディノイザーに与えないことも同様に、あるいはよりよく機能すると提唱しており、高次元の入力空間では、$\mathbf{x}_t$ 自体から雑音レベルを推論できるためです。

実用上で経路が決して交差しないという事実は、ディノイザーを用いた記憶lessness な Traverse を可能にします。経路は ODE ベースのサンプリングの文脈では解の軌跡（solution trajectories）として知られており、常微分方程式を解くことによって Traverse されます。

経路が曲がっているため、サンプリングでは「落下」したり、異なる経路上で終わったりしないように確保するためには、理想的には無限数の微小ステップを踏むべきです。実際には、我々は小さな但有ったステップを踏み、これはサンプリングの過程で蓄積する可能性があります近似誤差をもたらします。近似の質は、我々が踏むステップ数と経路の曲がり具合に依存します。より曲がっているほど、良好な近似を得るために多くのステップが必要です。

幸运的是、通常、計算上処理可能な数のステップ（しばしば 100 未満）で十分な結果を得ることができます。それでもなお、人々は高速なサンプリングを可能にするために経路の曲がり具合を最小化しようとしてきました。これはフローマッチングの動機のいくつかであり（実際にどの程度達成できるかは激しく議論されています）、Reflow プロシージャにおいても同様です。Reflow は、「再配線」により、どのデータサンプルがどの雑音サンプルと接続されるかを変更することでより直線的な経路を取得します。

測地学：フローマップによる経路のマッピング

ディノイザーモデルを使用して経路上の任意の点における接線方向を予測することを学ぶことは、その経路を完全に特徴付ける方法の一つです。しかし、この目標を達成するための唯一の方法ではありません：フローマップは魅力的な代替手段を提供します。経路上の任意の点において、それらの別の点を予測できます。

すでに $f(\mathbf{x}_t, t)$ をディノイザーを表すために使用したので、$F(\mathbf{x}_s, s, t)$ をフローマップを表すために使用しましょう。ここで注意すべきことは、2 つの時間ステップを inputs として取ることです：$s$ と $t$ はそれぞれソースとターゲットの雑音レベルに対応します。データと雑音の間の双射がある場合、理想的なフローマップは経路上のどこからでも別の場所にジャンプすることを可能にします：$F(\mathbf{x}_s, s, t) = \mathbf{x}_t$ 。通常、我々はノイズからデータへ移動することに興味があるので $s > t$ ですが、必ずしもそうである必要はありません。実際には、拡散モデルを訓練する際にディノイザーを行うように、ニューラルネットワークでこの関数を近似します。

以下では、フローマッチングで一般的に使用されるノイズスケジュールを使用していると仮定します：$\mathbf{x}_t = (1 - t)\mathbf{x}_0 + t\mathbf{\varepsilon}$ および $T=1$、$\mathbf{\varepsilon} \sim \mathcal{N}(0, 1)$（標準ガウス雑音）。これはおそらく現在最も人気のある選択肢です。シンプルに保つためです。より一般的な設定ですべてを導出することは可能です（単に $\mathbf{x}_t = \alpha(t) \mathbf{x}_0 + \sigma(t)\mathbf{\varepsilon}$ および任意の $T$ を仮定すれば）が、それは数学を複雑にし、追うのが難しくなります。注意点として、我々は時間方向に関する従来の拡散の規約に従うことにします（つまり $t=0$ がデータ分布に対応し、$t=1$ が雑音に対応する）。これはフローマッチング論文で使用されている規約とは逆です。これらの選択の影響について詳しくは、ノイズスケジュールに関する私のブログ記事を参照してください。

これらの選択に基づき、ディノイザー $f(\mathbf{x}_t, t)$（クリーン入力 $\hat{\mathbf{x}}_0 = \mathbb{E}\left[\mathbf{x}_0 \mid \mathbf{x}_t\right]$ を予測する）が与えられている場合、経路に対する接線方向または速度 $\mathbf{v}_t$ は以下のようになります：

$$ \mathbf{v}_t = v(\mathbf{x}_t, t) = \dfrac{\mathbf{x}_t - f(\mathbf{x}_t, t)}{t} . $$

フローマッチングの設定では、通常ニューラルネットワークはクリーン入力ではなく関数 $v(\mathbf{x}_t, t)$ を直接予測するようにパラメータ化されますが、互いに線形関数であるため片方からもう片方を容易に得ることができます。

次に、積分によって時間区間を積分することで単純にフローマップを構築できます：

$$ F(\mathbf{x}_s, s, t) = \mathbf{x}s + \int_s^t v(\mathbf{x}\tau, \tau) \mathrm{d} \tau . $$

この積分は、経路に沿って無限数の微小ステップを踏むことを表し、進むにつれて予測された接線方向 $v(\mathbf{x}_t, t)$ を累積します。始点 $\mathbf{x}_s$ にこの積分を加えることで、$\mathbf{x}_t$ に着地します。

典型的なケースでは、ノイズからデータへ移動するため $s > t$ です（拡散の規約において $t = 0$ がデータ側に対応するため）。これは、拡散が前向き汚染プロセスとして定義されていることを反映しており、実際にはデータ分布からサンプリングすることは逆方向に進むことを意味します。我々は $\mathbf{v}_t$ を $\hat{\mathbf{x}}_0$ から $\mathbf{x}_t$ へ向かうように規約で定義しましたので、データ側へ移動するにはこのベクトルの逆方向に従う必要があります。

いくつかの特別なケースを強調する価値があります：

• $t=0$ と設定すると、経路上の任意の場所からデータ側の終端点へ直接ジャンプできます：$F(\mathbf{x}_s, s, 0) = \mathbf{x}_0$ 。これを正確に行うことができれば、これは単一ステップでのサンプリングを可能にします。これがまさに「一貫性モデル（consistency models）」が doing のことです。フローマップはそのアイデアの一般化であり、後で詳しく議論します。
• $s=t$ と設定すると、積分する区間の長さがゼロになるため、積分自体はゼロになり、したがって $F(\mathbf{x}_t, t, t) = \mathbf{x}_t$ となります。
• 通常、ノイズからデータへの経路 Traverse に興味がある（つまり $t < s$ ）場合でも、必ずしもそうである必要はありません。他の方向へ行くことも同じ式を使用して行うことができます（$t > s$ を選択することで）。例として、$F(\mathbf{x}_s, s, 1)$ は $\mathbf{x}_s$ を含む経路の雑音側の終端点を予測します。

おそらく、関数 $F(\mathbf{x}_s, s, t)$ を予測することを学習することは、$f(\mathbf{x}_t, t)$ を予測することを学習するよりも難しい任务であることは自明でしょう：それは単一の時間ステップ入力ではなく二つの時間ステップ入力を持つためです。少なくとも、これはデータと雑音サンプル間の経路のグローバルな特徴付けを提供します。また、実用的にも大いに有益です：フローマップを一度に持てば、サンプリング中に小さすぎるステップを踏む必要や、経路から落ちることを気にする必要はありません。実際、ニューラルネットワークの近似が十分良ければ、単に雑音 $\mathbf{\varepsilon}$ をサンプリングし、$s=1$ から $t=0$ へ直接一歩踏み出すだけで $\mathbf{x}_0$ に到着し、サンプリングは完了です！次のセクションでは、フローマップモデルをどう訓練するかについて議論します。

拡散モデルがクリーン入力の期待値 $\hat{\mathbf{x}}_0$ または速度 $\mathbf{v}_t$ を予測するようにパラメータ化することが一般的であるように、フローマップについても 2 つの同等なパラメータ化があります。これまで記述してきた $F(\mathbf{x}_s, s, t)$ は経路上の目的地を予測しますが、経路沿いの平均速度または平均フローも予測することもできます：

$$ V(\mathbf{x}s, s, t) = \dfrac{1}{t - s} \int_s^t v(\mathbf{x}\tau, \tau) \mathrm{d} \tau . $$

2 つのパラメータ化間の関係は以下の通りです：

$$ F(\mathbf{x}_s, s, t) = \mathbf{x}_s + (t - s) V(\mathbf{x}_s, s, t) . $$

ここで、極限ケース $s=t$ は $V(\mathbf{x}_t, t, t) = v(\mathbf{x}_t, t)$ を与え：長さゼロの区間上の平均速度は単に瞬間的な速度です。これは、フローマップにはディノイザーが含まれており、したがって標準的な拡散モデルとしても使用できることを示しています。

グローバルな視点がデータと雑音サンプル間の経路を提供するため、フローマップは拡散モデルを時代遅れにするものだと考えられがちですが、前者は後者の厳密な一般化であり、多くの実用的メリットがあります。しかし我々が見るように、これまでこのグローバルな視点を構築するために開発されたすべてのアプローチは、拡散モデルによって提供されるローカルな視点からのブートストラップによって機能します。時にはこの関係は明示的であり、時にはそれほど明白ではありませんが、常に存在します。機械学習では常に、「無料のランチ」はありません：フローマップを使用してサンプリングすることは拡散モデルからのサンプリングよりも安価ですが、フローマップを訓練するにははるかに複雑で、多くの場合拡散モデルを最初に訓練する必要があります。正確な地図を描画することがナビゲーションを容易にするが、事前に多くの手間がかかるのと同じです！

一貫性の三つの概念

フローマップを訓練するための多くの異なるアルゴリズムが提案されています。結局のところ、これらすべてのバリエーションは、密接に関連する 3 つの一貫性則の一つに基づいていることが判明しました：合成性（compositionality）、ラグランジュ的な一貫性（Lagrangian consistency）、そしてオイラー的な一貫性（Eulerian consistency）。このセクションでは、それぞれについてカバーし、それらをフローマップ訓練にどう使用するかについて議論します。

Boffi, Albergo, および Vanden-Eijnden は当初フローマップフレームワークを開発し、これら 3 つの規則（およびそこから導出された訓練手順）をフローマップマッチングと自己蒸留に関する 2 つの最近の論文で記述しました。彼らの仕事は「確率的インターポラント（stochastic interpolant）」の視点に根ざしていますが、ここではそれを採用せず、代わりにより従来の拡散の枠組みに留めます：多くの人がそれに関係があると考えられています。

合成性（Compositionality）

フローマップ $F(\mathbf{x}_s, s, t)$ を使用すると、経路上の $\mathbf{x}_s$ から同じ経路上的 $\mathbf{x}_t$ へ直接移動できます。同じ手順を繰り返してより遠くへ移動することもでき、$F(\mathbf{x}_t, t, u)$ を使用して $\mathbf{x}_u$ へ行けますが、1 つのステップ $F(\mathbf{x}_s, s, u)$ を使用することもできます。経路を Traverse するいずれかの方法でも同じ結果を得るべきです：

$$ F(F(\mathbf{x}_s, s, t), t, u) = F(\mathbf{x}_s, s, u) = \mathbf{x}_u . $$

つまり、フローマップは合成的です。私はこれを「合成性」と呼んでいます：非標準的な用語です。これは執拗であるが、文献で使用されているさまざまな名称を曖昧で混乱させるため、そうしています。あなたはこれを「半群性質（semigroup property）」、「ショートカット性質（shortcut property）」または「プログレッシブマッチング/蒸留（progressive matching / distillation）」と呼ばれることがあります。

帰結として、フローマップは自身で逆関数となります（最初の引文に関して）：

$$ F(F(\mathbf{x}_s, s, t), t, s) = \mathbf{x}_s . $$

この場合、我々は $s > t$ および $t > s$ の両方でフローマップが定義されていると仮定しています。しかし、多くの場合、フローマップは 1 つの方向（$s > t$：高ノイズレベルから低ノイズレベルへ）のみ訓練されます：それがサンプリング（データ分布への移動）に relevance があるためです。

合成性を使用して拡散モデルからのブートストラップでフローマップを訓練できます。我々は $\mathbf{x}_s$ から始め、拡散モデルを使用して経路上の次の点 $\mathbf{x}_t$ を予測し、わずかな距離先を進みます。次に、フローマップは出発点に関係なく常に同じ答えを与えるべきだという事実に頼ることができます：$F(\mathbf{x}_s, s, u) = F(\mathbf{x}_t, t, u)$、そして特殊なケース $t = u$ では：$F(\mathbf{x}_s, s, t) = F(\mathbf{x}_t, t, t)$ 。これらの等式が成り立つことを保証することで、フローに関する情報を小さな時間区間から大きな時間区間に転送できます。

ラグランジュ的な視点：ゴールポストを移動させる

フローマップ $F(\mathbf{x}_s, s, t)$ の一貫性を特徴付ける別の方法は、その出力がどのように変化するかを研究することです（ここでは $t$、つまり目的地のインデックスを徐々に変更します）。これにより、出力 $\mathbf{x}_t$ が経路を沿って移動するはずです。$t$ の無限小変化を検討すると、微分を使用して起こり得ることを特徴付けることができます：

$$ \dfrac{\mathrm{d}}{\mathrm{d} t} F(\mathbf{x}_s, s, t) = \dfrac{\mathrm{d}\mathbf{x}_t}{\mathrm{d} t} = \mathbf{v}_t . $$

つまり：フローマップ出力の瞬間的な変化は速度です。直感的には、$t$ を変更することは単に経路を Traverse することを意味し、速度はその軌跡に従うために進むべき正確な方向であることを示しています。

速度 $\mathbf{v}_t = v(\mathbf{x}_t, t) = v(F(\mathbf{x}_s, s, t), t)$ を展開でき、これは拡散モデル $v(\mathbf{x}_t, t)$ からのフローマップ学習をブートストラップするための別の手段を与えます。我々は単に以下の等式がすべてで成り立つことを保証する必要があります：

$$ \frac{\partial}{\partial t} F(\mathbf{x}_s, s, t) = v(F(\mathbf{x}_s, s, t), t) , $$

ここでフローマップの $t$ に関する全微分が偏導関数に等しい（他の引数は $t$ に依存しないため）ことを使用しています：$\frac{\mathrm{d}F}{\mathrm{d}t} = \frac{\partial F}{\partial t}$ 。

ラグランジュ的な一貫性の別の解釈は、それが合成性の特殊ケースであることです（我々は 2 つ目の時間区間を無限小に縮小させ、極限振る舞いを検討している）：$u \to t + \Delta t$ と置き換えて明示的にします。

$$ F(F(\mathbf{x}_s, s, t), t, t + \Delta t) = F(\mathbf{x}_s, s, t + \Delta t) . $$

これは $\Delta t = 0$ でも真です：

$$ F(F(\mathbf{x}_s, s, t), t, t) = F(\mathbf{x}_s, s, t) . $$

元の等式からこの特殊ケースを減算し、$\Delta t$ で割ることによって：

$$ \dfrac{F(F(\mathbf{x}_s, s, t), t, t + \Delta t) - F(F(\mathbf{x}_s, s, t), t, t)}{\Delta t} = \dfrac{F(\mathbf{x}_s, s, t + \Delta t) - F(\mathbf{x}_s, s, t)}{\Delta t} . $$

最後に、$\Delta t \to 0$ の極限を取り、微分の定義を使用します：

$$ \left. \dfrac{\mathrm{d}}{\mathrm{d} u} F(F(\mathbf{x}s, s, t), t, u) \right\vert{u=t} = \dfrac{\mathrm{d}}{\mathrm{d} t} F(\mathbf{x}_s, s, t) . $$

左辺を単純化するために、元のフローマップ定義 $F(\mathbf{x}_s, s, t) = \mathbf{x}s + \int_s^t v(\mathbf{x}\tau, \tau) \mathrm{d} \tau$ を思い出して対応する導関数を取ります：

$$ \dfrac{\mathrm{d}}{\mathrm{d} t} F(\mathbf{x}_s, s, t) = \dfrac{\mathrm{d}}{\mathrm{d} t} \left( \mathbf{x}s + \int_s^t v(\mathbf{x}\tau, \tau) \mathrm{d} \tau \right) = v(\mathbf{x}_t, t) , $$

ここで $\frac{\mathrm{d}}{\mathrm{d}t} \mathbf{x}_s = 0$ と、微積分基本定理を使用しています。この単純化を適用することで、再び見つけます：

$$ v(F(\mathbf{x}_s, s, t), t) = \dfrac{\partial}{\partial t} F(\mathbf{x}_s, s, t) . $$

オイラー的な視点：賞を眼下に見る（Eyes on the prize）

ターゲット時間ステップ $t$ の影響を見る代わりに、$s$ 即ち出発点がどのように変わるかを研究することもできます。一見するとさらにはっきりとしています：

$$ \dfrac{\mathrm{d}}{\mathrm{d} s} F(\mathbf{x}_s, s, t) = 0 . $$

出発点を変更しますが、ターゲット時間ステップ $t$ は同じままである場合、目的地はまったく変化しないはずです。したがって、その導関数はゼロである必要があります。簡単で十分、ですか？しかし、この表面的な単純さは欺くものです。我々は現在、$\mathbf{x}_s$ および実際の入力空間における出発位置の 2 つが入力として $s$ に依存しています。

3 つの関数入力のうち 2 つが $s$ に依存するためには、多変数の連鎖律を使用して解く必要があります：

$$ \dfrac{\mathrm{d}}{\mathrm{d} s} F(\mathbf{x}s, s, t) = \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, t) \dfrac{\mathrm{d} \mathbf{x}_s}{\mathrm{d}s} + \dfrac{\partial}{\partial s} F(\mathbf{x}_s, s, t) = 0. $$

これは基本的には 2 つの変化を組み合わせたものです：出発時間ステップ $s$ の変更から入力空間で起こる変化と、出発時間ステップ自体の変更です。

$\frac{\mathrm{d} \mathbf{x}_s}{\mathrm{d}s} = v(\mathbf{x}_s, s) = \mathbf{v}_s$ であることに注意し、拡散モデルからのフローマップ学習をブートストラップするための別の等式を得ることができます：

$$ \dfrac{\partial}{\partial s} F(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) = 0. $$

ラグランジュ的な一貫性の場合と同様に、オイラー的な一貫性を合成性の特殊ケースとして解釈できます。今度は第 1 の時間区間を無限小に縮小させ、$s \to t$ を与えます：合成のルールをもう一度見直して、$t$ を $s + \Delta s$ に置き換えます。

$$ F(F(\mathbf{x}_s, s, s + \Delta s), s + \Delta s, u) = F(\mathbf{x}_s, s, u) . $$

$\Delta s$ が非常に小さいため、以下の近似を使用できます：

$$ F(\mathbf{x}_s, s, s + \Delta s) = \mathbf{x}s + \int_s^{s + \Delta s} v(\mathbf{x}\tau, \tau) \mathrm{d} \tau \approx \mathbf{x}_s + v(\mathbf{x}_s, s) \Delta s, $$

ここで $v(\mathbf{x}_\tau, \tau)$ が積分区間を越えて一定であると仮定しています。$\Delta s \to 0$ とする予定があるため、これは正当な仮定です。我々は次に：

$$ F(\mathbf{x}_s + v(\mathbf{x}_s, s)\Delta s, s + \Delta s, u) = F(\mathbf{x}_s, s, u) . $$

左辺で $(\mathbf{x}_s, s)$ における一次多変数テイラー展開を実行して：

$$ F(\mathbf{x}s, s, u) + \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, u) v(\mathbf{x}_s, s)\Delta s + \dfrac{\partial}{\partial s} F(\mathbf{x}_s, s, u) \Delta s . $$

注意してください $F(\mathbf{x}_s, s, u)$ は最初の項として現れ、かつ我々の前の式の右側にも出現するため、これらは相殺されます。残っているのは：

$$ \nabla_{\mathbf{x}_s} F(\mathbf{x}_s, s, u) v(\mathbf{x}_s, s)\Delta s + \dfrac{\partial}{\partial s} F(\mathbf{x}_s, s, u) \Delta s = 0 . $$

今では単に $\Delta s$ で除してオイラー的な一貫性のルールを回復します：

$$ \dfrac{\partial}{\partial s} F(\mathbf{x}s, s, u) + \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, u) v(\mathbf{x}_s, s) = 0. $$

明示的に極限 $\Delta s \to 0$ をどこでも取っていませんが、非常に小さい場合에만 valid な近似に頼っています。

オイラー的とラグランジュ的な一貫性は、結局のところ同じものの異なる視点であり、異なる参照系を使用しています。ラグランジュ的な一貫性では、特定のノイズ入力例に焦点を当て、フローマップの出力が時間とともにどのように進化するかを追跡します。オイラー的な一貫性では、ターゲット時間ステップを固定し、入力の変更とともに物事がどのように変化するかを評価します。フローが川であれば、ラケット（カヌー）に乗ってその経路に従うこと（ラグランジュ的）と、橋の上で下を眺めること（オイラー的）の間の違いです。

等式からのロス関数の構築

これら 3 つの一貫性ルールを表す方程式は、ある程度冗長であり、ほぼ自明のように感じられます：任意の有効なフローマップに対して真であることが明確です。しかしニューラルネットワークは $s, t$, および $\mathbf{x}_s$ の 3 つの入力のほとんど任意の関数を学習する十分な柔軟性があり、これらの可能性のほとんどは有効なフローマップが持つように一貫した方法ではありません。したがって、フローマップを学ぶ際には、明示的に一貫性ルールを実行することが有用です。

実はどれでも良いことが判明しました：我々が先ほど議論した 3 つの一貫性のルールのいずれかに従う関数は、適切な境界条件と組み合わせて自動的に有効なフローマップとなります。これは実際には、フローマップを訓練するためのロス関数を構築する際により多くのオプションを与えます。

一貫性ルールはすべて等式です：これらをロス関数に変換するのは比較的簡単です：すべての項を左側に移動して右側をゼロにします。左側は現在残差（residual）であり、我々が達成しきれない距離を測定しています。次に、単にその残差をペナルティにして、損失が最小化されたときにゼロに近いようにするだけです。それを実現する最も簡単な方法は、単に左辺を二乗して、すべての可能な時間ステップ組み合わせ（およびトレーニングデータセット）で平均化して損失関数を得ることです。

3 つの一貫性ルールに対して、それぞれ得られます：

$$ \mathcal{L}_{\mathrm{compositional}} = \mathbb{E} \left[ \left( F(F(\mathbf{x}_s, s, t), t, u) - F(\mathbf{x}_s, s, u) \right)^2 \right], $$

$$ \mathcal{L}_{\mathrm{Lagrangian}} = \mathbb{E} \left[ \left( \frac{\partial}{\partial t} F(\mathbf{x}_s, s, t) - v(F(\mathbf{x}_s, s, t), t) \right)^2 \right], $$

$$ \mathcal{L}_{\mathrm{Eulerian}} = \mathbb{E} \left[ \left( \dfrac{\partial}{\partial s} F(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right)^2 \right]. $$

これらのすべてのロス関数の極小値は一貫性を保証します。実際にはこれらの関数を完璧に最小化する事ができなくても、通常は期待通りに機能するだけ十分に近づけることができます。

何か有用なものを学習するために、$F(\mathbf{x}_t, t, t) = \mathbf{x}_t$ を制約し、$v(\mathbf{x}_t, t)$ が有意義な速度に対応することを保証します。これは最初に拡散モデルを訓練して参照（蒸留）として使用することで達成できますが、意味を持つ速度を実際に制約する他の方法もあります（第 4 セクションでフローマップから始められる方法）。

残差を二乗することはある程度の任意の選択です：絶対値をペナルティすることも、Huber loss のようなエキゾチックなものをを使用することもできます。いくつかの場合には、後ほど見るようにカテゴリカルクロスエントロピーさえ使用できます。平均二乗誤差（MSE）アプローチにはいくつかの実用的な利点があります：勾配降下法によって比較的容易に最適化でき、いくつかのスクラッチ訓練方法にとって必須です（第 4.2 セクション）。

バックプロパゲーションすべきか、するべきでないか？

これらのロス関数を詳しく見てみると、いくつかの異質な点があります：

• そのうち 2 つは我々が学習しようとする関数 $F$ の微分を含んでいます（ラグランジュ的とオイラー的）。これは、勾配ベースの学習が潜在的に高階導関数に関与することを意味します。
• 他のバリエーションは $F$ の複数の順次適用を含み、訓練中に順次的な前向きおよび後ろ向きパスを必要とする可能性があります。

機械学習で使用されるほとんどのロス関数は、モデル予測と静的ターゲット（「真実」）の間の差を測定する一方、これらはその移動ターゲットを含み自己参照的です。理論的には、勾配ベースの最適化はこのことに関心はありません：単に与えられた関数の極大値を探そうとします（通常は局所極大値）。しかし、フローマップ訓練を静的ターゲットを持つより従来の機械学習フレームワークに転換することで、いくつかの障害を克服できます：高階導関数を計算する必要を回避することなどです。

勾配の流れを塞ぐ（Flow of gradients を塞ぐ）

表現学習からの灵感を得ることができます：これらの自己参照的な移動ターゲットを含むロス関数のタイプがますます一般的になっています。ここでは、1 つのネットワークが別のネットワークの出力を模倣するように学習し、蒸留における生徒と教師のような関係です。教師は生徒と同じパラメータを使用して構築されます。通常、指数移動平均（EMA）がパラメータに使用され、損失の教師側を通じて勾配がバックプロパゲーションされないことが多く、退化的な解への収束を回避するのに役立ちます。

同様のトリックをフローマップ訓練を安定化し簡略化するために使用できます。損失の一部を停止勾配（stop-gradient）操作でラップできます。これはバックプロパゲーション中の勾配の流れをブロックし、それ以外ではパススルーとして機能します：

$$ \mathcal{L}_{\mathrm{Lagrangian}} = \mathbb{E} \left[ \left( \frac{\partial}{\partial t} F(\mathbf{x}_s, s, t) - v(\mathrm{sg} \left[ F(\mathbf{x}_s, s, t) \right], t) \right)^2 \right], $$

$$ \mathcal{L}_{\mathrm{Eulerian}} = \mathbb{E} \left[ \left( \dfrac{\partial}{\partial s} F(\mathbf{x}s, s, t) + \mathrm{sg} \left[ \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right] \right)^2 \right], $$

ここで $\mathrm{sg}[\cdot]$ は停止勾配操作を示します。ラップされた内部のものはすべてバックプロパゲーションの目的では一定として扱われるため、例えば $\nabla_{\mathbf{x}_s} F(\mathbf{x}_s, s, t)$ を通ってバックプロパゲーションする必要を回避できます。同様に合成の場合、順次的な後ろ向きパスを回避するために停止勾配操作を使用できます：

$$ \mathcal{L}_{\mathrm{compositional}} = \mathbb{E} \left[ \left( \mathrm{sg} \left[ F(F(\mathbf{x}_s, s, t), t, u) \right] - F(\mathbf{x}_s, s, u) \right)^2 \right]. $$

これはエレガントな解釈があります：2 つの順次的フローマップステップを使用してターゲットを計算し、それを真実として扱って凍結し、その後フローマップを更新して 1 ステップでそこに到達する方法を学習します。

停止勾配操作でラップされた損失の一部は実質的に静的として扱われる（実際には技術的にはそうではない）ため、EMA パラメータを使用して計算することで訓練を安定化できる場合があります。これは、それが training の過程でよりゆっくり変化することを保証し、それが静的であるという暗黙的な仮定をより重大でないものにします。

停止勾配操作を導入することは興味深い帰結を持っています：損失の一部のみを通じてバックプロパゲーションしたことで計算された「勾配」方向は実際には勾配方向ではありません！少なくとも、最適化しようとする損失の勾配方向ではありません：他のロス関数に対して有効な勾配である可能性があります。これは有時半勾配（semigradient）と呼ばれています。これは、いくつかの理論的保証が勾配ベースの最適化について窓の外に出ることを意味します。幸運にも、注意深く行う場合、理論的な基盤からの安全性を放棄することは実践的に大きな問題を引き起こさないようです（ニューラルネットワークのようなことが多くあるように）。

上記に与えられた損失バリエーションは単なる例です：どの部分の損失式が停止勾配操作でラップされているか、または EMA パラメータを使用して安定化されているかは、フローマップ訓練のさまざまな風味を特徴づけるものです。第 5 セクションでこの設計空間を詳しく探求します。

「平均速度」の視点

ここで、以前 1.3 セクションで議論したフローマップの平均速度パラメータ化を思い出すと便利になります：ラグランジュ的とオイラー的な一貫性のルールの中の微分と興味深い方法で相互作用するためです：

$$ V(\mathbf{x}s, s, t) = \dfrac{1}{t - s} \int_s^t v(\mathbf{x}\tau, \tau) \mathrm{d} \tau , $$

$$ F(\mathbf{x}_s, s, t) = \mathbf{x}_s + (t - s) V(\mathbf{x}_s, s, t) . $$

置換によってラグランジュ的な一貫性のルールを $V$ について表現できます：

$$ \frac{\partial}{\partial t} \left( \mathbf{x}_s + (t - s) V(\mathbf{x}_s, s, t) \right) = v( F(\mathbf{x}_s, s, t), t) . $$

$v$ の最初の引数の置換を実行しなかったのは、何も単純化しないためです。次に左辺の時間導関数を計算できます：積の法則を必要とします：

$$ \frac{\partial}{\partial t} \left( \mathbf{x}_s + (t - s) V(\mathbf{x}_s, s, t) \right) = V(\mathbf{x}_s, s, t) + (t - s) \dfrac{\partial}{\partial t} V(\mathbf{x}_s, s, t) . $$

$V$ の時間導関数の他に、$V$ 自体がこの式に含まれることに注意してください。項を整理して $V$ を孤立させると：

$$ V(\mathbf{x}_s, s, t) = v( F(\mathbf{x}_s, s, t), t) - (t - s) \dfrac{\partial}{\partial t} V(\mathbf{x}_s, s, t) . $$

これは次のように解釈できます：$s$ と $t$ の間の時間区間上の平均速度は、終端点での速度から、ターゲット時間ステップ $t$ に関する平均速度自体の導関数を含む補正項を減じたものです。我々がこれを使用して損失を構築する場合、停止勾配操作で右辺全体をラップできます。つまり、時間導関数をバックプロパゲーションする必要を気にする必要はなく、損失を最適化するために高階微分は必要ありません。

オイラー的な一貫性のルールでも完全相同的ことができます：

$$ \dfrac{\partial}{\partial s} \left( \mathbf{x}_s + (t - s) V(\mathbf{x}s, s, t) \right) + \nabla{\mathbf{x}_s} \left( \mathbf{x}_s + (t - s) V(\mathbf{x}_s, s, t) \right) v(\mathbf{x}_s, s) = 0. $$

積の法則（2 回）を使用して：

$$

• V(\mathbf{x}_s, s, t) + (t - s) \dfrac{\partial}{\partial s} V(\mathbf{x}_s, s, t) + v(\mathbf{x}s, s) + (t - s) \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) = 0. $$

項を整理して $V$ を孤立させると：

$$ V(\mathbf{x}_s, s, t) = v(\mathbf{x}_s, s) + (t - s) \left( \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right) . $$

これは平均速度を、出発点での速度から、$s$ 時間ステップに関する平均速度自体の導関数を含む補正項を加算したものを表します。我々は再び停止勾配操作で右辺全体をラップできます：MeanFlow の基礎となるものになります：

$$ \mathcal{L}_\mathrm{MF} = \ \mathbb{E} \left[ \left( V(\mathbf{x}_s, s, t) - \mathrm{sg} \left[ v(\mathbf{x}_s, s) + (t - s) \left( \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right) \right] \right)^2 \right] . $$

前向きおよび逆モードの微分化（Forward- and reverse-mode differentiation）

ニューラルネットワーク訓練のための現代フレームワークは勾配を計算します：ほとんど心配する必要はありませんが、これを可能にする自動微分機械はかなり複雑です。

深層計算グラフに対する勾配を計算するには、主に 2 つの方法があります：前向きモードと逆モードの微分です。それらは入力で出力へ、および出力で入力へとグラフを Traverse します。それらの間の選択は入力と出力の次元性に帰着します：出力が入力よりも高次元の場合、前向きモードが効率的です。他の場合、逆モードが勝ります。ニューラルネットワークを訓練する場合、入力は通常高次元ですが、微分する計算グラフの最終的な出力はほぼ必ず単一のロス値です。それはスカラーなので、出力次元は入力次元よりも低く、逆モードが正しい選択です。これはフレームワークがデフォルトで使用します。

前向きモードは時々登場しますが：Jacobian ベクトル積（JVP）を効率的に計算するために使用できます。そのような積はオイラー的な一貫性のルールで発生します：

$$ \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) = \left[ \dfrac{\partial V}{\partial \mathbf{x}_s} , \dfrac{\partial V}{\partial s} , \dfrac{\partial V}{\partial t} \right] \left[ v, 1, 0 \right] . $$

$V$ の Jacobian はすべての偏導関数からなり、それに乗算されるのは「接線ベクトル（tangent vector）」$[v, 1, 0]$ です。JAX では

jax.jvp

を使用してこれを計算できます：それは前向きパスと微分を同時に効率的に計算し、メモリ内の完全な Jacobian マトリックスを明示的に実装しないように避けます。これは贅沢ではありません：$V$ は $\mathbf{x}_s$ と同じ形状なので、両方ともサイズ $K$ のベクトルであれば、$\frac{\partial V}{\partial \mathbf{x}_s}$ は $K \times K$ マトリックスです！

有限差分近似（Finite-difference approximations）

勾配の流れを完全に停止させる代わりに、導関数を扱わないために避けられる別の一般的なトリックは、それらを有限差分近似で置き換えることです：微分の定義を使用できます：

$$ \dfrac{\mathrm{d}}{\mathrm{d}x} f(x) = \lim_{h \rightarrow 0} \dfrac{f(x + h) - f(x)}{h} \approx \dfrac{f(x + \Delta x) - f(x)}{\Delta x} . $$

ここで、$\Delta x$ は有限差分です。それが小さい限り、近似はかなり良いことができます。小さな値は浮動小数点精度の問題を引き起こす可能性があります（特に現在、低精度ニューラルネットワーク訓練が流行しているため）、したがってこのアプローチを使用する際には注意が必要です。

一つ時間区間を無限小に縮小させることでラグランジュ的とオイラー的な一貫性のルールを合成性のルールから導出したのを覚えていますか？いずれかのそれらに有限差分近似を適用することは、その区間を実質的に有限に戻すことに効果があります：これにより、使用する一貫性に基づいた方法の分類をある程度曖昧にします。

実用的考慮事項

本当にこの面倒な勾配操作を必要としているか問いかける価値があります。なぜすべてを通じてバックプロパゲーションすることは問題なのでしょうか？現代フレームワークは確かにこれを可能にし、広範なケースでは簡単ですが、それが常に良いアイデアであるという意味ではありません：

• 高階導関数を計算することは、浮動小数点演算（FLOPS）の点数においてコストがかかりますが、また特にメモリの点においても。それは通常大きなテンソルを長時間保持することを含むため、計算グラフ内の複数の場所で再利用されるためです。
• 通常、現代ニューラルネットワークの高階導関数は非常に意味がありません：2 次導関数は曲率を取り込み、これは入力空間を通じて滑らかに変化しないことが多いです。極端な例として、ReLU 非線形性しかないネットワークは実質的にピースバイース線形的であり、その曲率はほぼどこでゼロです。これも拡散モデルをスカラーエネルギー関数の勾配としてパラメータ化しない典型的な理由です（確かにできます）。
• ますます頻繁に、特定の操作のための特殊な高速カーネルを使用します（例：FlashAttention）。それらは通常、訓練をサポートするための同等に効率的な逆パスの実装を持ってきます。前向きモード微分と高階導関数は通常実装されていませんが、遅い実装へのフォールバックを必要とします。

フローマップ訓練の異なる実装は各トレーニング反復で異なる数の前向きおよび後ろ向きパスを必要とします（例えば有限差分近似は通常逆パスを 2 つの前向きパスに置き換えます）：前向きモード微分や高階導関数を必要とする場合があります、またそうでない場合があります。顕著なケースは Terminal Velocity Matching (TVM) で、ラグランジュ的な一貫性に基づく実装であり、高階導関数を回避するために停止勾配操作や他の近似を使用しません。著者は明示的にこのサポートのためにカスタムアテンションカーネルを開発したと述べます。我々は第 5 セクションでさまざまな実装を詳しく議論します。

スケッチからのフローマップ訓練

雑音とデータサンプル間の経路を表すフローマップを構築することは、何らかの形でのブートストラップを必要とします：例えば、拡散モデルを訓練することにより速度 $v(\mathbf{x}_t, t)$ を私たちに提供し、それだけでその経路を完全に記述するに十分です。次に、それがフローマップ訓練のための開始点として使用でき、それは実質的に蒸留の形に変えます。

しかし、フローマップをスケッチから訓練したい場合はどうすればよいですか？単一ステージの訓練プロセスを好む多くの良い理由があります。いかなる順次依存性も複雑さを大幅に追加し、我々はそれが最終結果の品質を有意に改善する場合のみそれを許容すべきです（偶然にも、これが潜在拡散の場合にそれを許容する理由）。

自己蒸留（Self-distillation）

前述のように、平均速度によってパラメータ化されたフローマップには、特殊ケースとして速度予測器が含まれています：$V(\mathbf{x}_t, t, t) = v(\mathbf{x}_t, t)$ 。したがって、訓練中に時々 $s=t$ をサンプリングし、選択した一貫性ベースのロス関数と、それらの場合に適用された標準拡散ロスを組み合わせることで、$v(\mathbf{x}_t, t)$ の「真実」を提供する事前トレーニング済みモデルを必要としません。両方の損失をバランスさせることで、モデルは瞬間的な速度と有限時間ステップ区間全体のその積分を同時に学習します。例として、ラグランジュ的な一貫性ロスを拡散ロスと組み合わせることができます：

$$ \mathcal{L}_\mathrm{flow,map} = \overbrace{\mathbb{E}\left[ \left( V(\mathbf{x}_t, t, t) - (\mathbf{\varepsilon} - \mathbf{x}_0) \right)^2 \right]}^{\mathrm{diffusion,loss}}\+ \underbrace{ \mathbb{E} \left[ \left( V(\mathbf{x}_s, s, t) - V( F(\mathbf{x}_s, s, t), t, t) + (t - s) \dfrac{\partial}{\partial t} V(\mathbf{x}s, s, t) \right)^2 \right] }{\mathrm{Lagrangian,consistency,loss}} . $$

注意として、ラグランジュ的な一貫性損失項での $v(\mathbf{x}_t, t)$ の出現を $V(\mathbf{x}_t, t, t)$ で置き換えてスケッチから訓練できるようにしました。

このデュアルロス設定を事前トレーニング済み拡散モデルと組み合わせることもできます：$\mathbf{\varepsilon} - \mathbf{x}_0$ をその速度見積もりで置き換えて拡散ロスト項の分散を減らすことができますが、これは厳密には必要ありません。それでも、これを自己蒸留の一種として解釈する意味があります：モデルは同時に教師としても自らへも蒸留されるように訓練されています。

ニューラルネットワーク訓練セットアップにおける私の経験（教師訓練と生徒蒸留が順次ではなく同時である）は、実践的にかなりよく機能し得るということです（私の同僚らはある時点でこのアイデアを表現学習のために使用しました）。結果は通常、2 つの順次トレーニングステージ（まず教師、次に生徒）と同じかほぼ同じですが、順次依存性によって引き起こされる多くの面倒なしです。

周辺からの条件付き学習（Marginal-from-conditional learning）

いくつかのフローマップ訓練の定式化は、スケッチから訓練するために単一のみの一貫性ベースロスのみを必要とする代替アプローチを認めます：これを可能にするために、拡散訓練がどう機能するかを見直す価値があります：ディノイザーは $\mathbb{E}\left[\mathbf{x}_0 \mid \mathbf{x}_t\right]$ を予測する学習をするのに、訓練中に $p(\mathbf{x}_0, \mathbf{x}_t)$ のサンプルで監督されています。それは条件付き期待値を直接予測するために監督されることは決してありませんが、$p(\mathbf{x}_0, \mathbf{x}_t)$ のすべての可能なサンプルにおいて単一の予測を実行させられるため、全体のエラーを最小化するための最良の方法として自動的に期待値に着地します。これはしばしば周辺化トリック（marginalisation trick）として知られており、それは $\mathbf{x}_0$ 条件付き速度から周辺速度の学習を可能にします。

この同じトリックをフローマップ訓練にどのように適用できますか？これを機能させるために 2 つのアプローチがあります：どちらもオイラー的な一貫性のルールから開始します：MeanFlow および改善された MeanFlow (iMF) です：MeanFlow 元のアプローチを見てみましょう。平均速度形式を使用して、我々は：

$$ V(\mathbf{x}_s, s, t) = v(\mathbf{x}_s, s) + (t - s) \left( \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right) . $$

この等式の右辺を学習のターゲットとして扱い、それを停止勾配操作でラップすることで、我々は条件付き速度 $v(\mathbf{x}_s, s)$ を周辺速度 $\mathbf{\varepsilon} - \mathbf{x}_0$（拡散の場合と同様）に置き換えることができます。これは MSE ロスの最小値を変更しないまま残します。なぜこれらを許容できるのかを正確に分解することに少し時間を取れる価値があります。それは 4 つの重要な特徴に基づいています：

1. 我々は平均二乗誤差（MSE）を使用します。
2. 速度は現在のノイズ入力 $\mathbf{x}_s$ で評価されます。
3. 予測ターゲットは速度 $v(\mathbf{x}_s, s)$ について線形です。
4. 停止勾配操作は、結果の更新方向が速度について線形であることを保証します。

この残差を $R$ と呼びましょう：これは一貫性のルールの左辺と右辺の違いです。$R$ は $v(\mathbf{x}_s, s)$ について線形です。我々が最小化する損失関数は単に $\mathbb{E}\left[R^2\right]$ です：この損失関数をモデルパラメータ $\theta$ に関する勾配を取ることによって、我々は得ます：

$$ G_\theta = \nabla_\theta \mathbb{E} \left[ R^2 \right] = \mathbb{E} \left[ 2R \nabla_\theta R \right] . $$

しかし予測ターゲットは停止勾配操作でラップされているため、これは実際上使用する更新方向ではありません。代わりに：

$$ \widetilde{G}\theta = \mathbb{E} \left[ 2R \nabla\theta V \right] , $$

我々は得ます：$V(\mathbf{x}s, s, t)$ は $R$ の唯一の停止勾配操作の外に座っている部分です。したがって、更新方向 $\widetilde{G}\theta$ も速度について線形です。$v(\mathbf{x}_s, s)$ を $\mathbf{\varepsilon} - \mathbf{x}_0$ で入れ替える場合、期待値演算子は結果が同じままになることを保証し、それは $\mathbf{x}_s$ 条件付きに期待されているためです。

停止勾配操作が存在しない場合、これは成り立たないことに注意してください：「正当な」勾配 $G_\theta$ は速度に依存する $R$ および $\nabla_\theta R$ の積を含み、この更新方向は全く速度について線形ではありません：周辺からの条件付き学習のトリックは完全に破綻します！

もし速度が $\mathbf{x}_s$ 以外の場所で評価された場合、それは機能しないでしょう：周辺速度を条件付き速度 $\mathbf{\varepsilon} - \mathbf{x}_0$ で置き換えるのは、$\mathbf{x}_s$ 条件付きに期待を計算しているためです。これがラグランジュ的な一貫性のルールを同様の扱いを与えることができない理由です：それは $t = F(\mathbf{x}_s, s, t)$ での速度評価が必要です。したがって、予測ターゲットも速度について線形であり、停止勾配操作を使用して更新方向が速度について線形であることを保証できるにもかかわらず、期待は置換に機能するために誤った変数条件付きです。

MeanFlow の停止勾配操作は二重任務を実行していると言えます：高階微分を回避し（導関数のバックプロパゲーションなし）、周辺からの条件付き学習を可能にします。一見すると訓練をより効率的にするための調整のように見えますが、それは実際に訓練が機能するために不可欠です。

改善された MeanFlow (iMF) については、同じ平均速度形式のオイラー的な一貫性のルールから開始しますが、 twists：$V(\mathbf{x}_s, s, t)$ と $v(\mathbf{x}_s, s)$ を入れ替えます：

$$ v(\mathbf{x}_s, s) = V(\mathbf{x}_s, s, t) + (t - s) \left( \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right) . $$

次に、$s$ での瞬間速度 $v$ のための式を得ます：$s$ と $t$ の間の区間全体平均速度 $V$ についてです。不幸にも自己参照的です：瞬間的な速度が右辺の Jacobian ベクトル積（JVP）の中に現れます。しかし、瞬間速度も空区間全体の平均速度に過ぎないことを思い出してください：$v(\mathbf{x}_s, s) = V(\mathbf{x}_s, s, s)$ なので：

$$ v(\mathbf{x}_s, s) = V(\mathbf{x}_s, s, t) + (t - s) \left( \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) V(\mathbf{x}_s, s, s) \right) . $$

次に、右側の式を単なる標準拡散またはフローマッチングモデルの再パラメータ化として解釈でき、それを一つのもののように訓練できます。つまり、我々は定義します：

$$ W(\mathbf{x}_s, s, t) = V(\mathbf{x}_s, s, t) + (t - s) \mathrm{sg} \left[ \dfrac{\partial}{\partial s} V(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} V(\mathbf{x}_s, s, t) V(\mathbf{x}_s, s, s) \right] . $$

（混乱を招くことに iMF 論文はこれのために $V$ の記号を使用しますが、私はすでにその文字を平均速度で使用しました。ごめんなさい！）JVP の計算をラップする停止勾配操作の注意してください：我々は通常の MSE ロスで $W$ を予測器として使用できます：

$$ \mathcal{L}_\mathrm{iMF} = \mathbb{E} \left[ \left( W(\mathbf{x}_s, s, t) - (\mathbf{\varepsilon} - \mathbf{x}_0) \right)^2 \right] . $$

「拡散モデル」$W$ を訓練することで、それは区間全体にわたる平均速度 $V$ を学習させられ、したがって完全にそのパラメータ化を通じてフローマップ。どれほど素晴らしいですか？

技術的には、これを機能させるために停止勾配のトリックを必要としませんが、実際には JVP はまだ高階微分を回避するために停止勾配操作でラップされます。従来の MeanFlow とは異なり、しかし停止勾配は方法が正しく機能するために全く必要ではありません！エレガントであるだけでなく、改善された MeanFlow ロスは実際にははるかに低い分散を持つ傾向があります。

実践におけるフローマップ

我々がフローマップがどのようなものか確立し、それらが拡散モデルとどう関係するか、そしてそれらを訓練する方法について議論した後、文献で記述されたいくつかの具体的な実装を詳しく見ていきましょう：いつものように、これは意見のある論文の選択であり、網羅的な概要を提供するつもりはありません。以下コメントに、目立つ欠落（または単に関係のある面白い仕事）を追加してください。これは比較的乾燥しており、私はあなたがセクションの最後に進んで、ここで私がすべてをテーブルで要約することをスキップしても傷つけないでしょう。

あなたが言及されたいくつかの論文を読むことを計画している場合、你可能会会遇到一些记号的变化：

• 時間の方向はデータ（$t=0$）から雑音（$t=1$）へ拡散規約に従うか、またはノイズ（$t=0$）からデータ（$t=1$）へフローマッチングの規約に従うことができます。私は前者に留まりましたが、多くの論文は後者を使用します。
• ソースとターゲット時間ステップは時々逆順で与えられ、まずターゲットを指定し、次にソース、つまり $F(\mathbf{x}_s, t, s)$ 代わりに $F(\mathbf{x}_s, s, t)$ を使用します：時折ターゲット時間ステップは固定されており、したがって省略（例：一貫性モデル）：$F(\mathbf{x}_t, t)$ 。
• 時間ステップは関数の引数（例：$F(\mathbf{x}s, s, t)$ ）ですが、しばしばインデックスとして指定されます（例：$F{s,t}(\mathbf{x}_s)$）。明示的な関数引数を私は好みます：我々はこれらの時間ステップについて（部分的）微分を時々取る必要があるためです。
• フローマップと拡散モデルを表す関数は小文字、大文字、またはギリシャ文字を使用できます：時間ステップはしばしば $s$ および $t$, $t$ および $s$, または $t$ および $r$ であり、標準的な規約はありません。私は $s$ は「ソース」、$t$ は「ターゲット」が好きなので、それにとどまりましたが、多くの論文は実際にはそれを反対方向で使用しています！

このブログ記事の執筆中は、これらの不一致が絶望的に混乱させたいくつかの瞬間がありました：それらを探し出して解釈を正しくすることを確保するために時間を費やすことで、あなたは自分自身に多くの面倒を節約できます。また、パラメータ化の選択（フローマップ $F$, 平均速度 $V$, または他のもの）を念頭にも置くことも重要です：一貫性のルールについて議論した時のように、この選択は式がかなり異なるように見せます。

拡散モデルを訓練することは、考えるとおそらく非常に簡単です：単にガウス雑音と平均二乗誤差ロスという非常に基本的な概念のみが必要です。我々が見たように、フローマップを訓練するとはるかに複雑で、比較します。頻繁には、またコストが高くても、1 つのトレーニングステップを実行するためにモデルを通じて複数のパスを必要とする場合もあります。

ラグランジュ的方法 🐱

Boffi らはラグランジュ的マッピング蒸留（LMD）を記述します：事前トレーニング済み教師モデルに与えられた速度を予測し、最小化します：

$$ \mathcal{L}_{\mathrm{LMD}} = \mathbb{E} \left[ \left( \frac{\partial}{\partial t} F(\mathbf{x}_s, s, t) - v(F(\mathbf{x}_s, s, t), t) \right)^2 \right] . $$

彼らは前向きモード微分（JVP と接線ベクトル $[0, 0, 1]$）を使用して $\frac{\partial}{\partial t} F(\mathbf{x}_s, s, t)$ および $F(\mathbf{x}_s, s, t)$ を同時に効率的に計算することを提唱します：停止勾配操作の欠乏に注意してください、したがって損失関数を最小化することは高階微分を必要とします。損失は $F$ について表現されているが、彼らは予測 $V$ を提唱します。スケッチから訓練のために、自己蒸留バリエーションは速度をフローマップ自身の予測で置き換えることで構築できます（停止勾配操作の導入に注意し）、標準拡散ロスを組み合わせます（第 4.1 セクション）：

$$ \mathcal{L}_{\mathrm{LSD}} = \ \mathbb{E} \left[ \left( \frac{\partial}{\partial t} F(\mathbf{x}_s, s, t) - \mathrm{sg} \left[ V(F(\mathbf{x}_s, s, t), t, t) \right] \right)^2 \right] + \mathbb{E}\left[ \left( V(\mathbf{x}_t, t, t) - (\mathbf{\varepsilon} - \mathbf{x}_0) \right)^2 \right] . $$

Align Your Flow は類似の蒸留アプローチ（AYF-LMD）を提案しますが、合成の視点からそれへ達します：大きなステップ $s$ から $t$ を取ることは、より小さなステップ $s$ から $t - \Delta t$ を取り、その後教師モデルを使用して残りの方法から $t$ へ行くこと（拡散サンプリングステップ）と同等であるべきです：

$$ F(\mathbf{x}_s, s, t) = F(\mathbf{x}_s, s, t - \Delta t) + \Delta t \cdot v(F(\mathbf{x}_s, s, t - \Delta t), t - \Delta t) . $$

彼らは右辺を停止勾配オペレーターでラップして残差の二乗を作り、その後 $\Delta t \to 0$ の極限を取り損失からこれを構築します：これは $\mathcal{L}_\mathrm{LMD}$ を回復することを示し（もちろん停止勾配を除き、それは高階微分を回避するのに役立ちます）：彼らはそれは玩具実験でより安定であることを示しますが、実際の画像では良い結果を生産しないことも指摘しています。

Terminal Velocity Matching (TVM) は類似のレシピに従いますが、自己蒸留を使用したスケッチから訓練をターゲットにします（第 4.1 セクション）。彼らの「終端速度条件」は正確にラグランジュ的な一貫性ルールであり、TVM ロスはラグランジュ的な一貫性項とフローマッチング（拡散）項で構成されます。興味深く、彼らはモデル呼び出しの一部の重みに停止勾配操作を使用することを提案し、そのうちの 1 つに指数移動平均（EMA）重みを適用します：しかし、彼らは一貫性損失項で現れる微分項をこの操作に適用しないため、訓練のために高階微分が必要です。彼らはこれはカスタム FlashAttention カーネルを開発する必要があることを指摘し、それらも公開しています：Lipschitz 連続制約などのアーキテクチャと最適化の調整もあります。

FreeFlow はラグランジュ的な一貫性を起点としてフローマップ蒸留を完全にデータフリーにするための巧妙な方法を見つけます：彼らはノイズ分布からサンプルを引き出すことで拡散モデルを成功してフローマップへ蒸留します。彼らはさらに、トレーニングデータ分布全体を要求しない理由について説得力のある主張を行います：それがその分布で訓練されたとしても、拡散モデルが生成できるサンプルを実質的に代表していない可能性があります！これはクラスファイアードガイダンスのような介入のためでもあり、単にモデルがデータ分布を超えて一般化するのを学んでいるためです。そして時々、元のデータ分布は蒸留の時点で利用できないことがあります。

フローマップ蒸留で使用されるデータ分布が我々がモデル化しようとするサンプリング軌跡を代表していない場合、それは明らかに非最適です。しかし、データなしにニューラルネットワークをどう訓練できますか？彼らは以下の 2 つの要素を組み合わせてこの業績を達成します：

1. ラグランジュ的な一貫性蒸留ロスを、平均速度形式を使用し、ソース時間ステップを $s=1$ に錨定します。彼らは常に純粋な雑音 $\mathbf{\varepsilon} \sim \mathcal{N}(0, 1)$ から始まり、最小化（微分に有限差分近似を使用）：

$$ \mathbb{E} \left[ \left( V(\mathbf{\varepsilon}, 1, t) - \mathrm{sg} \left[ v( F(\varepsilon, 1, t), t) - (t - 1) \dfrac{\partial}{\partial t} V(\varepsilon, 1, t) \right] \right)^2 \right] . $$

1. 補助的なディノイザーモデルは元の腐敗プロセスに従ってそれらを再雑音化することで、1 ステップフローマップサンプル $F(\mathbf{\varepsilon}, 1, 0)$ 上で並行に訓練されています：その後、このディノイザーによって予測された速度を教師速度と比較し、両者の違いを使用してフローマップを更新します。これは、フローマップによって暗示される分布 $p(\mathbf{x}_0)$ を地面に固定するのに役立ちます。

各コンポーネントは単独では良いフローマップモデルを学習するには不十分であることを示します：補助ディノイザーのみを使用すると収束しやすく、ラグランジュ的な一貫性ロスのみを使用すると誤りの蓄積しやすいです：FreeFlow は以前このブログで議論した、ラグランジュ的な一貫性に基づく以前のデータフリー蒸留方法 BOOT と密接に関連しています。

Physics Informed Distillation (PID) は物理情報を加えたニューラルネットワーク（PINNs）からの灵感を引き出します：人々は長期間にわたって微分方程式の解オペレーターを学習するためにニューラルネットワークを使用してきました。それらの方法は、拡散モデルからの決定論的サンプリングのために使用される ODE のように、物理現象を描く ODE のようにも適用可能です：これは別のデータフリー蒸留バリエーションをもたらします。FreeFlow のように、微分は有限差分近似を使用して処理されますが、ここでは停止勾配操作は単に教師速度をラップしています：

$$ \mathcal{L}_\mathrm{PID} = \mathbb{E} \left[ \left( V(\mathbf{\varepsilon}, 1, t) - \mathrm{sg} \left[ v( F(\varepsilon, 1, t), t) \right] + (t - 1) \dfrac{\partial}{\partial t} V(\varepsilon, 1, t) \right)^2 \right] . $$

彼らは教師のバックプロパゲーションを回避することが本質的であることを述べており、それは学生が教師の弱点を利用することを可能にする（悪意ある例に類似する現象）。

オイラー的方法 🐔

オイラー的マッピング蒸留（EMD）は、事前トレーニング済み教師モデルからの速度見積もりを使用してオイラー的な一貫性から単純に導出されたロスを使用方法します：

$$ \mathcal{L}_{\mathrm{EMD}} = \mathbb{E} \left[ \left( \dfrac{\partial}{\partial s} F(\mathbf{x}s, s, t) + \nabla{\mathbf{x}_s} F(\mathbf{x}_s, s, t) v(\mathbf{x}_s, s) \right)^2 \right]. $$

LMD のように、自己蒸留バリエーションは速度をフローマップ自身の予測で置き換えることで構築でき、標準拡散ロスを組み合わせることができます：停止勾配操作を追加して、Jacobian の空間的部分 $\nabla_{\mathbf{x}_s} F(\mathbf{x}_s, s, t)$ をラップします。

Align Your Flow はまた、オイラー的蒸留方法（AYF-EMD）も特徴としています：ラグランジュバージョンのように、彼らは大きなステップとわずかに小さなステップを比較することから始めます：

$$ F(\mathbf{x}s, s, t) = F(\mathbf{x}{s + \Delta s}, s + \Delta s, t) , $$

ここで $\mathbf{x}_{s + \Delta s} = \mathbf{x}_s + \Delta s \cdot v(\mathbf{x}s, s)$。右辺は停止勾配操作でラップされ、二乗残差が損失として使用されます：$\Delta s \to 0$ を与えることで $\mathcal{L}\mathrm{EMD}$ を回復することを示し、再び停止勾配操作を除き、これはこの場合 Jacobian の空間的部分を通ってバックプロパゲーションを回避するのに役立ちます。最高の結果のために、彼らはこれを教師への自動ガイダンス、線形性レギュラリゼーションを持つウォームアップ訓練フェーズ、および EMD ロスを悪意あるロスと組み合わせた悪意あるファインチューニングフェーズと組み合せます。

Solution Flow Models (SoFlow) は非常に類似のレシピに従いますが、2 つの主要な違いがあります：

• 彼らはスケッチからフローマップを学ぶことに焦点を合わせ、自己蒸留をそのメカニズムとして使用します（ AYF は事前トレーニング済み拡散モデルからの蒸留に焦点を合わせています）；
• Jacobian ベクトル積は有限差分近似を通じて回避されます（$\Delta s$ は無限小ではなく有限だが小さい）：停止勾配操作の一方をラップ。

実用上で有限差分近似がうまく機能させるために、彼らは損失の重みを調整し、訓練の過程で $\Delta s$ を徐々に減らすためのカリキュラムを使用します。

Flow-anchored consistency models (FACM) は AYF の精神も似ていますが、訓練の安定性を改善するための興味深いトリック「フローアンカリング（flow anchoring）」を使用します：FACM のベースバージョンは $t=0$ だけ考慮します：ターゲット時間ステップは一貫性モデルのように固定です。その後、ソース時間ステップ $s$ の範囲を $[0, 1]$ から $[0, 2]$ に拡張します：$s > 1$ で、モデルはディノイザーとして動作することを期待します。これはフローマップモードとディノイザーモードを持つ単一のモデルをもたらし、これら 2 つのタスクでパラメータを共有します。それはモデルのパラメータを「アンカー」すると述べられています：補助的なディノイザータスクはフローマップ学習のためのレギュラーライザーとして機能します。

フローアンカリングパラメータ化は効率的な JVP 実装と組み合わされます：$t$ を変化することを許容するバージョンも考慮しており、完全なフローマップ学習を可能にします：興味深いことに、その設定ではモデルはディノイザーを 2 回学習します：1 回は補助的なディノイザータスク（$s > 1$）で、もう 1 回はフローマップタスク（$t = s$）で。明らかな冗長性にもかかわらず、補助的なタスクはまたこの場合に役立つようです。

以前のアプローチとは異なり、MeanFlow (MF) は（自己）蒸留に依存せず、標準拡散またはフローマッチングモデルと同じように周辺からの条件付き学習を使用します：これは前のセクションで既に説明されています（改善された MeanFlow バリエーションも含まれます）。MF の実用的な実装は $s$ および $t$ が互いに近づくにつれて不安定性を回避するために適応的な重みを伴います。また、訓練中に $s = t$ ケースがモデルを地面に保つために大幅に過サンプリリングされます。

MeanFlow の多くのバリエーションと拡張は探索されました：いくつかはこちらです：

• AlphaFlow はカリキュラム学習アプローチを提案し、瞬時の速度（フローマッチング）から平均速度（MF）への学習へ滑らかに補間します。
• Decoupled MeanFlow (DMF) はアーキテクチャの調整を提案します：ネットワークの初期レイヤーは単にソース時間ステップ $s$ 条件化し、後期レイAYER だけターゲット時間ステップ $t$ です：これは事前トレーニング済みディノイザーを MeanFlow モデルに適応させることを非常に簡単になります：単に初期および後期レイヤーのための時間エンベディングを非結合し、その後ファインチューンするだけです。彼らはまた、MF ロスの Cauchy バリエーションを使用して外れ値を押圧することを提案しています。
• Rectified MeanFlow は次の観察から始めます：データと雑音間のすべての経路が完全に直線の場合、瞬間速度および平均速度（任意の区間）はすべてで一致します！経路が曲がるほど少ないほど、ディノイザーを MeanFlow モデルに適応させることはより簡単になります。彼らは MF 訓練との単一の Reflow ステージを組み合せることを提案します。
• Pixel MeanFlow は、モデルがサンプルを生産するために非常に少数のステップを必要とするとき、潜在空間での生成モデリングの計算的利益が薄れることに注意します：その時点で、入力空間で直接操作する相対的単純さは好ましいかもしれないので、この設定に適応するための iMF を探求します。

合成的方法 🐶

ショートカットモデルは平均速度 $V$ の用語に基づくロス関数を使用し、自己蒸留で基礎付けされた合成的な一貫性ルールに基づいています：

$$ \mathcal{L}_\mathrm{shortcut} = \mathbb{E}\left[ \left( V(\mathbf{x}s, s, s + 2h) - \mathrm{sg} \left[ \hat{V}\mathrm{s + 2h} \right] \right)^2 \right] + \mathbb{E}\left[ \left( V(\mathbf{x}_t, t, t) - (\mathbf{\varepsilon} - \mathbf{x}_0) \right)^2 \right] , $$

ここで $\hat{V}\mathrm{s + 2h} = \frac{V(\mathbf{x}s, s, s + h) + V(\hat{\mathbf{x}}{s + h}, s + h, s + 2h)}{2}$ および $\hat{\mathbf{x}}{s + h} = \mathbf{x}_s + h V(\mathbf{x}_s, s, s + h)$ です：これは最初に少し複雑ですが、それは単に長さ $2h$ の時間区間全体の平均速度は長さ $h$ の 2 つの中間時間区間全体の平均速度の平均であるべきということです：このステップサイズを倍増することによってブートストラップする戦略はプログレッシブ蒸留と非常に類似しています。注記：損失全体では微分機能は何も特徴しません。

SplitMeanFlow は、前のセクションに入っているように聞こえるかもしれませんが、実際にはショートカットモデルの一般化であり、合成される時間区間が同じ長さである必要はありません：彼らは蒸留に焦点を合わせます。Boffi らはこれをプログレッシブ自己蒸留（PSD）として回復します。

Flow Map Matching (FMM) はやや異なるアプローチを取ります：合成性から思い出してください：フローマップは自身で逆関数であり、$F(F(\mathbf{x}_s, s, t), t, s) = \mathbf{x}_s$ です：$s$ に関する偏導関数を取り：

$$ \frac{\partial}{\partial s}F(F(\mathbf{x}_s, s, t), t, s) = v(\mathbf{x}_s, s) . $$

これら 2 つの等式は FMM ロスを構築するために組み合わされます：

$$ \mathcal{L}_\mathrm{FMM} = \mathbb{E} \left[ \left( \frac{\partial}{\partial s} F(F(\mathbf{x}_s, s, t), t, s) - (\mathbf{\varepsilon} - \mathbf{x}_0) \right)^2 \right] + \mathbb{E} \left[ \left( F(F(\mathbf{x}_s, s, t), t, s) - \mathbf{x}_s \right)^2 \right]. $$

第一項の周辺からの条件付き学習の使用に注意してください：スケッチからフローマップ訓練を可能にします。不幸に、この項も時間微分を導入しますが、それは $s$ に関する偏導関数であり、$F(\mathbf{x}_s, s, t)$ にバックプロパゲーションする必要はありません：彼らはこの方法が時間区間 $\mid t - s \mid$ が大きすぎないよう制限されたときに最もよく働くことを発見し、1 ステップでサンプリングすることを学ぶには適していません。

後者を解決するために、彼らはまたプログレッシブフローマップマッチング（PFMM）を提案します：

$$ \mathcal{L}\mathrm{PFMM} = \mathbb{E} \left[ \left( F(\mathbf{x}s, s, u) - F\mathrm{pre}(F\mathrm{pre}(\mathbf{x}_s, s, t), t, u) \right)^2 \right] , $$

ここで $F_\mathrm{pre}$ は限られた時間区間全体での事前トレーニング済みフローマップを表します：これは合成的な一貫性ルールを適用する最も純粋な方法かもしれませんが、既存の部分的フローマップが動作することを必要とします（FMM または別の方法を通じて獲得できます）。

一貫性モデルについてはどうですか？

2023 年にさかのぼる長期的な一貫性モデルの作業ラインがあります：以前の記事でいくつかについて書きました：元の Consistency Models の論文は連鎖反応を引き起こし、人々がノイズとデータの間の経路を特徴付ける方法の一つだけが速度を予測するに過ぎないことに realization しました。しかし「フローマップ」のフレーミングまで来なかったとしても、私はそのブログ記事のためにそれを使用することにしました：我々がすべてこの作業が互いに関連していることを理解するための役立つフレームワークを提供するためです。多くの最近の作業もそれを採用しています。

言及すると、いくつかのこれらの元の作業を現代のフレームワークに適合させるために時間を取る価値があります：Consistency Distillation (CD) はターゲット時間ステップ $t=0$（データ側）に錨定されたフローマップを生産します：

$$ \mathcal{L}_\mathrm{CD} = \mathbb{E} \left[ \left( F(\mathbf{x}s, s, 0) - \mathrm{sg} \left[ F(\hat{\mathbf{x}}{s - \Delta s}, s - \Delta s, 0) \right] \right)^2 \right] , $$

ここで $\hat{\mathbf{x}}_{s - \Delta s} = \mathbf{x}_s - \Delta s \cdot v(\mathbf{x}_s, s)$ は出力です。