RL para Modelos de Consistência: Geração de Imagens a partir de Texto Orientada por Recompensa de Forma Mais Rápida

RL for Consistency Models: Faster Reward Guided Text-to-Image Generation

March 25, 2024

Autores: Owen Oertell, Jonathan D. Chang, Yiyi Zhang, Kianté Brantley, Wen Sun

cs.AI

Resumo

O aprendizado por reforço (RL) aprimorou a geração guiada de imagens com modelos de difusão ao otimizar diretamente recompensas que capturam qualidade de imagem, estética e capacidade de seguir instruções. No entanto, as políticas gerativas resultantes herdam o mesmo processo iterativo de amostragem dos modelos de difusão, o que causa geração lenta. Para superar essa limitação, os modelos de consistência propuseram aprender uma nova classe de modelos gerativos que mapeiam diretamente ruído para dados, resultando em um modelo capaz de gerar uma imagem em tão poucas quanto uma iteração de amostragem. Neste trabalho, para otimizar modelos gerativos de texto para imagem em recompensas específicas de tarefas e permitir treinamento e inferência rápidos, propomos uma estrutura para ajuste fino de modelos de consistência via RL. Nossa estrutura, chamada Aprendizado por Reforço para Modelo de Consistência (RLCM), enquadra o processo iterativo de inferência de um modelo de consistência como um procedimento de RL. O RLCM melhora os modelos de difusão ajustados por RL em capacidades de geração de texto para imagem e troca computação durante o tempo de inferência por qualidade de amostra. Experimentalmente, mostramos que o RLCM pode adaptar modelos de consistência de texto para imagem a objetivos que são desafiadores de expressar com prompts, como compressibilidade de imagem, e aqueles derivados de feedback humano, como qualidade estética. Comparado aos modelos de difusão ajustados por RL, o RLCM treina significativamente mais rápido, melhora a qualidade da geração medida sob os objetivos de recompensa e acelera o procedimento de inferência ao gerar imagens de alta qualidade com tão poucas quanto duas etapas de inferência. Nosso código está disponível em https://rlcm.owenoertell.com.

English

Reinforcement learning (RL) has improved guided image generation with diffusion models by directly optimizing rewards that capture image quality, aesthetics, and instruction following capabilities. However, the resulting generative policies inherit the same iterative sampling process of diffusion models that causes slow generation. To overcome this limitation, consistency models proposed learning a new class of generative models that directly map noise to data, resulting in a model that can generate an image in as few as one sampling iteration. In this work, to optimize text-to-image generative models for task specific rewards and enable fast training and inference, we propose a framework for fine-tuning consistency models via RL. Our framework, called Reinforcement Learning for Consistency Model (RLCM), frames the iterative inference process of a consistency model as an RL procedure. RLCM improves upon RL fine-tuned diffusion models on text-to-image generation capabilities and trades computation during inference time for sample quality. Experimentally, we show that RLCM can adapt text-to-image consistency models to objectives that are challenging to express with prompting, such as image compressibility, and those derived from human feedback, such as aesthetic quality. Comparing to RL finetuned diffusion models, RLCM trains significantly faster, improves the quality of the generation measured under the reward objectives, and speeds up the inference procedure by generating high quality images with as few as two inference steps. Our code is available at https://rlcm.owenoertell.com