技術的サマリー：LeapAlign

問題定義

本論文は、フローマッチングモデルをポストトレーニングを通じて人間の嗜好に整合させるという課題に取り組んでいる。直接勾配法は、フローマッチングサンプリングの微分可能性を活用することで、方策勾配アプローチ（GRPO など）に対する有望な代替手段を提供するが、長い生成軌道に適用される際には重大な障壁に直面する。具体的には、ODE サンプリングステップの全シーケンスを通じて報酬勾配を逆伝播させることは、莫大なメモリコストを招き、勾配爆発のリスクを伴う。その結果、既存の直接勾配法は通常、最終画像に近い単一のタイムステップへの更新に制限するか、不安定さを回避するためにネストされた勾配を破棄している。この制限により、最終画像のグローバルなレイアウトと構造的構成を決定する上で重要な初期生成ステップの最適化が妨げられている。

手法：LeapAlign

著者らは、過度なメモリオーバーヘッドを招くことなく、初期タイムステップを含む任意の生成ステップへ効率的かつ安定した直接勾配伝播を可能にするファインチューニング手法LeapAlignを提案する。中核的な手法は以下の 3 つの主要コンポーネントから構成される。

1. 2 段階リープ軌道構築

完全な生成軌道全体を通じて逆伝播する代わりに、LeapAlign は標準的なフルラン軌道から導出された 2 段階のみで構成される短縮されたリープ軌道を構築する。

• プロセス: 各イテレーションにおいて、ノイズ ($x_1$) からクリーンな画像 ($x_0$) までの完全な軌道をサンプリングする。2 つのタイムステップ $k$ と $j$（ただし $k > j$）をランダムに選択する。
• リープ: 整流フローモデルの 1 段階リープ予測特性を利用して、未来の潜在変数を 1 段階で推定する。
  • リープ 1: $x_k$ における速度場を用いて $x_k$ から潜在変数 $x_j$ を予測する：$\hat{x}_{j|k} = x_k - (k-j)v_\theta(x_k)$。
  • リープ 2: $x_j$ における速度場を用いて $x_j$ から最終潜在変数 $x_0$ を予測する：$\hat{x}_{0|j} = x_j - j v_\theta(x_j)$。
• 潜在変数コネクタ: 予測された状態をフル軌道からサンプリングされた実際の状態と整合させつつ微分可能性を維持するため、「潜在変数コネクタ」が stop_gradient 操作を用いて導入される（例：$x_j = \hat{x}_{j|k} + \text{stop\_gradient}(x_j - \hat{x}_{j|k})$）。
• 利点: これにより逆伝播パスが 2 段階に短縮され、元の軌道の長さに関わらずメモリコストを一定に保ちながら、$(k, j)$ のランダムな選択を通じて生成プロセスの任意の部分を更新することを可能にする。

2. 勾配割引

有用な学習シグナルを保持しつつ勾配爆発に対処するため、LeapAlign は勾配割引メカニズムを導入する。

• 課題: 2 段階を通じて逆伝播することは、タイムステップ間の相互作用を捉える「ネストされた勾配」項（$\frac{\partial v_\theta(x_j)}{\partial x_j} \frac{\partial v_\theta(x_k)}{\partial \theta}$）を導入する。DRTune などの従来手法は、モデル入力において勾配を停止させることでこの項を完全に排除するが、これにより貴重な依存関係情報が失われる。
• 解決策: LeapAlign は、ネストされた勾配項を完全に除去するのではなく、係数 $\alpha \in [0, 1]$ によってスケーリングする。これは、2 番目のリープにおける速度予測への入力を変更することで達成される：$\hat{x}_{0|j} = x_j - j v_\theta(\alpha x_j + (1-\alpha)\text{stop\_gradient}(x_j))$。
• 効果: これにより大きな勾配項の大きさを低減してトレーニングを安定させつつ、ネストされた勾配を含む完全な勾配構造を保持する。

3. 軌道類似度重み付け

短縮されたリープ軌道が信頼性の高い監督を提供することを保証するため、本手法は真の生成経路に密接に類似するリープに高い重みを割り当てる。

