RL para Modelos de Consistencia: Generación Guiada por Recompensas más Rápida de Imágenes a partir de Texto

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

Resumen

El aprendizaje por refuerzo (RL, por sus siglas en inglés) ha mejorado la generación guiada de imágenes con modelos de difusión al optimizar directamente recompensas que capturan la calidad de la imagen, la estética y la capacidad de seguir instrucciones. Sin embargo, las políticas generativas resultantes heredan el mismo proceso de muestreo iterativo de los modelos de difusión, lo que provoca una generación lenta. Para superar esta limitación, los modelos de consistencia propusieron aprender una nueva clase de modelos generativos que mapean directamente el ruido a los datos, lo que resulta en un modelo capaz de generar una imagen en tan solo una iteración de muestreo. En este trabajo, para optimizar modelos generativos de texto a imagen en función de recompensas específicas de tareas y permitir un entrenamiento e inferencia rápidos, proponemos un marco para ajustar modelos de consistencia mediante RL. Nuestro marco, denominado Aprendizaje por Refuerzo para Modelos de Consistencia (RLCM, por sus siglas en inglés), enmarca el proceso de inferencia iterativa de un modelo de consistencia como un procedimiento de RL. RLCM mejora los modelos de difusión ajustados con RL en capacidades de generación de texto a imagen y equilibra el cómputo durante el tiempo de inferencia con la calidad de las muestras. Experimentalmente, demostramos que RLCM puede adaptar modelos de consistencia de texto a imagen a objetivos que son difíciles de expresar mediante indicaciones, como la compresibilidad de la imagen, y aquellos derivados de la retroalimentación humana, como la calidad estética. En comparación con los modelos de difusión ajustados con RL, RLCM entrena significativamente más rápido, mejora la calidad de la generación medida bajo los objetivos de recompensa y acelera el procedimiento de inferencia al generar imágenes de alta calidad con tan solo dos pasos de inferencia. Nuestro código está disponible en 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