TRL を用いた DDPO による Stable Diffusion のファインチューニング

Hugging Face は DDPOTrainer を介して trl ライブラリに Denoising Diffusion Policy Optimization (DDPO) を統合し、強化学習 (RL) を使用して Stable Diffusion モデルをファインチューニングし、人間の好みと美的基準にさらに適合させることができます。

DDPO と報酬重み付き回帰の比較

DDPO は、計算オーバーヘッドを削減し、近似誤差を排除することで、拡散モデルに対する以前の RL 手法（例えば、報酬重み付き回帰 (RWR)）を改善します。RWR はデノイジング損失関数を再利用し、最終サンプルの報酬に基づいて重み付けします—これにより中間ステップが実質的に無視されます—一方、DDPO はデノイジングプロセス全体を多段階マルコフ決定プロセス (MDP) とみなします。

プロセスを MDP と捉え、固定サンプラーを使用することで、DDPO はエージェントポリシーを各方向同一のガウス分布とすることができます。これにより、最終サンプルの近似尤度に頼る代わりに、各デノイジングステップの正確な尤度を計算できるようになり、パフォーマンスが向上し、より複雑な目的にも対応できるようになります。

DDPO アルゴリズムと RLHF ワークフロー

DDPO は最適化問題に取り組むために、方策勾配法のうち特に近位方策最適化 (PPO) を使用します。DDPO 実装における主なカスタマイズは、PPO アルゴリズムの軌道収集部分です。

人間からのフィードバックによる強化学習 (RLHF) ワークフローに組み込むと、拡散モデルを人間の好みに合わせるプロセスは以下の 3 ステップに簡素化されます：

事前学習済みモデル: 事前学習済みの拡散モデル（例えば、Stable Diffusion）から開始します。
報酬モデル: 好みのデータを収集し、信号として機能する報酬モデルを訓練します。
DDPO ファインチューニング: 報酬モデルを信号として使用して、DDPO を用いてモデルをファインチューニングします。

画像の美観を向上させる場合、Hugging Face は Aesthetic Visual Analysis (AVA) データセットで訓練された訓練可能な回帰ヘッドを持つ凍結済み CLIP モデルを使用し、報酬指標として機能させました。

実装とトレーニングガイドライン

DDPO を用いて Stable Diffusion をトレーニングするには、相当なハードウェアリソースが必要です；NVIDIA A100 GPU がメモリ不足 (OOM) エラーを回避するための最小要件です。実装は、trl ライブラリの DDPOTrainer クラスと DDPOConfig クラスを介して行われます。

推奨ハイパーパラメータ

単一 GPU トレーニングの場合、以下のハイパーパラメータが推奨されます：

パラメータ	推奨値
`num_epochs`	200
`train_batch_size`	3
`sample_batch_size`	6
`gradient_accumulation_steps`	1
`sample_num_steps`	50
`sample_num_batches_per_epoch`	4
`per_prompt_stat_tracking`	`True`
`per_prompt_stat_tracking_buffer_size`	32
`mixed_precision`	`True`
`train_learning_rate`	3e-4

主な発見と学んだ教訓

Hugging Face の実験結果によると、DDPO チューニング済みモデルは、トレーニングプロンプトセットが最小限でも、幅広いプロンプトに対してよく一般化されます。

LoRA とフルファインチューニングの比較

LoRA: 推奨され、複数のテストで安定していることが証明されています。
フルファインチューニング (non-LoRA): LoRA よりも複雑な画像を生成できますが、安定させるのははるかに困難です。non-LoRA の実行では、学習率を低く設定すること（約 1e-5）および mixed_precision を None に設定することをお勧めします。

現在の制限事項

trl ライブラリにおける DDPOTrainer の現在の実装は、バニラ Stable Diffusion モデルのファインチューニングに限定されます。LoRA が主な焦点でありうまく機能しますが、フルトレーニングも可能ですが、より精密なハイパーパラメータ調整が必要です。