【DL輪読会】Flow-OPD: On-Policy Distillation for Flow Matching Models

DEEP LEARNING JP [DL Papers] On-Policy Distillation for Diffusion / Flow Matching Ku Onoda, Matsuo-Iwasawa Lab M2 http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 論文A：Flow-OPD 「On-Policy Distillation for Flow Matching Models」 • Z. Fang, W. Huang(Lead), Y. Zeng, Y. Zhao, S. Chen, K. Feng, Y. Lin, L. Chen, Z. Chen, S. Cao, F. Zhao • USTC / UCLA / CUHK / Xiaohongshu • arXiv:2605.08063 論文B：D-OPSD 「On-Policy Self-Distillation for Continuously Tuning Step-Distilled Diffusion Models」 • D. Jiang, X. Jin, D. Liu, Z. Wang, M. Zheng, R. Du, X. Yang, Q. Wu, Z. Li, P. Gao, Harry Yang, S. Hoi • HKUST / Alibaba（Z-Image Team）/ UCSD / CUHK • arXiv:2605.05204 どちらもOn-policy Distillation を Diffuison / Flow Modelで扱った研究

On-Policy Distillation とは？ 通常の蒸留 / SFT（off-policy） • 教師の軌道や正解画像（に乗せたノイズ）の状態で生徒を学習 • 推論時に生徒が実際に訪れるのは生徒自身が生成する状態 → 学習時と推論時の分布のずれが発生 On-policy Distillation • studentモデルがまず自分で生成（roll-out） • そのstudenモデルが訪れた状態の上で教師の出力に合わせる → 学習と推論の分布が一致 • ARモデルの場合

論文A｜概要：Flow-OPD • Flow Matching に初めて On-Policy Distillation を統合した post-training フレ ームワーク • 複数の「単一報酬でpost-traininしたexpert」を 1つの生徒モデルに統合 • やりたいこと： – マルチタスクの text-to-image の報酬モデルへのアラインメント（構図・文字描画・美的品質 …）を相互干渉なく1モデルに • アプローチ：スカラー報酬の代わりに教師からの密な軌道レベルの教師信号を使 う – 2段階：① 専門家教師を作る（GRPO）→ ② マルチ教師 OPD で生徒へ統合

論文A｜背景 • マルチタスク整合の「シーソー効果」 – 構図・文字・美しさ…目的が衝突し、1指標を上げると別指標が落ちる • GRPO の限界（スカラー報酬） – 報酬希薄性：多次元の対立を1次元スカラーに圧縮 – 勾配干渉：T₁ を最適化する更新が監視外の Tₖ を 損なう • 予備実験 – GenEvalを最適化したモデルに OCR 報酬を足すと GenEval が 5%低下 – 報酬を素朴に混ぜると新報酬ごとに既存指標が 3〜9%低下 ＝ catastrophic forgetting • 着想 – LLMでは は OPD（専門家からの蒸留）で解決済み →「Flow Matching でも同じことができないか？」

論文A｜手法① 全体像（2段階） Stage 1：Cold Start • 各タスクの専門家教師を単一報酬 GRPO で作る（各々単独で性能上限まで） • 生徒の初期化は2通り：教師軌道で SFT / 全教師の Model Merging Stage 2：マルチ教師 OPD • 生徒が on-policy サンプリング • task routing で条件に応じ専門家を選ぶ • その状態で velocity を一致

論文A｜手法②：reverse-KL → velocity の L2 LLM の OPD（reverse-KL）を連続フローへ移すことを考える ODE を SDE 化（各 denoising ステップを確率的な遷移として扱う） • 生徒と教師の遷移は 共分散が共通（SDE のノイズスケジュールで決まる） • 共分散が等しい2ガウスの KL は L2 距離だけ に縮約 ⟹ velocity（速度場）の重み付き L2 距離に帰着 • LLM のような高分散な policy gradient が不要（LLMでは勾配をモンテカルロ近似する必要あり） policy gradient の期待値が 勾配にちょうど一致 → 分散ゼロ の閉形式で最適化できる。

論文A｜手法③ Routing と MAR Task-Specific Routing • テキスト条件 → 対応するexpertにハードルーティング • 目標 velocity は マッチした専門家のみ から → 目的どうしの干渉を構造的に分離 Manifold Anchor Regularization (MAR) • 純粋な RL 最適化による 美的劣化を防止 • 凍結した美的最適化した教師への KL ペナルテ ィで品質多様体を保つ

論文A｜結果①

