Poly-OPD: Heterogeneous Multi-Teacher On-Policy Distillation for Capability-Selectable Flow Models

技術要約: Poly-OPD

問題提起
現在のオープンなテキスト・トゥ・イメージ（text-to-image）モデルは、しばしば相補的な強みを示します。例えば、一部のモデルは人間が好む審美性（FLUX.1-devなど）に優れている一方で、他のモデルは数え上げ、空間関係、属性結合といった構成的な指示への厳密な追従（Z-Imageなど）に長けています。しかし、これらの能力を単一のモデルに統合しようとすると、アーキテクチャの異質性が障壁となります。これらの特化型モデルは、異なるオートエンコーダ、デノイザー、およびノイズスケジュールのを利用しており、それゆえに潜在空間（latent space）が互いに理解不能な状態にあります。既存の蒸留手法（スコア、速度、または軌跡マッチングなど）は、共通の座標系を必要とするため、このようなケースには適用できません。さらに、教師モデルが生成した画像を用いた標準的なオフポリシー（off-policy）学習は、訓練と推論のミスマッチ（train-inference mismatch）を引き起こします。また、複数の能力を同時に学習させると、相反する勾配方向による破壊的な干渉が生じることがよくあります。

手法: Poly-OOD
著者らは、異種混合で潜在空間が互換性のない教師モデルを、単一のコンパクトなフローマッチング・スチューデント（具体的には2.5BのSD3.5-Mediumモデル）へと蒸留するためのフレームワークであるPoly-OPD（Poly-Objective On-Policy Distillation）を提案しています。このフレームワークは、以下の3つの核心的な課題に対処しています。

1. ピクセル・ブリッジによる異種混合オンポリシー蒸留:
   互換性のない潜在空間を橋渡しするために、Poly-OPDは「ピクセル・ブリッジ（pixel bridge）」を採用しています。スチューデントがサンプルを生成し、それをピクセル空間へとデコードします（ピクセル空間は、異種混合モデル間で唯一共有される座標系です）。この画像は、次にアクティブな教師モデルのVAEエンコーダによって再エンコードされます。教師は、タイムステップのインデックスではなく「大きさ（magnitude）」に基づいてノイズレベルが一致したブリッジ済みの潜在変数から、修正されたターゲット画像を生成します。教師による監督（supervision）は、教師の潜在空間ではなく、凍結されたDINOv2の特徴量空間（CLS埋め込み）において行われます。このセマンティックな特徴量空間は、モデル固有のピクセル統計量に依存しないため、スチューデントは互換性のある潜在変数を必要とすることなく、教師による修正のセマンティックな内容を学習できます。このプロセスが「オンポリシー」である理由は、ターゲットがスチューデント自身の生成サンプルから導出されるためであり、これにより訓練と推論のミスマッチを軽減しています。

2. 勾配の互換性を備えた選択可能なアダプタ:
   異なる能力（例：審美性と構成性）間の破壊的な干渉を防ぐために、Poly-OPDはモジュール式のアダプタ設計を利用しています。著者らは、どのパラメータを共有し、どのパラメータを隔離すべきかを経験的に決定するための勾配互換性診断（gradient compatibility diagnostic）を導入しました。測定の結果、アテンションの勾配はモード間で整合しているのに対し、フィードフォワードネットワーク（FFN）の勾配はしばしば衝突することが明らかになりました。その結果、Poly-OPDはすべての教師間で単一のアテンションLoRAを共有しつつ、特定の能力ごとにFFNアダプタを隔離しています。これにより、単一の凍結されたバックボーン上で複数の能力をホストでき、推論時にはアダプタの切り替えだけで対応が可能になります。

3. ギャップを考慮した適応的サンプリング:
   構成的なサブスキル（例：数え上げ、位置関係）が飽和する速度はそれぞれ異なることを認識し、本フレームワークは、生のプロンプトの難易度ではなく、残存する教師とスチューデントのギャップに基づいて訓練予算を割り当てるカリキュラムを採用しています。プローブセットを用いて、各カテゴリにおけるパフォーマンスのギャップを推定します。サンプリング分布は、スチューデントが教師に対して依然として大きく遅れをとっているカテゴリに集中するように動的に調整され、ギャップが縮まるにつれて一様サンプリングへと漸近していきます。

主な貢献

• 定式化: 本論文は、アーキテクチャが異なり潜在空間が互換性のない教師モデルからの、マルチ能力統合の問題を定式化しています。これは、既存の潜在空間蒸留が失敗する設定です。
• フレームワーク: Poly-OPDは、ピクセル・ブリッジによるオンポリシー監督、勾配測定に基づいたアダプタ共有戦略、およびギャップ認識型のサンプリング・カリキュラムを導入しています。
• 性能: 手法は、FLUX.1-dev (12B) と Z-Image (6B) の強みを、単一の 2.5B SD3.5-Medium スチューデントへと統合することに成功しました。

実験結果
標準的なベンチマークで評価した結果、Poly-OPDスチューデントは、それぞれの領域において両方のより大きな教師モデルを上回りました。

• 構成性 (GenEval): スチューデントは 73.3 を達成し、Z-Image (69.4) および FLUX.1-dev (65.2) を上回りました。特に、2オブジェクトの構成 (+14.93) や属性結合 (+8.25) といった構造化されたカテゴリで顕著な向上が見られました。
• 審美性 (DrawBench): スチューデントは HPSv3 スコア 11.35 および ImageReward 1.168 を達成し、SD3.5-Medium のベースモデルを上回り、12B の FLUX.1-dev 教師との差を大幅に縮めました。
• アブレーション研究: 選択可能なアダプタを取り除いた場合（共有FFNを使用した場合）、性能が劇的に低下（例：GenEval で -11.0 ポイント）し、相反する勾配を隔離する必要性が確認されました。同様に、ウォームスタート段階やギャップ認識サンプリングを削除すると結果が著しく低下し、提案された訓練パイプラインの妥当性が証明されました。

意義と主張
本論文は、Poly-OPDが、複数の大規模モデルをデプロイするメモリやレイテンシのコストをかけることなく、相補的な能力を単一の切り替え可能なモデルへと統合できることを主張しています。著者らが強調する重要な発見は、蒸留されたスチューデントが特定の指標（GenEval や ImageReward など）において自身の教師を凌駕することです。著者らは、この「凌駕」する能力を、訓練のオンポリシーな性質に特有のものと考えています。教師が生成したサンプルの分布に縛られるオフポリシー手法とは異なり、オンポリシー蒸留では、スチューデンドが教師による修正に基づいて自身の軌跡を洗練させることができるため、潜在的に教師の生の出力以上の品質を実現できる可能性があります。本フレームワークは、監督が共有されたセマンティックな特徴空間に固定されており、訓練が勾配認識型のアーキテクチャとカリキュラム設計を通じて管理されている限り、異種混合の蒸留が（共有の潜在空間なしでも）実現可能であることを示しています。