Alinhando Modelos de Difusão de Texto para Imagem com Retropropagação de Recompensa

Aligning Text-to-Image Diffusion Models with Reward Backpropagation

October 5, 2023

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

cs.AI

Resumo

Modelos de difusão texto-imagem surgiram recentemente na vanguarda da geração de imagens, impulsionados por conjuntos de dados de treinamento texto-imagem em grande escala, não supervisionados ou fracamente supervisionados. Devido ao seu treinamento não supervisionado, controlar seu comportamento em tarefas subsequentes, como maximizar a qualidade da imagem percebida por humanos, o alinhamento imagem-texto ou a geração ética de imagens, é desafiador. Trabalhos recentes ajustam modelos de difusão para funções de recompensa subsequentes usando aprendizado por reforço tradicional, conhecido pela alta variância dos estimadores de gradiente. Neste artigo, propomos o AlignProp, um método que alinha modelos de difusão a funções de recompensa subsequentes usando retropropagação end-to-end do gradiente de recompensa através do processo de remoção de ruído. Embora uma implementação ingênua dessa retropropagação exigiria recursos de memória proibitivos para armazenar as derivadas parciais de modelos modernos texto-imagem, o AlignProp ajusta módulos de pesos de adaptadores de baixa classificação e utiliza checkpointing de gradiente, tornando seu uso de memória viável. Testamos o AlignProp no ajuste fino de modelos de difusão para vários objetivos, como alinhamento semântico imagem-texto, estética, compressibilidade e controlabilidade do número de objetos presentes, bem como suas combinações. Mostramos que o AlignProp alcança recompensas mais altas em menos etapas de treinamento do que as alternativas, sendo conceitualmente mais simples, tornando-o uma escolha direta para otimizar modelos de difusão para funções de recompensa diferenciáveis de interesse. Código e resultados de visualização estão disponíveis em 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/.