PRDP: Proximale Belohnungsdifferenzvorhersage für die großskalige Belohnungsfeinabstimmung von Diffusionsmodellen

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

February 13, 2024

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

cs.AI

Zusammenfassung

Reward-Finetuning hat sich als vielversprechender Ansatz zur Ausrichtung von Foundation-Modellen auf nachgelagerte Ziele erwiesen. Im Sprachbereich wurden bemerkenswerte Erfolge erzielt, indem Reinforcement Learning (RL) eingesetzt wurde, um Belohnungen zu maximieren, die menschliche Präferenzen widerspiegeln. Im visuellen Bereich sind jedoch bestehende RL-basierte Reward-Finetuning-Methoden durch ihre Instabilität im großflächigen Training eingeschränkt, was sie unfähig macht, komplexe, unbekannte Prompts zu generalisieren. In diesem Artikel schlagen wir Proximal Reward Difference Prediction (PRDP) vor, das erstmals stabiles Black-Box-Reward-Finetuning für Diffusionsmodelle auf großflächigen Prompt-Datensätzen mit über 100.000 Prompts ermöglicht. Unsere zentrale Innovation ist das Reward Difference Prediction (RDP)-Ziel, das die gleiche optimale Lösung wie das RL-Ziel hat, jedoch eine bessere Trainingsstabilität aufweist. Konkret ist das RDP-Ziel ein überwachtes Regressionsziel, bei dem das Diffusionsmodell die Aufgabe hat, die Belohnungsdifferenz von generierten Bildpaaren aus ihren Denoising-Trajektorien vorherzusagen. Wir beweisen theoretisch, dass das Diffusionsmodell, das eine perfekte Belohnungsdifferenzvorhersage erreicht, genau der Maximierer des RL-Ziels ist. Wir entwickeln weiterhin einen Online-Algorithmus mit proximalen Updates, um das RDP-Ziel stabil zu optimieren. In Experimenten zeigen wir, dass PRDP die Belohnungsmaximierungsfähigkeit etablierter RL-basierter Methoden im kleinflächigen Training erreichen kann. Darüber hinaus erreicht PRDP durch großflächiges Training auf Text-Prompts aus dem Human Preference Dataset v2 und dem Pick-a-Pic v1-Datensatz eine überlegene Generierungsqualität bei einer Vielzahl komplexer, unbekannter Prompts, während RL-basierte Methoden vollständig versagen.

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.