Ausrichtung von Text-zu-Bild-Diffusionsmodellen durch Belohnungsrückpropagation

Aligning Text-to-Image Diffusion Models with Reward Backpropagation

October 5, 2023

papers.authors: Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki

cs.AI

papers.abstract

Text-to-Image-Diffusionsmodelle sind kürzlich an die Spitze der Bildgenerierung aufgestiegen, angetrieben durch sehr groß angelegte, unüberwachte oder schwach überwachte Text-to-Image-Trainingsdatensätze. Aufgrund ihres unüberwachten Trainings ist es schwierig, ihr Verhalten in nachgelagerten Aufgaben zu steuern, wie z. B. die Maximierung der vom Menschen wahrgenommenen Bildqualität, die Ausrichtung von Bild und Text oder die ethische Bildgenerierung. Aktuelle Arbeiten feinen Diffusionsmodelle an nachgelagerte Belohnungsfunktionen an, indem sie klassisches Reinforcement Learning verwenden, das für die hohe Varianz der Gradientenschätzer bekannt ist. In diesem Artikel schlagen wir AlignProp vor, eine Methode, die Diffusionsmodelle an nachgelagerte Belohnungsfunktionen anpasst, indem sie den Belohnungsgradienten durch den Denoising-Prozess mittels End-to-End-Backpropagation ausrichtet. Während eine naive Implementierung einer solchen Backpropagation prohibitive Speicherressourcen für die Speicherung der partiellen Ableitungen moderner Text-to-Image-Modelle erfordern würde, feint AlignProp niedrigrangige Adapter-Gewichtsmodule an und verwendet Gradient-Checkpointing, um den Speicherverbrauch tragbar zu machen. Wir testen AlignProp beim Feinabstimmen von Diffusionsmodellen auf verschiedene Ziele, wie die semantische Ausrichtung von Bild und Text, Ästhetik, Komprimierbarkeit und die Steuerbarkeit der Anzahl vorhandener Objekte sowie deren Kombinationen. Wir zeigen, dass AlignProp in weniger Trainingsschritten höhere Belohnungen erzielt als Alternativen, während es konzeptionell einfacher ist, was es zu einer naheliegenden Wahl für die Optimierung von Diffusionsmodellen für differenzierbare Belohnungsfunktionen von Interesse macht. Code und Visualisierungsergebnisse sind verfügbar unter https://align-prop.github.io/.

English

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/.