DPOK: Aprendizaje por Refuerzo para el Ajuste Fino de Modelos de Difusión de Texto a Imagen

DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models

May 25, 2023

Autores: Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, Kimin Lee

cs.AI

Resumen

Aprender a partir de la retroalimentación humana ha demostrado mejorar los modelos de texto a imagen. Estas técnicas primero aprenden una función de recompensa que captura lo que los humanos valoran en la tarea y luego mejoran los modelos basándose en la función de recompensa aprendida. Aunque se han investigado enfoques relativamente simples (por ejemplo, muestreo por rechazo basado en puntajes de recompensa), el ajuste fino de los modelos de texto a imagen con la función de recompensa sigue siendo un desafío. En este trabajo, proponemos utilizar el aprendizaje por refuerzo (RL) en línea para ajustar finamente los modelos de texto a imagen. Nos enfocamos en los modelos de difusión, definiendo la tarea de ajuste fino como un problema de RL y actualizando los modelos preentrenados de texto a imagen utilizando el gradiente de política para maximizar la recompensa entrenada con retroalimentación. Nuestro enfoque, denominado DPOK, integra la optimización de políticas con la regularización KL. Realizamos un análisis de la regularización KL tanto para el ajuste fino con RL como para el ajuste fino supervisado. En nuestros experimentos, demostramos que DPOK es generalmente superior al ajuste fino supervisado en cuanto a la alineación texto-imagen y la calidad de la imagen.

English

Learning from human feedback has been shown to improve text-to-image models. These techniques first learn a reward function that captures what humans care about in the task and then improve the models based on the learned reward function. Even though relatively simple approaches (e.g., rejection sampling based on reward scores) have been investigated, fine-tuning text-to-image models with the reward function remains challenging. In this work, we propose using online reinforcement learning (RL) to fine-tune text-to-image models. We focus on diffusion models, defining the fine-tuning task as an RL problem, and updating the pre-trained text-to-image diffusion models using policy gradient to maximize the feedback-trained reward. Our approach, coined DPOK, integrates policy optimization with KL regularization. We conduct an analysis of KL regularization for both RL fine-tuning and supervised fine-tuning. In our experiments, we show that DPOK is generally superior to supervised fine-tuning with respect to both image-text alignment and image quality.