RL für Konsistenzmodelle: Schnellere belohnungsgesteuerte Text-zu-Bild-Generierung

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

March 25, 2024

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

cs.AI

Zusammenfassung

Das Reinforcement-Learning (RL) hat die gesteuerte Bildgenerierung mit Diffusionsmodellen verbessert, indem es direkt Belohnungen optimiert, die die Bildqualität, Ästhetik und die Fähigkeit zur Befolgung von Anweisungen erfassen. Die resultierenden generativen Richtlinien erben jedoch den gleichen iterativen Prozess der Stichprobenentnahme von Diffusionsmodellen, der zu einer langsamen Generierung führt. Um diese Einschränkung zu überwinden, schlagen Konsistenzmodelle vor, eine neue Klasse generativer Modelle zu erlernen, die direkt Rauschen in Daten abbilden, wodurch ein Modell entsteht, das ein Bild in nur einer Stichprobeniteration generieren kann. In dieser Arbeit schlagen wir zur Optimierung von Text-zu-Bild-generativen Modellen für aufgabenspezifische Belohnungen und zur Ermöglichung schnellen Trainings und Inferenz ein Framework zur Feinabstimmung von Konsistenzmodellen über RL vor. Unser Framework, genannt Reinforcement Learning for Consistency Model (RLCM), fasst den iterativen Inferenzprozess eines Konsistenzmodells als ein RL-Verfahren auf. RLCM verbessert die Fähigkeiten von RL-feinabgestimmten Diffusionsmodellen in der Text-zu-Bild-Generierung und tauscht Rechenleistung während der Inferenzzeit gegen Probenqualität. Experimentell zeigen wir, dass RLCM Text-zu-Bild-Konsistenzmodelle an Zielen anpassen kann, die schwer mit Aufforderungen auszudrücken sind, wie z.B. Bildkomprimierbarkeit, und solche, die sich aus menschlichem Feedback ableiten, wie z.B. ästhetische Qualität. Im Vergleich zu RL-feinabgestimmten Diffusionsmodellen trainiert RLCM signifikant schneller, verbessert die Qualität der Generierung gemessen an den Belohnungszielen und beschleunigt das Inferenzverfahren, indem hochwertige Bilder mit nur zwei Inferenzschritten generiert werden. Unser Code ist verfügbar unter 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