FewStepRL：少ステップ拡散モデルを強化学習で微調整する方法

拡散モデルのポストトレーニングでは、生成結果を報酬で評価し、その信号を使ってモデルを更新する強化学習が広く使われています。しかし、現在の方策を毎回ロールアウトする方式では、サンプリングステップが多いほどロールアウト自体が訓練コストの大部分を占めます。少ステップ拡散モデルはこの負担を減らせる一方、既存の強化学習目的をそのまま適用できるとは限りません。

問題の核心は速度の扱いです。多ステップの拡散モデルでは、輸送過程のある時点における瞬間速度を前提にすることが多いのに対し、少ステップモデルは有限区間にわたる平均速度を表現します。この違いを無視すると、目的関数が想定するモデルの挙動と、実際に更新するモデルの表現が一致しなくなります。

FewStepRLは、特定の新しい報酬関数を提案するのではなく、複数の目的を少ステップ設定へ移植する共通手順を示します。構成は次の三段階です。

• x空間への変換：各目的を共通の方策予測として書き直します。原則としてクリーンなサンプル側で表現し、遷移目的やフローマップ目的では対応する中間レベルを明示します。これにより、密度、速度、報酬を基礎にした目的を同じ枠組みで比較できます。
• 勾配の分解：目的の勾配を、参照先へ向かう自己回帰的な回帰項と、報酬による増分項に分けます。アルゴリズム間の違いは、勾配レベルでは主にアンカーと重みの選び方として整理できます。
• アンカーの適応：少ステップモデルの輸送終端、またはその接線方向からアンカーを計算します。瞬間速度が必要な場合は、モデル自身から復元するか、補助ネットワークを利用します。

この設計により、FlowGRPO、AWM、DiffusionNFT、DiffusionOPSD、ReFLの少ステップ版が一つのモジュール型コードベースに実装されました。研究者は手法ごとに訓練パイプラインを作り直すのではなく、アンカー、重み、速度推定の方法を組み合わせながら比較できます。

実験結果は、すべての条件で一つの手法が勝つわけではないことを示します。報酬が微分可能で、勾配計算のコストも許容できる場合は、報酬勾配に基づく手法が最も効果的でした。一方、報酬が微分不可能だったり、勾配を求める負担が大きかったりする場合、DiffusionNFT系は比較的少ない反復で高い報酬に到達する傾向を示しました。ただし、1回の反復コストが高くなる場合もあります。

学習時より推論ステップを増やす実験では、全体として一貫した改善は確認されませんでした。それでも、一部の出力は学習時の二倍から四倍程度のステップで目に見えて改善します。少ステップ化による高速化と品質向上は、単純なステップ数の増減だけで決まらないと考えられます。

FewStepRLの意義は、少ステップ拡散モデルの強化学習を個別の実装作業から、再利用可能な分析枠組みへと整理した点にあります。高速なサンプリングを活かすには、ステップ数を減らすだけでなく、速度の定義、報酬信号、方策予測の位置、回帰先となる参照値を一貫させる必要があります。

出典：Hugging Face Daily Papers