• メカニズム: 損失は、接続点（$x_j$ および $x_0$）における予測状態（$\hat{x}$）と実際の状態（$x$）との間の平均絶対差の逆数によって重み付けされる。
• 数式: $w_{sim} = \frac{1}{\max(d_j, \tau) + \max(d_0, \tau)}$。ここで $d$ は距離を表し、$\tau$ はほぼ同一のペアへの過度な強調を防ぐための最小閾値である。
• 目的: これにより、真のダイナミクスから大きく逸脱するリープが勾配更新に与える影響を低減し、監督の安定性と忠実性を向上させる。

4. ファインチューニング目的関数

本手法は、報酬ハッキングを防ぎつつ報酬を最大化するためのヒンジ型損失を採用する：$L_{raw} = \max(0, \lambda - r(x_0))$。重要なのは、リープ予測から報酬を推定する手法とは異なり、LeapAlign はフル軌道によって生成された実際の最終画像（$x_0$）に対して報酬を評価することであり、これにより報酬モデルが真の視覚的および意味的品質を評価することを保証している。

主な貢献

1. LeapAlign フレームワーク: 2 段階リープ軌道を構築し、一定のメモリコストで初期生成ステップへの直接勾配逆伝播を可能にする新規ポストトレーニング手法。
2. 勾配安定化技術: ネストされた勾配を除去するのではなくスケーリングする勾配割引の導入と、一貫したリープを優先する軌道類似度重み付けの導入。これにより、従来の直接勾配法における不安定性とシグナル損失の問題に対処した。
3. 包括的な評価: 人間の嗜好スコア（HPSv2.1、HPSv3、PickScore）および構成的整合ベンチマーク（GenEval）を含む複数の指標において、最先端の方策勾配（GRPO ベース）および直接勾配法を上回る LeapAlign の性能を実証。

結果

著者らはFlux.1-Devモデルを LeapAlign でファインチューニングし、DanceGRPO、MixGRPO、ReFL、DRaFT-LV、DRTune などのベースラインと比較評価を行った。

• 一般的な嗜好整合: LeapAlign は、ドメイン内（HPSv2.1）およびドメイン外（HPSv3、PickScore、UnifiedReward）の両方の指標で最高スコアを達成した。特に、3 つの報酬モデルを使用する MixGRPO を上回りながら、単一の報酬モデル（HPSv2.1）のみでトレーニングされた点で注目すべきである。
• 構成的整合: GenEval ベンチマークにおいて、LeapAlign は「2 つの物体」、「色」、「位置」、「属性バインディング」タスクで顕著な向上を示した。これらのカテゴリでの改善は、グローバルなレイアウトを決定する初期生成ステップの最適化における本手法の成功を浮き彫りにしている。
• 定性的分析: 視覚的比較により、LeapAlign はテキストプロンプトとの整合を改善するために画像のグローバルな構造を大幅に変更するのに対し、後期ステップ更新に制限された手法（ReFL など）は事前学習モデルのレイアウトを保持することが多いことが示された。
• アブレーション研究: 実験により、2 段階リープ設計が性能とメモリの面で最適なトレードオフを提供することが確認された。$\alpha=0.3$ での勾配割引が最適な結果をもたらし、$[0, 1]$ の全範囲にわたるタイムステップのランダム選択は固定範囲よりも優れており、初期ステップ更新の重要性を浮き彫りにした。

意義と主張

本論文は、LeapAlign がフローマッチングポストトレーニングにおけるメモリ効率と勾配カバレッジの間のトレードオフを成功裡に解決したと主張している。過大な計算コストを伴うことなく初期生成ステップの更新を可能にすることで、複雑な構成的タスクに不可欠なグローバルな画像構造の最適化を可能にする。著者らは、LeapAlign を確率的方策勾配法および不安定な直接勾配アプローチの両方に対する安定かつ効果的な代替手段として位置づけており、優れた画像 - テキスト整合と生成品質を達成している。この研究は、ネストされた勾配を破棄するのではなく保持してスケーリングすることが、直接勾配整合の性能向上の鍵となる要因であることを示唆している。今後の課題として、これらの原則を動画生成へ拡張することが挙げられている。