DPOK: Aprendizado por Reforço para o Ajuste Fino de Modelos de Difusão de Texto para Imagem

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

Resumo

Aprender com o feedback humano tem se mostrado eficaz para melhorar modelos de texto para imagem. Essas técnicas primeiro aprendem uma função de recompensa que captura o que os humanos consideram importante na tarefa e, em seguida, aprimoram os modelos com base na função de recompensa aprendida. Embora abordagens relativamente simples (por exemplo, amostragem por rejeição com base em pontuações de recompensa) tenham sido investigadas, o ajuste fino de modelos de texto para imagem com a função de recompensa ainda é desafiador. Neste trabalho, propomos o uso de aprendizado por reforço online (RL) para ajustar modelos de texto para imagem. Focamos em modelos de difusão, definindo a tarefa de ajuste fino como um problema de RL e atualizando os modelos pré-treinados de difusão de texto para imagem usando gradiente de política para maximizar a recompensa treinada com feedback. Nossa abordagem, denominada DPOK, integra otimização de política com regularização KL. Realizamos uma análise da regularização KL tanto para ajuste fino com RL quanto para ajuste fino supervisionado. Em nossos experimentos, mostramos que o DPOK geralmente supera o ajuste fino supervisionado em termos de alinhamento texto-imagem e qualidade da imagem.

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.