PRDP: Predicción de Diferencia de Recompensa Proximal para el Ajuste Fino a Gran Escala de Recompensas en Modelos de Difusión

PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models

February 13, 2024

Autores: Fei Deng, Qifei Wang, Wei Wei, Matthias Grundmann, Tingbo Hou

cs.AI

Resumen

El ajuste fino mediante recompensas ha surgido como un enfoque prometedor para alinear modelos base con objetivos específicos. Se ha logrado un éxito notable en el dominio del lenguaje al utilizar aprendizaje por refuerzo (RL) para maximizar recompensas que reflejan preferencias humanas. Sin embargo, en el dominio visual, los métodos existentes de ajuste fino basados en RL están limitados por su inestabilidad en entrenamientos a gran escala, lo que les impide generalizar a indicaciones complejas y no vistas. En este artículo, proponemos Predicción de Diferencia de Recompensa Proximal (PRDP), permitiendo por primera vez un ajuste fino estable de recompensas en modelos de difusión a gran escala con más de 100K indicaciones. Nuestra innovación clave es el objetivo de Predicción de Diferencia de Recompensa (RDP), que tiene la misma solución óptima que el objetivo de RL, pero con mayor estabilidad en el entrenamiento. Específicamente, el objetivo RDP es un objetivo de regresión supervisada que requiere que el modelo de difusión prediga la diferencia de recompensa entre pares de imágenes generadas a partir de sus trayectorias de eliminación de ruido. Teóricamente demostramos que el modelo de difusión que logra una predicción perfecta de la diferencia de recompensa es exactamente el maximizador del objetivo de RL. Además, desarrollamos un algoritmo en línea con actualizaciones proximales para optimizar de manera estable el objetivo RDP. En los experimentos, demostramos que PRDP puede igualar la capacidad de maximización de recompensas de métodos bien establecidos basados en RL en entrenamientos a pequeña escala. Además, a través de entrenamientos a gran escala en indicaciones de texto del Conjunto de Datos de Preferencias Humanas v2 y el Conjunto de Datos Pick-a-Pic v1, PRDP logra una calidad de generación superior en un conjunto diverso de indicaciones complejas y no vistas, mientras que los métodos basados en RL fracasan por completo.

English

Reward finetuning has emerged as a promising approach to aligning foundation models with downstream objectives. Remarkable success has been achieved in the language domain by using reinforcement learning (RL) to maximize rewards that reflect human preference. However, in the vision domain, existing RL-based reward finetuning methods are limited by their instability in large-scale training, rendering them incapable of generalizing to complex, unseen prompts. In this paper, we propose Proximal Reward Difference Prediction (PRDP), enabling stable black-box reward finetuning for diffusion models for the first time on large-scale prompt datasets with over 100K prompts. Our key innovation is the Reward Difference Prediction (RDP) objective that has the same optimal solution as the RL objective while enjoying better training stability. Specifically, the RDP objective is a supervised regression objective that tasks the diffusion model with predicting the reward difference of generated image pairs from their denoising trajectories. We theoretically prove that the diffusion model that obtains perfect reward difference prediction is exactly the maximizer of the RL objective. We further develop an online algorithm with proximal updates to stably optimize the RDP objective. In experiments, we demonstrate that PRDP can match the reward maximization ability of well-established RL-based methods in small-scale training. Furthermore, through large-scale training on text prompts from the Human Preference Dataset v2 and the Pick-a-Pic v1 dataset, PRDP achieves superior generation quality on a diverse set of complex, unseen prompts whereas RL-based methods completely fail.