Alineación de Modelos de Difusión de Texto a Imagen con Retropropagación de Recompensas

Aligning Text-to-Image Diffusion Models with Reward Backpropagation

October 5, 2023

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

cs.AI

Resumen

Los modelos de difusión de texto a imagen han surgido recientemente a la vanguardia de la generación de imágenes, impulsados por conjuntos de datos de entrenamiento de texto a imagen a muy gran escala, no supervisados o débilmente supervisados. Debido a su entrenamiento no supervisado, controlar su comportamiento en tareas posteriores, como maximizar la calidad de la imagen percibida por humanos, la alineación entre imagen y texto, o la generación ética de imágenes, es difícil. Trabajos recientes ajustan los modelos de difusión a funciones de recompensa posteriores utilizando aprendizaje por refuerzo básico, conocido por la alta varianza de los estimadores de gradiente. En este artículo, proponemos AlignProp, un método que alinea los modelos de difusión a funciones de recompensa posteriores mediante la retropropagación de extremo a extremo del gradiente de recompensa a través del proceso de eliminación de ruido. Si bien una implementación ingenua de dicha retropropagación requeriría recursos de memoria prohibitivos para almacenar las derivadas parciales de los modelos modernos de texto a imagen, AlignProp ajusta módulos de pesos adaptadores de bajo rango y utiliza la verificación de gradientes para hacer viable su uso de memoria. Probamos AlignProp en el ajuste fino de modelos de difusión para varios objetivos, como la alineación semántica entre imagen y texto, la estética, la compresibilidad y la controlabilidad del número de objetos presentes, así como sus combinaciones. Mostramos que AlignProp logra mayores recompensas en menos pasos de entrenamiento que las alternativas, siendo conceptualmente más simple, lo que lo convierte en una elección directa para optimizar modelos de difusión para funciones de recompensa diferenciables de interés. El código y los resultados de visualización están disponibles en 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/.