PRDP: Predição de Diferença de Recompensa Proximal para Ajuste Fino em Larga Escala de Recompensas em Modelos de Difusão

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

Resumo

O ajuste fino baseado em recompensas surgiu como uma abordagem promissora para alinhar modelos de base com objetivos específicos. Sucessos notáveis foram alcançados no domínio da linguagem ao usar aprendizado por reforço (RL) para maximizar recompensas que refletem preferências humanas. No entanto, no domínio da visão, os métodos existentes de ajuste fino baseados em RL são limitados por sua instabilidade em treinamentos em larga escala, tornando-os incapazes de generalizar para prompts complexos e não vistos. Neste artigo, propomos a Predição de Diferença de Recompensa Proximal (PRDP), permitindo pela primeira vez o ajuste fino estável de recompensas em modelos de difusão em grandes conjuntos de dados de prompts com mais de 100 mil prompts. Nossa principal inovação é o objetivo de Predição de Diferença de Recompensa (RDP), que tem a mesma solução ótima que o objetivo de RL, mas com maior estabilidade no treinamento. Especificamente, o objetivo RDP é uma tarefa de regressão supervisionada que exige que o modelo de difusão preveja a diferença de recompensa entre pares de imagens geradas a partir de suas trajetórias de remoção de ruído. Provamos teoricamente que o modelo de difusão que obtém uma predição perfeita da diferença de recompensa é exatamente o maximizador do objetivo de RL. Desenvolvemos ainda um algoritmo online com atualizações proximais para otimizar de forma estável o objetivo RDP. Nos experimentos, demonstramos que o PRDP pode igualar a capacidade de maximização de recompensa de métodos bem estabelecidos baseados em RL em treinamentos em pequena escala. Além disso, por meio de treinamento em larga escala em prompts de texto do Human Preference Dataset v2 e do Pick-a-Pic v1, o PRDP alcança qualidade de geração superior em um conjunto diversificado de prompts complexos e não vistos, enquanto métodos baseados em RL falham completamente.

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.