Flow-GRPO の学習方法

Flow-GRPO: Training Flow Matching Models via Online RL

Flow-GRPO の強み

DDPO はユーザーの選好データを集めるコストが高い。Flow-GRPO は以下のような画像評価モデルを使用し、教師画像を使わずモデルのみで画質やプロンプトの追従性を向上させられる。ただし Flow-GRPO は非常に計算コストが高く、予算もしくは豊富な計算資源を持っている場合に使える。

評価モデル例

なぜ SDE を使うのか

Flow Matching は通常 ODE でサンプリングを行うが、ODE は決定的でサンプリングの際にノイズがないため、RL の報酬探索を行うと方策が過度に確定しやすくなる。SDE に変換すると同じマージナル分布を保ちながらノイズを導入でき、統計的サンプリングと方策勾配の更新が可能になるため、オンライン RL で効果的に学習できるようになる。

大まかなアルゴリズム

複数枚画像を生成
生成した画像を報酬モデルで評価
生成した画像の平均スコアを計算
平均以上の画像はプラスの学習、平均以下の画像はマイナスの学習を行う

なぜ複数枚画像を生成するのか

方策勾配の計算

学習の方向を決めるために、複数枚画像生成が絶対に必要になる。要は平均スコアがあれば、平均よりスコアの高い画像は良い方向に学習し、平均以下の画像はマイナスの方向に学習できる。

勾配の分散低減 (Variance Reduction)

生成画像枚数を増やせば分散が減少し、計算時間は増えるが学習が安定する。

並列計算

バッチ処理で効率的に計算できるうえ、各画像の計算は独立しているので複数の GPU で生成することもできる。

計算負荷

計算を高速化するために 10 step で画像を生成している。40 step だと ODE と品質が区別できず、5 step では品質が悪すぎてうまく学習できない。

論文では１回の学習に生成する画像数 G = 24（A800 24台使用）なので、１回学習するのに必要な推論回数は 240 回。SFT では１回の学習に必要な推論回数は１回だけだ。

論文では最大 5,000 step まで学習させている。

実装

Flow-GRPO は SDE サンプラーを実装する必要があり、KL 正則化項を計算するための mu と sigma とを計算する必要がある。

SDE 推論

実装は flow_grpo/diffusers_patch/sd3_sde_with_logprob.py を参照。

sde_step_with_logprob

sde_step_with_logprob は Flow Matching のモデル出力（速度ベクトル）を使って以下の処理を行う。

逆 SDE（SDE デノイズ）を１ステップ進める
サンプルを計算し対数尤度を取得
KL 正則化で使う prev_sample_mean と std_dev_t とを計算

std_dev_t

std_dev_tは以下の手順で計算される。

サンプリングするタイムステップ t を決める
t を使ってノイズスケジューラーから sigma を取得
sigma を使って std_dev_t を計算

prev_sample_mean

prev_sample_mean はモデルの出力にノイズを付与しただけ。

MixGRPO

MixGRPO: Unlocking Flow-based GRPO Efficiency with Mixed ODE-SDE

MixGRPO はスライディングウィンドウメカニズムを使用し、ノイズ除去ステップの一部のみを SDE でサンプリングすることで、訓練時間を50%近く短縮する。

実装は sd3_pipeline_with_logprob_fast.py を参照。

Pref-GRPO

Pref-GRPO: Pairwise Preference Reward-based GRPO for Stable Text-to-Image Reinforcement Learning

従来の GRPO はスコアの差が小さいと Reward Hacking が発生する。なぜならスコア差が小さい場合、スコア（Advantage）を正規化する際に差が実際以上にスケールされてしまうからだ。

そこでスコア差ではなく、勝率によって報酬を決定する Pairwise Preference Reward-based GRPO によって Reward Hacking を抑制し、安定した学習に成功した。