Allineamento dei modelli di diffusione testo-immagine con retropropagazione del reward

Aligning Text-to-Image Diffusion Models with Reward Backpropagation

October 5, 2023

Autori: Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki

cs.AI

Abstract

I modelli di diffusione text-to-image sono recentemente emersi in prima linea nella generazione di immagini, alimentati da dataset di addestramento text-to-image su larga scala, non supervisionati o debolmente supervisionati. A causa del loro addestramento non supervisionato, controllare il loro comportamento in task downstream, come massimizzare la qualità percepita dell'immagine, l'allineamento immagine-testo o la generazione etica di immagini, è difficile. Recenti lavori ottimizzano i modelli di diffusione per funzioni di reward downstream utilizzando il classico reinforcement learning, noto per l'elevata varianza degli stimatori del gradiente. In questo articolo, proponiamo AlignProp, un metodo che allinea i modelli di diffusione alle funzioni di reward downstream utilizzando la backpropagation end-to-end del gradiente del reward attraverso il processo di denoising. Mentre un'implementazione ingenua di tale backpropagation richiederebbe risorse di memoria proibitive per memorizzare le derivate parziali dei moderni modelli text-to-image, AlignProp ottimizza moduli di peso adattatori a basso rango e utilizza il gradient checkpointing, rendendo così fattibile l'uso della memoria. Testiamo AlignProp nell'ottimizzazione di modelli di diffusione per vari obiettivi, come l'allineamento semantico immagine-testo, l'estetica, la comprimibilità e la controllabilità del numero di oggetti presenti, nonché le loro combinazioni. Mostriamo che AlignProp raggiunge reward più elevati in meno passi di addestramento rispetto alle alternative, pur essendo concettualmente più semplice, rendendolo una scelta diretta per ottimizzare i modelli di diffusione per funzioni di reward differenziabili di interesse. Codice e risultati di visualizzazione sono disponibili su 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/.