技術要約: AREX: フロー・マッチングにおける少ステップ・サンプリングのためのアフィン・残差指数積分器

問題提起

フロー・マッチング（FM）および関連するODEベースの生成モデルは、単純な参照分布（通常は等方的なガウス分布）を複雑なデータ分布へと輸送するために、学習された速度場を積分することに依存しています。効果的ではあるものの、高精度なサンプリングには多くの場合、多数のニューラルネットワーク評価（NFE）が必要となり、これが計算上のボトルネックとなります。既存のトレーニングフリーのサンプラー（DDIM、DPM-Solverファミリー、UniPCなど）は、ダイナミクスを線形部分（解析的に積分可能）と残差部分（数値的に積分可能）に分解することで、NFEを削減しようと試みています。しかし、これらの手法は通常、線形部分が単位行列のスカラ倍（等方的）であることを前提としています。この仮定は、現実世界のデータ分布における異方的な幾何構造（方向によって分散が大きく異なること）を捉えることができません。その結果、ニューラルネットワークは非ガウス的な残差構造と方向依存のスケーリング（異方性）の両方を補償しなければならず、これは低NFE領域において困難な課題となります。

手法: AREX

著者らは、ターゲット分布の第1および第2モーメント（平均と共分散）を明示的に利用して、異方性を解析的に扱うトレーニングフリーのサンプラーであるAREX（Affine-Residual Exponential Integrator）を提案します。

コアとなる分解

AREXは、学習された速度場 $v_\theta$ をブラックボックスとして扱うのではなく、以下の2つのコンポーネントに分解します：

1. モーメント一致アフィン成分 ($v_G$): ターゲットデータと同じ平均 ($m_1$) と共分散 ($\Sigma_1$) を持つガウス分布から導出された速度場。
   • 理論的正当性：著者らは、$v_G$ があらゆるタイムステップにおいて、真の母集団速度場に対する$L_2$最適アフィン近似であることを証明しています。
   • このコンポーネントは、ターゲットの固有値によって決定される方向依存の膨張または収縮率を捉えます。
   • 極めて重要な点として、このアフィン場によって生成されるフローは、閉形式の解（行列値プロパゲータ $F(t, s)$）を持ち、これは数値積分を行うことなく正確に計算可能です。
2. 残差コンポーネント ($r$): 学習された速度とアフィン成分の差 ($r = v_\theta - v_G$)。
   • この項は、非ガウス的な構造とニューラルネットワークの近似誤差を捉えます。
   • 理論的には、この残差は真のスコアと、モーメント一致ガウス分布のスコアとの差（非ガウススコア補正）に比例することが示されています。

積分戦略

AREXは、定数変化の公式を用いてダイナミクスを積分します：
$$x(t) = \underbrace{t m_1 + F(t, s)(x(s) - s m_1)}_{\text{正確なアフィン輸送}} + \underbrace{\int_s^t F(t, \tau) r(\tau, x(\tau)) d\tau}_{\text{数値的な残差補正}}$$

• 正確な伝播: アフィン部分は、ターゲット共分散 $\Sigma_1$ の固有値分解に依存する事前計算されたプロパゲータ $F(t, s)$ を用いて解析的に伝播されます。これにより、方向依存のスケーリングを数値的に近似する必要がなくなります。
• 数値的求積: 残差積分のみが数値的に近似されます。著者らは、残差の積分に対して**アダムス・バッシュフォース（AB）**スキーム（AREX-AB）を採用しています。
  • 効率性: この手法は、ステップごとにわずか1回の新しいネットワーク評価しか必要としません。過去のステップでの残差評価を再利用して、残差積分の高次多項式近似を構築します。
  • 残差プロファイル: トラジェクタリ全体にわたる残差の大きさの変化を扱うために、「残差マグニチュード・プロファイル」がオフラインで推定されます。このプロファイルは求積重みを調整し、追加のネットワークコールを行うことなく、系統的なスケール変化を考慮できるようにします。

条件付き生成 (AREX-C)

条件付きタスク（クラス条件付きやテキスト・トゥ・イメージなど）の場合、ターゲット分布は条件によって変化します。