論文A｜結果② 汎化 ・OOD評価 最適化していないベンチマークにおいても高い性能を発揮 通常のGRPOのみでは性能劣化（ catastrophic forgetting ）

論文A｜結果② Ablation • Cold start ablation Merge > SFT > w/o cold start • MAR（美的experのKL） 品質を保った状態でGenevalタスクに最適化 11

論文B｜概要：D-OPSD • few-step 蒸留済みモデル（例：Z-Image-Turbo, FLUX.2-klein）を 少ステップ 推論能力を壊さずに 継続 fine-tune する on-policy 自己蒸留（外部報酬は不要） • やりたいこと：4〜8 step の高速モデルを、手元の画像–テキストペアだけで追加 学習（概念/スタイル付与） • 着想： 同一モデルを 生徒（テキスト条件） 教師（テキスト＋目標画像のマルチモーダ ル条件） の2役に （LLMのOPSDで教師モデルにのみ答えを与えることと同じ着想）

論文B｜背景・課題 • few-step 蒸留モデルの弱点 – 4〜8 step で動くよう蒸留された dynamics – 普通に vanilla SFT すると… 正解画像にノイズを乗せた状態で学習＝生徒が通らないパス → 分布ずれで few-step 能力が崩壊 • 既存手法の不足 / 求められる設計 – Online RL：few-step は保てるが 報酬関数の設計が必要 – 求められる設計 ① 生徒自身の roll-out 状態で更新し ② その同じ状態にペアデータの教師信号を載せる

論文B｜encoderのin-contex能力 in-context 能力 最近の拡散モデル（LLM/VLM エンコーダ）は、テキストだけの特徴を テキスト ＋目標画像のマルチモーダル特徴 に差し替えるだけで、追加学習なしでも その概 念/スタイルを保った生成ができる

論文B｜手法① 同じモデルに異なる条件付け • 生徒：テキストのみ条件 • 教師：マルチモーダル条件 c_s = f_text(y) c_t = f_mm(y, x₀)（プロンプト＋目標画像） 15

論文B｜手法② 損失と「壊れない理由」 生徒が few-step 軌道を生成し、その訪れた各状態で生徒・教師が velocity を予測。生徒を教師（ stop-grad）に合わせる few-step を保てる理由 • vanilla SFT と違い、生徒が決して通らない「正解画像の 状態」を無理に当てさせない ・学習は常に 生徒の roll-out 状態の上 で行われ、分布ず れが大幅に減る 報酬が要らない理由 • 教師信号は 目標画像を条件に入れた同一モデルの出力 そ のもの、外部の reward model を設計せず自己蒸留で成立

論文B｜結果① LoRA（小データ） • 少数のデータ（4枚のコンセプトサンプル）から新しい概念を獲得し、汎化能力が向上

論文B｜結果② Full-FT • 新しい概念を獲得しながら、すでに獲得している能力は忘却しない

論文B｜結果③ Ablation • SFTだと、段階的にfew-step能力が破綻するのに対 し、self distillatioでこの問題を軽減する • freezeしたbase modelをteacheにすることで安定し て学習可能 • EMAでは大きなモメンタム係数が必要（0.9999の設 定が最良だった） →なるべく教師モデルの信号は変えない方がいい 19

まとめ 2本の比較 観点 論文A：Flow-OPD 論文B：D-OPSD 目的 マルチタスク整合（干渉の解消） few-step 能力を保った継続学習 教師 複数の報酬専門家（GRPO） 同一モデル（マルチモーダル条件） 監督 報酬ベース ＋ OPD 報酬フリーの自己蒸留 数式のキモ reverse-KL → velocity L2（閉形式・分散0） roll-out 状態での velocity L2 付加要素 task routing / MAR in-context 条件付け Flow Matching（SD-3.5-M） step-distilled（Z-Image-Turbo 等） 対象 共通点 • どちらも「生徒の roll-out 状態の上で velocity を教師に合わせる」＝ off-policy / vanilla SFT の train–test mismatch を回避 する 連続モデル版 On-Policy Distillation。

課題と今後の方向性 • 課題 – Flow-OPD • 教師モデルの数の比例してだけ訓練コストが増大 • 教師モデルの生成能力に依存 – D-OPSD • 通常のSFTに比べて、オンポリシーのロールアウトコストが高い – （FLOPs4倍, wall clock2倍） • In-context能力に依存 • 今後の方向性 – 動画 / 3D / 編集 など強力な条件信号の利用 21