Exploring the Design Space of Reward Backpropagation for Flow Matching

技術要約：Flow Matchingにおける報酬バックプロパゲーションの設計空間の探索

1. 問題提起

学習済みのテキスト・トゥ・イメージ（text-to-image）フローマッチングモデルは高品質な画像を生成するが、人間の好みに沿う（アライメント）という点ではしばえることが多い。方策勾配法（policy-gradient methods）も存在するが、**直接的な報酬バックプロパゲーション（direct reward backpropagation）**は、フローサンプラーの微分可能性を利用してエンド・ツー・エンドで報酬勾配を伝播させることで、よりサンプル効率の高い代替案を提供する。しかし、このアプローチは、現代のモデルスケールにおける2つのよく知られた病理によって阻害されている：

1. 活性化メモリ（Activation Memory）: フルサンプリング・トラジェトリ全体にわたって活性化を保存することは不可能である。
2. 勾配爆発（Gradient Explosion）: 複数のステップにわたるヤコビアン積の連鎖により、報酬勾配が逆方向に伝わる際に増幅され、初期タイムステップの更新を不安定にする。

既存の解決策は、フルバックワード・トラジェトリを短い「固定された（pinned）」パス（例：straight-throughコネクターの使用）に置き換えることで、これらの問題を軽減しようと試みている。しかし、これらの手法は、バックワードパスがいかにフォワード・ロールアウトを近似するかに関して、特定の孤立した設計上の選択に依存していることが多い。これにより、特に長い間隔において、グラディエントの品質が近似の精度に強く依存するというトレードオフが生じ、基礎となる設計空間が未探索のまま残されている。

2. 手法：FlowBPフレームワーク

著者らは、バックワード・トラジェトリ自体を主要な設計対象として扱う統一されたサロゲート・トラジェトリ・フレームワークであるFlowBPを提案する。FlowBPは、勾配計算を行わないキャッシュされたロールアウトに基づいた、軽量なサロゲート・トラジェトリを構築することで、$N$ステップのオイラー・ロールアウト全体に対してバックプロパゲーションを行う代わりに、これを用いる。

このフレームワークは、4つの独立した設計軸によって定義される：

1. 報酬モデルの入力（Reward-Model Input）: 報酬が、実際のキャッシュされたロールアウト・サンプル（$x_0$）に対して評価されるか、あるいは事後平均推定値（$\hat{x}_0$）に対して評価されるかを決定する。
2. アクティブ・セット（Active Set）: どの速度評価（velocity evaluations）が勾配と共に再フォワードされるかを指定し、それ以外の部分は定数として扱う。
3. 積分重み（Integration Weights）: アクティブな速度がどのように結合されるかを定義する（例：オイラー重み vs ラグランジュ求積法）。
4. ブリッジ・カップリング（Bridge Coupling）: 分割インデックス $j$ を用いて、前後セグメントを接続するかどうか、およびこの分割における入れ子状の勾配パスの強度（$\alpha$）を制御する。

主要なメカニズム:

• メモリ効率: メモリ使用量は、ロールアウトの長さ（$N$）ではなく、アクティブ・セットのサイズ（$|A|$）に比例する。
• 勾配の安定性: 勾配伝播は、最大でも単一のヤコビアン因子（ブリッジ・カップリング経由）に制限され、多ステップの連鎖による指数関数的な増幅を回避する。
• サロゲート構築: フレームワークは、2種類のサロゲート推定をサポートする：
  • リープ推定（Leap Estimate）: コネクターを使用して、予測された状態をキャッシュされた状態に固定する（残留成分を導入する）。
  • 再構成推定（Reconstruction Estimate）: キャッシュされた状態と再フォワードされたアクティブな速度を用いて離散的なオイラー間隔を再現し、フォワードの値がキャッシュされたロールアウトと正確に一致することを保証する。

3. 主要な貢献

本論文は、FlowBPフレームワークの3つの具体的なインスタンスを紹介し、それぞれ設計空間の異なる領域を探索している。

