RL pour les modèles de cohérence : Génération d'images à partir de texte guidée par récompense plus rapide

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

March 25, 2024

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

cs.AI

Résumé

L'apprentissage par renforcement (Reinforcement Learning, RL) a amélioré la génération d'images guidée avec des modèles de diffusion en optimisant directement des récompenses qui capturent la qualité de l'image, l'esthétique et la capacité à suivre des instructions. Cependant, les politiques génératives résultantes héritent du même processus d'échantillonnage itératif des modèles de diffusion, ce qui entraîne une génération lente. Pour surmonter cette limitation, les modèles de cohérence ont proposé d'apprendre une nouvelle classe de modèles génératifs qui mappent directement le bruit aux données, aboutissant à un modèle capable de générer une image en aussi peu qu'une seule itération d'échantillonnage. Dans ce travail, afin d'optimiser les modèles de génération texte-à-image pour des récompenses spécifiques à une tâche et de permettre un entraînement et une inférence rapides, nous proposons un cadre pour le réglage fin des modèles de cohérence via l'apprentissage par renforcement. Notre cadre, appelé Reinforcement Learning for Consistency Model (RLCM), formalise le processus d'inférence itératif d'un modèle de cohérence comme une procédure d'apprentissage par renforcement. RLCM améliore les capacités de génération texte-à-image par rapport aux modèles de diffusion réglés finement par RL et échange du temps de calcul pendant l'inférence contre la qualité des échantillons. Expérimentalement, nous montrons que RLCM peut adapter des modèles de cohérence texte-à-image à des objectifs difficiles à exprimer par des invites, tels que la compressibilité de l'image, et ceux dérivés de retours humains, comme la qualité esthétique. Comparé aux modèles de diffusion réglés finement par RL, RLCM s'entraîne significativement plus vite, améliore la qualité de la génération mesurée sous les objectifs de récompense, et accélère la procédure d'inférence en générant des images de haute qualité en aussi peu que deux étapes d'inférence. Notre code est disponible à l'adresse 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