[論文レビュー] Aligning Text-to-Image Diffusion Models with Reward Backpropagation

Mihir Prabhudesai, Anirudh Goyal+2 more|arXiv (Cornell University)|Oct 5, 2023

Generative Adversarial Networks and Image Synthesis被引用数 5

ひとこと要約

この論文では、ノイズ除去プロセスを通じて微分可能報酬勾配をエンドツーエンドバックプロパゲーションすることで、テキストから画像への拡散モデルをファインチューニングする方法であるAlignPropを提案する。低ランクアダプタと勾配チェックポイントを用いることで、強化学習ベースラインと比較して少ないトレーニングステップでより高い報酬性能を達成し、画像・テキストの整合性、アート性、制御性といった多様な目的の効率的最適化を可能にする。

ABSTRACT

Text-to-image diffusion models have recently emerged at the forefront of image generation, powered by very large-scale unsupervised or weakly supervised text-to-image training datasets. Due to their unsupervised training, controlling their behavior in downstream tasks, such as maximizing human-perceived image quality, image-text alignment, or ethical image generation, is difficult. Recent works finetune diffusion models to downstream reward functions using vanilla reinforcement learning, notorious for the high variance of the gradient estimators. In this paper, we propose AlignProp, a method that aligns diffusion models to downstream reward functions using end-to-end backpropagation of the reward gradient through the denoising process. While naive implementation of such backpropagation would require prohibitive memory resources for storing the partial derivatives of modern text-to-image models, AlignProp finetunes low-rank adapter weight modules and uses gradient checkpointing, to render its memory usage viable. We test AlignProp in finetuning diffusion models to various objectives, such as image-text semantic alignment, aesthetics, compressibility and controllability of the number of objects present, as well as their combinations. We show AlignProp achieves higher rewards in fewer training steps than alternatives, while being conceptually simpler, making it a straightforward choice for optimizing diffusion models for differentiable reward functions of interest. Code and Visualization results are available at https://align-prop.github.io/.

研究の動機と目的

• 画像品質、公平性、テキスト・画像の整合性といった、画像生成品質やフェアネスを向上させる下流の報酬関数と、事前学習済みのテキストから画像への拡散モデルを一致させる課題に対処すること。
• 拡散モデルにおける長いノイズ除去チェーンを経由するエンドツーエンドバックプロパゲーションの、記憶容量を圧倒的に増加させるコストを克服すること。
• 強化学習の高分散勾配を避けるために、報酬勾配を直接微分可能に最適化できる手法を開発すること。
• LoRAアダプタのパラメータ平均化により複数の報酬関数間の補間を可能とし、マルチオブジェクティブな整合性を実現すること。

提案手法

• ノイズとテキストプロンプトから画像へのマッピングを、微分可能で再帰的なポリシーとして扱い、報酬からモデル重みへの勾配伝播を可能にする。
• 低ランクアダプタ（LoRA）モジュールを用いて、パラメータの小さなサブセットのみをファインチューニングすることで、フルファインチューニングと比較して記憶使用量を著しく削減する。
• 勾配チェックポイントを採用し、部分的勾配を事前に保存せず、必要に応じて計算することで、テラバイト規模の記憶容量を回避する。
• 固定されたコンテキスト長Kを用いた、時間方向に切り詰めたバックプロパゲーション（BPTT）を適用し、勾配品質と記憶効率のバランスを取る。
• 別々に異なる報酬関数でファインチューニングされたモデルのLoRA重みを平均化することで、モデルの補間を可能とし、マルチオブジェクティブな生成を実現する。
• CLIPに基づく画像・テキストの整合性、アートスコア、JPEG圧縮性といった微分可能な報酬関数を用い、最適化をガイドする。

実験結果

リサーチクエスチョン

• RQ1 拡散サンプリングプロセスを経由するエンドツーエンドバックプロパゲーションを、大規模なテキストから画像へのモデルに実用的になるほど記憶効率よく実装できるか？
• RQ2 サンプル効率および最終的な報酬性能の観点で、AlignPropは強化学習ベースラインと比べてどのように異なるか？
• RQ3 別々に訓練されたモデルのLoRAアダプタを組み合わせることで、AlignPropは複数の報酬関数の間で効果的に補間できるか？
• RQ4 低ランクアダプタと勾配チェックポイントの使用は、性能を維持しつつ、記憶容量と計算コストを削減できるか？

主な発見

• AlignPropは、ReFL や DDPO といったベースラインと比較して顕著に高い平均報酬を達成し、HPSv2ベンチマークでは平均報酬3.30を記録した（K=1のとき2.91と比較）。
• 高い報酬性能に到達するためのトレーニングステップ数を削減でき、強化学習手法と比較して優れたサンプル効率を示した。
• アート性と圧縮性で事前にファインチューニングされたモデルのLoRA重みを平均化することで、ハイブリッドモデルが得られ、混合係数α=0.5で最高の全体的報酬を達成した。
• アブレーションスタディの結果、LoRAファインチューニングはフルUNetファインチューニングを上回り、EDICTレイヤーは記憶容量を削減するが、性能を低下させ、推論時間を延長することが判明した。
• K=10で切り詰めたバックプロパゲーションはK=1よりも優れた性能を示しており、より長いコンテキストが勾配品質を向上させつつ、記憶コストを著しく増加させないことが示された。
• パラメータ補間を用いることで、画像・テキストの整合性とアート性を同時に向上させるような、複雑なマルチオブジェクティブ報酬に対しても、効果的な整合性が実現された。