1. FlowBP-Sparse:

   • 戦略: フル・ロールアウトに対するスパースなオイラー再構成を使用し、ブリッジは持たない（$j=0$）。
   • メカニック: キャッシュされたオイラー更新を再生し、オートグラッド（autograd）にさらされるのは中程度のサイズのアクティブ・セットのみとする。ブリッジを完全に排除することで、クロスステップのヤコビアンを持たない完全にデカップルされた勾配を実現する。
   • 利点: 無条件の安定性を保証し、報酬を真のイン分布・サンプル上で評価しながら、メモリをアクティブ・セットのサイズに抑える。

2. FlowBP-Bridge:

   • 戦略: フル・オイラー求積トラジェトリ再構成と、構造的に正確なブリッジを組み合わせる。
   • メカニック: ロールアウトを、調整可能な結合強度（$\alpha$）を持つブリッジによって結合された前・後セグメントに分割する。ブリッジは、バイアスのある単一速度コネクターではなく、オイラー合成によって形成される。
   • 利点: FlowBP-Sparseのデカップルされた安定性と、ブリッジされた手法のクロスステップ・クレジット割り当ての間を補完し、制御された入れ子状の結合を可能にする。

3. FlowBP-Lagrange:

   • 戦略: リープ求積の次数を高めることで、コネクターベースのアプローチ（LeapAlignに類似）を強化する。
   • メカニック: 単一速度のホップを、2セグメント構造内での高次ラグランジュ求積則に置き換える。
   • 利点: 長い間隔における積分精度を向上させることで、コネクターの残留成分（$d_0$ および $d_j$）を数値的に減少させ、従来のコネクターベースの手法で見られる「エンドポイント・残留失敗モード」を軽減する。

4. 実験結果

著者らは、3つのフローマッチング・バックボーン（SD3.5-M, FLUX.1-dev, FLUX.2-Klein-base）を用いて、3つのバリアントを評価した。

• パフォーマンス: 嗜好性（HPSv2.1, PickScore, ImageReward）、品質（UR-Align, UR-IQ）、および構成性（GenEval）の指標において、FlowBPの各バリアントは、直接勾配ベースライン（ReFL, DRaFT-LV, DRTune, LeapAlign）を一貫して上回った。
• 変動性: すべての指標やモデルにおいて支配的な単一のバリアントは存在せず、最適な選択はバックボーンや指標によって回転した。例えば、SD3.5-MにおいてはFlowBP-LagrangeがHPSv2.1で優れており、FLUX.1-devにおいてはFlowBP-BridgeがPickScoreでリードした。
• 構成的生成: FLUX.1-devにおいて、FlowBP-LagrangeはGenEvalの総合スコアを63.25から69.88に向上させ、嗜好性へのアライメントがプロンプト追従性を低下させないことを示した。
• アブレーションによる洞察:
  • エンドポイント再構成: 再構成（via reconstruction）を通じて、実際のロールアウト・サンプル上で報酬を評価することは、ReFLのような既存の手法に対してもパフォーマンスを向上させた。
  • コネクターの精度: 高次のコネクター（Lagrange）は、単一速度のリープと比較してエンドポイント予測誤差を大幅に減少させ、大きな残留成分による学習の不安定性を防いだ。
  • ブリッジ・カップリング: 適度な入れ子状の結合（$\alpha$）は、勾配の不安定性を再導入することなく、有用な信号伝播を提供した。

5. 意義と主張

本論文は、FlowBPが、以前は未探索であった設計構成を露出させつつ、先行する直接勾配手法を特殊なケースとして包含する統一的な視点を提供すると主張している。その主な意義は以下の通りである：

• サンプリングと最適化の分離: バックワード・トラジェトリを設計オブジェクトとして扱うことで、サンプリングプロセス（キャッシュされたロールアウト）と最適化グラフ（サロゲート）を分離する。
• スケーラビリティ: 設計レベルでメモリと勾配爆発を制限することで、大規模モデルに対する直接的な報酬バックプロパゲーションを可能にする。
• 実行可能な設計: 4つの設計軸（入力、アクティブ・セット、重み、結合）は、勾配の忠実度、メモリ使用量、および学習の安定性の間のトレードオフを制御するための実用的なノブを提供する。

著者らは、明示的なサロゲート・トラジェトリの設計が、フローベースの生成モデルにおける報酬アライメントをスケールアップするための、孤立したヒューリスティックを超えた、実践的かつ解釈可能なルートであると結論付けている。