[論文レビュー] Directly Fine-Tuning Diffusion Models on Differentiable Rewards

Kevin Clark, Paul Vicol+2 more|arXiv (Cornell University)|Sep 29, 2023

Generative Adversarial Networks and Image Synthesis被引用数 7

ひとこと要約

本稿では、微分可能報酬関数を用いて、全拡散サンプリングプロセスを直接バックプロパゲーションすることで、拡散モデルを微調整するDirect Reward Fine-Tuning (DRaFT) を提案する。勾配チェックポイントとLoRAを活用することで、強化学習ベースのベースラインと比較して200倍以上高速な性能を達成し、DRaFT-$K$ やDRaFT-LVといった変種により、訓練の安定性と分散低減が向上する。

ABSTRACT

We present Direct Reward Fine-Tuning (DRaFT), a simple and effective method for fine-tuning diffusion models to maximize differentiable reward functions, such as scores from human preference models. We first show that it is possible to backpropagate the reward function gradient through the full sampling procedure, and that doing so achieves strong performance on a variety of rewards, outperforming reinforcement learning-based approaches. We then propose more efficient variants of DRaFT: DRaFT-K, which truncates backpropagation to only the last K steps of sampling, and DRaFT-LV, which obtains lower-variance gradient estimates for the case when K=1. We show that our methods work well for a variety of reward functions and can be used to substantially improve the aesthetic quality of images generated by Stable Diffusion 1.4. Finally, we draw connections between our approach and prior work, providing a unifying perspective on the design space of gradient-based fine-tuning algorithms.

研究の動機と目的

• 微分可能報酬関数を用いて、モデルパラメータを勾配最適化することで、事前学習済みの拡散モデルと望ましい行動（美しさや人間の好みなど）との乖離を是正する。
• 拡散モデルの整合性に向けた強化学習ベースの微調整手法の非効率さとサンプルの複雑さを克服する。
• 微分可能報酬関数から得られる解析的勾配を効率的に活用し、ブラックボックス報酬最適化の必要性を回避する手法を開発する。
• 画像品質、オブジェクト検出、敵対的例生成など、多様な報酬関数を効率的かつスケーラブルに扱える、組み合わせ可能な微調整を実現する。
• 複数の目的を統合的にサポートする、一元化された勾配ベースのフレームワークを提供し、LoRA重みの補間と合成を可能にする。

提案手法

• DRaFTを提案する。この手法は、最終的に生成された画像の報酬から、展開された全拡散サンプリングプロセスにわたり勾配をバックプロパゲートすることで、拡散モデルのエンドツーエンド微調整を可能にする。
• メモリコストを低減するため勾配チェックポイントを用い、パラメータの小さな部分のみを微調整するLoRAを活用することで、計算効率を維持する。
• DRaFT-$K$ を導入し、最後の$K$ 個のノイズ除去ステップにのみバックプロパゲーションを制限することで、訓練の安定性と勾配爆発の低減を図る。
• DRaFT-LVを提案する。これはDRaFT-1の変種であり、複数のノイズサンプルを用いてより低い分散を持つ勾配推定値を計算することで、訓練の安定性と速度を向上させる。
• LoRA重みの適応を用いて、Stable Diffusion 1.4をLAION Aesthetics、PickScore、HPSv2といったさまざまな微分可能報酬関数で微調整する。
• 異なる報酬目的で学習されたLoRA重みを混合またはスケーリングすることで、モデルの合成と補間を可能にし、再トレーニングなしに多目的生成を実現する。

実験結果

リサーチクエスチョン

• RQ1 全拡散サンプリングプロセスを経由した勾配ベースの微調整は、微分可能報酬関数において強化学習よりも優れた性能を達成できるか？
• RQ2 最後の$K$ ステップにのみバックプロパゲーションを制限する（DRaFT-$K$）ことで、完全なバックプロパゲーションと比較して訓練の安定性と収束性にどのような影響を与えるか？
• RQ3 複数のノイズサンプルの平均化といった分散低減技術（DRaFT-LV）は、訓練の効率性と収束速度を著しく向上させられるか？
• RQ4 DRaFTは、画像の圧縮性、オブジェクトの除去、敵対的例生成といった多様な報酬関数にどの程度一般化可能か？
• RQ5 微調整済みモデルは、LoRA重みの混合により合成または補間可能であり、複数の目的における性能を維持できるか？

主な発見

• LAION Aestheticsスコアを最大化する際、DRaFTは強化学習ベースのベースラインと比較して200倍以上高速な訓練速度を達成した。
• DRaFT-LVは、勾配推定の分散が低いため、先行する勾配ベース手法ReFLと比較して約2倍の高速な訓練が可能となった。
• DRaFT-$K$ で$K=2$ を使用した場合、実際の勾配爆発の問題により、完全なバックプロパゲーションよりも優れた性能を達成した。
• PickScoreおよびHPSv2で微調整したモデルは、人間の好みスコアにおいて顕著な向上を示し、LoRA重みスケーリングにより過学習のリスクが低減された。
• 異なる報酬で学習されたLoRA重みの間で、効果的な補間が可能であり、再トレーニングなしに合成生成が実現した。
• DRaFTは、多様なプロンプトにおいて画像品質の向上に成功し、画像の非圧縮性や敵対的例生成といった複雑な目的に対しても対応した。