LeapAlignとは？2ステップ軌跡でFlow Matchingモデルを効率的に後訓練するアライメント手法

Flow Matchingモデルの長い生成軌跡を2ステップに圧縮することで、メモリ効率を保ちながら任意のステップから後訓練が可能
軌跡類似度スコアで整合性の高いサンプルに高い訓練重みを割り当て、勾配の安定性と学習品質を向上
FLUX.1-devを用いたHPSv2.1評価で0.4092を達成し、GRPOベース・直接勾配法の既存手法を全指標で上回る

研究の背景と課題

近年の画像生成モデルの主流は、拡散モデルからFlow Matching（フローマッチング、連続時間の最適輸送に基づく生成フレームワーク）へと移行しつつあります。Flux（Black Forest Labsが開発したFlow Matchingベースの高性能画像生成モデル）はその代表例であり、テキストから高品質な画像を生成する能力で注目を集めています。

しかし、こうした生成モデルを人間の好みや特定の要件に合わせて調整する「後訓練アライメント」には大きな課題がありました。報酬関数の勾配を生成プロセス全体に逆伝播させる直接勾配法（Direct Gradient Method）が提案されていますが、100ステップを超える長い生成軌跡を通じたバックプロパゲーションは膨大なGPUメモリを消費するうえ、勾配爆発（Gradient Explosion）のリスクを伴います。

また、GRPOなどの強化学習ベースのアプローチは報酬を最大化する方向に学習を進めますが、強化学習における多様性の維持やサンプル効率という別の問題が生じます。LeapAlignはこれらの課題を一挙に解決する、新しいアライメントフレームワークです。

LeapAlignの提案手法

LeapAlignの核心は、長い生成軌跡を「2ステップ軌跡」に圧縮するという発想です。通常のFlow Matchingでは、ノイズ状態x₁から最終画像x₀まで多数のステップを経て生成が進みます。LeapAlignはこの長い軌跡全体でバックプロパゲーションを行う代わりに、2つのタイムステップk（初期の高ノイズ側、1に近い値）とj（kより小さい中間ステップ）を[0,1]の全範囲からランダムに選択します。ランダム選択は固定距離での選択よりも性能が高く、特定のステップへの過学習を防ぐ効果があります。

具体的には、ステップkの潜在表現xₖから中間状態の予測x̂ⱼ|ₖを生成し、さらにそこから最終画像の予測x̂₀|ⱼを導出するという、2段階のリープ（飛躍）によって軌跡を圧縮します。こうして得られた短縮軌跡のみを通じて勾配を逆伝播するため、メモリ消費と計算コストを大幅に削減しながら任意の生成ステップを訓練対象にできます。

もう一つの重要な仕組みが軌跡類似度スコアによる重み付けです。圧縮した軌跡が本来の長い軌跡と大きく乖離している場合、その勾配信号は信頼性が低くなります。そこでLeapAlignは、予測状態と実際の軌跡状態との平均絶対差として dⱼ = mean(|xⱼ − x̂ⱼ|ₖ) および d₀ = mean(|x₀ − x̂₀|ⱼ) を算出し、重み付け因子 wsim = 1 / (max(dⱼ, τ) + max(d₀, τ))（τ=0.1）を計算します。類似度が高いほど訓練重みが大きくなり、信頼性の低いサンプルの影響を自然に抑制できます。

さらに、勾配安定化のための工夫として、値が大きい勾配項（ネストされた高次勾配）を完全に除去するのではなく、その大きさを低減する処理を加えています。既存の直接勾配法であるDRTuneはこれらの有用な勾配項を完全に削除しているため、学習効率が落ちる場面があります。LeapAlignはこの構造を保持することで、より豊富な学習シグナルを活用できます。

実験結果

FLUX.1-devを対象としたアライメント実験では、LeapAlignは主要な評価指標で既存手法を上回る性能を示しました。画像品質を測るHPSv2.1スコアでは、LeapAlignが0.4092を達成し、直接勾配法のDRTune（0.3882）・ReFL（0.3852）、そしてGRPOベースのMixGRPOを大きく上回っています。ベースモデルのスコア0.3078からの改善幅は約0.03となっており、後訓練の効果は顕著です。

テキストと画像の整合性を測るGenEvalスコアでも同様の傾向が見られ、LeapAlignは0.7420を記録し、MixGRPO（0.7232）やDRTune（0.7101）を上回りました。アブレーション実験では、タイムステップのランダム選択・軌跡類似度重み付け・勾配安定化のいずれもが性能向上に寄与していることが確認されています。特にタイムステップの固定距離選択をランダム選択に変更するだけで有意な改善が得られており、シンプルかつ実装しやすい設計であることが分かります。

まとめと今後の展望

LeapAlignは、Flow Matchingモデルの後訓練アライメントにおける計算コストと勾配安定性という2つの根本課題を、2ステップ軌跡という明快な設計で解決した研究です。CVPR 2026への採択を経て公開されており、Fluxをはじめとする最新の画像生成モデルの微調整に直接応用できる実用性を持ちます。

今後の課題としては、テキスト整合性のさらなる向上や、動画生成モデルなどFlow Matchingを採用する他のドメインへの展開が考えられます。また、タイムステップ選択の戦略を特定の生成ステップに特化した形で最適化することで、スタイル・構図・細部描写など特定の生成特性をターゲットにしたファインチューニングへの応用も期待されます。