PRDP: Proximale Beloningsverschilvoorspelling voor Grootschalige Beloningsafstemming van Diffusiemodellen

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

February 13, 2024

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

cs.AI

Samenvatting

Reward finetuning is naar voren gekomen als een veelbelovende aanpak om foundationmodellen af te stemmen op downstreamdoelstellingen. Opmerkelijke successen zijn behaald in het taalgebied door reinforcement learning (RL) te gebruiken om beloningen te maximaliseren die menselijke voorkeuren weerspiegelen. In het visuele domein worden bestaande RL-gebaseerde reward finetuning-methoden echter beperkt door hun instabiliteit bij grootschalige training, waardoor ze niet in staat zijn om te generaliseren naar complexe, onbekende prompts. In dit artikel introduceren we Proximal Reward Difference Prediction (PRDP), waarmee voor het eerst stabiele black-box reward finetuning mogelijk wordt gemaakt voor diffusiemodellen op grootschalige promptdatasets met meer dan 100K prompts. Onze belangrijkste innovatie is het Reward Difference Prediction (RDP)-doel, dat dezelfde optimale oplossing heeft als het RL-doel, maar met betere trainingsstabiliteit. Specifiek is het RDP-doel een supervised regressiedoel dat het diffusiemodel de taak geeft om het beloningsverschil van gegenereerde beeldparen te voorspellen op basis van hun denoising-trajecten. We bewijzen theoretisch dat het diffusiemodel dat perfecte beloningsverschilvoorspelling bereikt, precies de maximalisator is van het RL-doel. We ontwikkelen verder een online algoritme met proximale updates om het RDP-doel stabiel te optimaliseren. In experimenten tonen we aan dat PRDP de beloningsmaximalisatiecapaciteit van goed ingeburgerde RL-gebaseerde methoden kan evenaren bij kleinschalige training. Bovendien bereikt PRDP door grootschalige training op tekstprompts uit het Human Preference Dataset v2 en het Pick-a-Pic v1-dataset superieure generatiekwaliteit op een diverse set van complexe, onbekende prompts, terwijl RL-gebaseerde methoden volledig falen.

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.