• 共分散: AREX-Cは、アフィンバックボーンとして条件内平均共分散 ($\bar{S}$) を使用しており、これはすべての条件間で共有されます。
• 平均: 条件固有の平均は、トラジェクタリの最初のネットワーク評価を用いて、オフラインで適合させたキャリブレーション係数でスケーリングすることで動的に推定されます。
• このアプローチにより、各条件に対して個別の平均を保存することなく、アフィン分解の利点を維持したまま条件付き生成を実現します。

主な貢献

1. 手法: 学習された速度場を、解析的に解けるモーメント一致アフィン成分と、数値的に解ける残差に分離する、新しいトレーニングフリーの指数積分器であるAREXの導入。これはAREX-Cを通じて条件付き生成にも拡張されます。
2. 理論:
   • 最適性: モーメント一致アフィン場が、母集団速度場に対する$L_2$最適アフィン近似であることを証明しました。
   • 解釈: 残差項を非ガウススコア補正に関連付け、ターゲットがガウス分布であれば残差が消失することを示しました。
   • 誤差境界: プロパゲータの閉形式の公式を導出し、サンプラーの分布誤差（ワッサースタイン距離）の境界を提供しました。これにより、主要な誤差項が全トラジェクタリの曲率ではなく、残差の物質微分にのみ依存することを実証しました。
3. 結果: CIFAR-10、ImageNet-256、MJHQ-30Kを用いた実験的評価により、AREXは低NFE領域（画像で2–8 NFE、テキスト・トゥ・イメージで4–10 NFE）において、一貫して最高のFIDスコアを達成していることが示されました。

実験結果

論文では、3つのベンチマークにおいて、AREXを最先端のトレーニングフリー・サンプラー（Euler, DPM-Solver++, UniPC, CAB, STORK, DPM-Solver-v3など）と比較評価しています。

• CIFAR-10 (無条件): AREX-AB2は、NFE 2において39.24のFIDを達成し、次に優れたDPM-Solver-v3（112.04）を大幅に上回りました。NFE 4において、AREX-AB2（12.20）はDPM-Solver++（43.18）に対し約3.5倍の性能を示しました。
• ImageNet-256 (クラス条件付き): SiT-XL/2を使用した場合、AREX-C-AB3はNFE 2, 4, 8のすべてで最高のFIDを記録しました。NFE 4において、8.18に達し、DPM-Solver-v3（9.77）やCAB（12.37）を凌駕しました。
• MJHQ-30K (テキスト条件付き): SANA-0.6Bを使用した場合、AREX-C-AB2はNFE 4から10のすべての範囲で最高のパフォーマンスを発揮しました。NFE 4において、8.15のFIDを達成し、DPM-Solver-v3の10.56を上回りました。

定性的なサンプルからも、低NFEにおいてAREXがベースラインと比較して、より高い忠実度、優れたグローバル構造、および微細なセマンティック詳細を生み出すことが確認できます。

意義と主張

本論文は、AREXが少ステップ生成に対して原理的な視点を提供するものであると主張しています。一般的な蒸留（distillation）における目標である「フローを『真っ直ぐ』にする」ことに頼るのではなく、AREXは解析的に解くことが可能な、データの幾何学に内在する構造化されたダイナミクス（異方性）を活用します。

• 補完的なアプローチ: トレーニング時の効率化を目指す手法を、再学習を必要としない推論時の効率化という観点から補完します。
• 効率性: アフィンバックボーンを解析的に伝播させることで、数値ソルバーは方向依存のスケーリングを捉える負担から解放され、純粋に「真に困難な」残差ダイナミクスの捕捉に集中できるようになります。
• 汎用性: 本手法は、学習済みのあらゆるフロー・マッチング・モデルに適用可能であり、共分散推定のために学習データにアクセスする必要はありません（ただし、SANAのように自己サンプリングを利用することは有益です）。

著者らは、ターゲットのモーメントによって定義される扱いやすい構造と残差ダイナミクスを推論時に分離することで、限られたNFE予算の有用性を最大化し、効果的な少ステップ・サンプリングが可能になると結論付けています。