技術要約：境界条件による Rectified Flow の改善

問題提起

Rectified Flow (RF)、別名 Flow Matching は、常微分方程式 (ODE) を介してノイズ分布からターゲットデータ分布へと変換するための速度場を学習する、強力な生成モデリングフレームワークとして台頭しています。テキスト・トゥ・イメージやビデオ生成などのタスクにおけるその成功にもかかわらず、著者らは標準的な RF 実装における決定的な限界を指摘しています。それは、学習された速度場が理論的に導出された境界条件を満たさないことが多いという点です。

具体的には、ノイズ ($t=0$) とデータ ($t=1$) の間の線形補間において、最適な速度場は以下の条件を満たす必要があります：

1. 左境界 ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. 右境界 ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

標準的なニューラルネットワークによる速度場のパラメータ化は、これらの制約を本質的に強制しません。論文内の可視化が示すように、$t=1$ における予測速度は、期待される恒等写像から逸脱することがよくあります。この逸脱は、確率的サンプリング（例：ランジュバン動力学）において特に有害です。なぜなら、Tweedie の公式によるスコア関数の推定には $(1-t)$ による除算が含まれるため、$t=1$ 付近での不正確な速度はスコア推定を発散させるからです。これはサンプリングの不安定化を招き、特に大きなステップサイズやノイズスケールを使用する場合、画像の過度な平滑化や「カートゥーン風」のアーティファクトを引き起こします。

手法

これに対処するため、著者らは Boundary-enforced Rectified Flow Model (Boundary RF Model) を提案しています。このアプローチは、2 つの異なるバリアントを通じて、速度場のパラメータ化に境界制約を明示的に組み込みます。

1. マスクベースの Boundary RF Model

このバリアントは、速度場を境界項と学習可能なニューラルネットワークの出力の線形結合として構成することで、両方の境界条件 ($t=0$ および $t=1$) を強制します：
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
ここで：

• $m_\theta(\mathbf{x}, t)$ は標準的なニューラルネットワーク（例：U-Net または DiT）です。
• $C = \mathbb{E}[\mathbf{X}_1]$ です。
• $f(t), g(t), h(t)$ は、特定の境界で消失または活性化するように設計されたスカラー関数です（例：$g(1)=f(0)=h(0)=h(1)=0$ および $g(0)=f(1)=1$）。
• 著者らは「標準的なコサイン（Standard Cosine）」の選択肢として、$g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, $h(t) = \sin(\pi t)$ を提案しています。

この構成により、ニューラルネットワークの出力に関わらず、設計によって境界条件が満たされることが保証されます。

2. 減算ベースの Boundary RF Model

$t=1$ の条件のみを強制することが大きな経験的利益をもたらすという観察に基づき、著者らはより単純な単一境界バリアントを提案しています：
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
この定式化は、$t=1$ において項 $m_\theta(\mathbf{x}, 1)$ が相殺されるため、本質的に $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ を満たします。この手法は設計の簡潔さと柔軟性を提供しますが、すべての評価に対して $m_\theta(\mathbf{x}, 1)$ を計算するための追加のフォワードパスを必要とするため、計算コストが生じます。

確率的サンプリングへの影響

$\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ を強制することにより、Boundary RF Model は $t \to 1$ において Tweedie の公式の分子 ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) がゼロに近づくことを保証します。これにより、スコア関数 $\nabla \log \rho_t$ の発散を防ぎ、大きなノイズスケールやステップサイズを使用する場合でも、SDE サンプラー（確率的サンプラー）を安定させ、堅牢な生成を可能にします。

主な貢献

1. 限界の特定: 本論文は、制約のない速度場が Rectified Flow の理論的境界条件に違反し、不正確な ODE 軌道と不安定な確率的サンプリングを引き起こすことを明らかにしました。
2. 新しいパラメータ化: 最小限のコード修正でこれらの条件を強制できる、2 つの実用的なバリアント（マスクベースおよび減算ベース）を提供する Boundary RF Model を導入しました。
3. スコア関数の安定化: 境界の強制が $t=1$ 付近でのスコア関数の安定化に直接寄与し、バニラな RF を悩ませる発散問題を軽減することを実証しました。
4. 実証的検証: 異なるデータセットとモデルスケールにわたって、生成品質とサンプリングの安定性が一貫して向上していることを示す包括的な実験を行いました。

実験結果

著者らは、Diffusion Transformer (DiT) アーキテクチャを用いて、CIFAR-10 および ImageNet (256$\times$256 および 512$\times$512) でモデルを評価しました。

• 定量的な改善:

  • ImageNet 256$\times$256 において、Subtraction-based Boundary RF Model は、バニラ RF のそれぞれ 6.87 および 6.79 に対し、FID 6.32 (ODE サンプリング) および 6.38 (SDE サンプリング) を達成しました。
  • Mask-based モデルは、6.63 (ODE) および 6.18 (SDE) の FID を達成しました。
  • 両方のデータセットにおいて、Inception Score (IS)、Precision、Recall においても改善が見られました。
  • Subtraction-based バリアントは、一貫して最高の IS スコアを達成しました。

• サンプリングの堅牢性:

  • 様々なノイズスケール ($c$) を用いた Overshooting Sampler を使用した実験において、Boundary RF モデルは、確率性が増大してもバニラ RF と比較して、性能の低下が著しく少なく、大幅に安定したパフォーマンスを示しました。
  • 定性的な結果では、Boundary RF モデルは、特にバニラ RF が過度な平滑化を起こしやすい確率的サンプリングの設定において、アーティファクトが少なく、よりシャープな画像と優れたオブジェクト忠実度を生成することが示されました。

• スケーラビリティ:

  • Boundary RF Model の恩恵は、より大きなモデル (DiT-L, DiT-XL) や高解像度 (512$\times$512) にスケールアップした場合でも一貫していました。

意義と主張

本論文は、Boundary RF Model が Rectified Flow の根本的な限界に対する、シンプルで容易に実装可能な解決策を提供すると主張しています。境界条件を明示的に強制することで、この手法は以下を実現します：

1. 生成品質の向上: FID およびその他の指標において、一貫してバニラ RF を上回ります。
2. 安定した確率的サンプリングの実現: 終端時刻付近でのスコア関数の発散を解決し、より堅牢なサンプリング戦略を可能にします。
3. 適応の容易さ: 最小限のアーキテクチャ変更で、既存の RF モデル（U-Net や DiT を使用するもの）に統合できます。

著者らは、本研究を、特に高品質な確率的サンプリングを必要とするシナリオにおいて、Rectified Flow の実用的な展開に必要な洗練であると位置付けています。彼らは、焦点は RF にありますが、境界条件を強制するという核心的な概念は、微調整を加えることで拡散モデルにも拡張できる可能性があると示唆しています。今後の課題として、適応的な境界強制や、テキスト・トゥ・ビデオ生成への応用が挙げられています。