worldbench/DiffusionOPSD
🔥 On-Policy Self-Distillation in Diffusion Models
📚 DiffusionOPSD とは何ですか？
DiffusionOPSD（オンポリシー自己蒸留）は、論文 「DiffusionOPSD: Diffusion モデルにおけるオンポリシー自己蒸留」（arXiv 2608.24646）に記載されたアルゴリズムを実装した研究用コードベースです。大規模な拡散画像生成モデル（例：Stable Diffusion 3.5‑M および Z‑Image‑Turbo）を微調整する方法を提供し、生成画像が PickScore、CLIPScore、HPS など、さまざまな学習済み画像品質または好み評価指標でより高いスコアを獲得できるようにします。
この手法の鍵となるアイデアは、報酬がノイズ除去の全過程終了後にのみ観測されるという通常の監視ギャップを埋めることです。DiffusionOPSD は繰り返し以下のステップを実行します：
- 現在の（固定された）「行動」ポリシーから低ノイズ状態を収集します。
- 各状態のクリーン出力アンカーの周囲で、いくつかの正規化された報酬勾配ステップを取ることで、明示的なポジティブおよびネガティブなターゲットを作成します。
- 限定的な最適化更新回数内で、これらの分離されたターゲットに一致する新しい「学習可能な」ポリシーを訓練します。
- 学習可能なポリシーの指数移動平均（EMA）で行動ポリシーを更新し、その後繰り返します。
ターゲット構築とモデル適合が分離されているため、DiffusionNFT や FlowGRPO などの先行手法と公平に検証・ベンチマーク・比較が可能です。
🚀 主な特徴
| 特徴 | なぜ重要か |
|---|---|
| オンポリシークエリ収集 | ターゲットは現在のモデルが実際に訪れた状態から構築されるため、オフラインまたは人工的にノイズを加えた状態を使用する場合に生じる不一致を回避できます。 |
| 明示的な報酬誘導ターゲット | ポジティブ／ネガティブなアンカーは報酬勾配に沿って制限付きステップを取ることで生成されるため、モデルが改善すべき明確な方向性が与えられます。 |
| 分離された有限適合 | ポリシー更新の前に報酬／デコーダーのグラフが分離されるため、ターゲットの品質はモデルの学習の程度とは独立して測定可能です。 |
| EMAベースの行動更新 | 各外ループの後、行動ポリシーは指数移動平均で更新されるため、訓練全体を通して新鮮な監視が保証されます。 |
| 2つのバックボーンをサポート | Stable Diffusion 3.5‑M（512²）および Z‑Image‑Turbo（1024²）の両方に対応しており、標準的および少数ステップの拡散レジームをカバーします。 |
| シングルおよびミックス報酬学習 | 論文で使用された7つの公開評価器のための事前定義済みプリセットに加え、任意の正の重み付き和（例：PickScore/26 + CLIPScore + HPSv2.1）を組み合わせる機能も備えています。 |
| 効率性の向上 | DiffusionNFT と比較して GPU 時間を 40 %–63 % 削減しながら、20の報酬マッチ設定のうち19でより高いホールドアウトスコアを達成しています。 |
| 豊富なスクリプト | 公開学習、ベースライン、スモークテスト、ミックス報酬実験用の準備済み起動スクリプトを提供。また、報酬モデルのセットアップを確認するためのユーティリティも用意されています。 |
🛠️ インストールとクイックスタート
前提条件 – CUDA 対応 GPU を搭載した Linux、Python 3.10‑3.11、および対応する PyTorch ビルド。
# リポジトリをクローン
git clone https://github.com/worldbench/DiffusionOPSD.git
cd DiffusionOPSD
# 新しい conda 環境を作成（オプションだが推奨）
conda create -n diffusionopsd python=3.11 -y
conda activate diffusionopsd
# コアパッケージをインストール（reward スタック拡張を含む）
pip install -e "[rewards]"
# ImageReward は古い timm バージョンをピン止めしているため、依存関係なしでインストール
pip install --no-deps 'image-reward==1.5'
# 公開プリセットで使用される報酬モデルのチェックポイントをダウンロード
export REWARD_CKPT_PATH="$PWD/reward_ckpts"
bash scripts/download_reward_weights.sh
# Pick‑a‑Pic プロンプトマニフェストを準備（論文の実験で使用）
python scripts/prepare_pickapic_prompts.py
Z‑Image‑Turbo で学習したい場合は、ZImagePipeline を含む Diffusers ビルドが必要です：
git clone https://github.com/huggingface/diffusers.git
pip install -e "./diffusers[torch]"
オプションの環境設定
export HF_HOME=/path/to/huggingface-cache # モデル重みのキャッシュ場所
export WANDB_MODE=offline # デフォルトでオンラインログを無効化
▶️ 小規模スモークテストの実行
このリポジトリには、1つの GPU ノード上で1回の最適化更新を実行する小さな「スモーク」スクリプトが同梱されており、報酬勾配パイプラインが自分のマシンで正しく動作するかを確認するのに役立ちます。
# SD3.5‑M バックボーン、HPSv2.1 報酬（高速チェック）
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh sd35 hpsv2
# Z‑Image‑Turbo バックボーン、CLIPScore 報酬
SMOKE_TEST=1 NPROC=8 UPDATES=1 bash scripts/train_public.sh zimage clipscore
また、HPSv3 や DeQA などの重い報酬もテストできます。これらは7〜8Bの大きなモデルに依存するため、別環境が必要です：
SMOKE_TEST=1 NPROC=7 UPDATES=1 bash scripts/train_public.sh zimage hpsv3
SMOKE_TEST=1 NPROC=7 UPDATES=1 bash scripts/train_public.sh zimage deqa
📊 本格的なトレーニング例
以下は、論文で使用された100更新の予算（約100回の最適化ステップ ≒ 数百回のロールアウト）で、公開報酬固有のモデルをトレーニングする典型的なコマンドです。
# PickScore/26 + CLIPScore + HPSv2.1 の混合目的で SD3.5‑M をトレーニング
NPROC=8 UPDATES=100 OUTPUT_DIR=outputs/sd35_mixed_opsd \
bash scripts/train_public.sh sd35 mixed
# HPSv3 評価器で Z‑Image‑Turbo をトレーニング（6ポリシー ランク + 1報酬サーバーランク）
NPROC=7 UPDATES=100 OUTPUT_DIR=outputs/zimage_hpsv3 \
bash scripts/train_public.sh zimage hpsv3
チェックポイントは10更新ごとに保存され（設定可能）、最終モデルは自動的に書き出されます。
📦 リポジトリの内容
| ディレクトリ／ファイル | 目的 |
|---|---|
scripts/ |
報酬のダウンロード、プロンプトの準備、スモークテストの実行、分散トレーニングの起動用のヘルパー スクリプト。 |
config/ |
各バックボーン、各報酬、および混合報酬設定用の YAML 形式の構成ファイル。 |
scripts/train_*.py |
分散実行のセットアップ、行動ポリシーと学習可能なポリシーの読み込み、軌道の収集、ターゲットの構築、最適化ステップの実行を行うエントリーポイント。 |
scripts/check_reward_setup.py |
長時間ジョブを実行する前に、必要な報酬モデルのチェックポイントとライブラリバージョンが存在するかを検証。 |
scripts/smoke_reward_gradient.py |
単一の報酬モデルを読み込み、画像空間の非ゼロ勾配が計算可能かどうかを確認。 |
assets/ |
README で使用される図（定性的ギャラリー、トレーニング曲線、アブレーション）。 |
LICENSE |
Apache 2.0 – フレキシブルなオープンソースライセンス。 |
📜 ライセンス
コードは Apache 2.0 ライセンス の下でリリースされており、自由な使用、改変、配布が許可されています（出典表記が必要）。モデル重み自体は上流の拡散モデルおよび報酬モデルのライセンスに準拠しており、リポジトリ内では再配布されていません。
🎯 どのような人に向いているか？
- 報酬誘導拡散微調整、自己蒸留、または生成モデル向けのオンポリシー強化学習スタイル手法を研究している 研究者。
- 特定の美学や好み指標に調整された高品質で制御可能な拡散モデルが必要で、再現可能なベースラインを求める 実務家。
- 微分可能な報酬モデル（大規模言語視覚評価器を含む）を拡散トレーニングパイプラインに統合する方法の具体的なエンドツーエンド例を探している 学生。
📚 補足リーディング
- 完全な論文（arXiv 2608.24646）では、数学的定式化、理論的動機、および広範なアブレーションが説明されています。
- プロジェクトページ（https://diffusionopsd.github.io/）には追加の視覚的結果と FAQ が掲載されています。
- 競合手法の背景については、引用された DiffusionNFT および FlowGRPO を参照してください。
TL;DR: DiffusionOPSD は、複数の学習済み報酬関数に基づいて拡散画像生成を改善する、新しいオンポリシー自己蒸留技術の、ドキュメントが整ったオープンソース実装です。先行手法と比較してトレーニングコストを削減しています。実行可能なスクリプト、2つの主要な拡散バックボーンのサポート、柔軟な報酬ミキシングシステムを備えています。
関連
- プロジェクト
- プロジェクト
- プロジェクト
- プロジェクト