PRDP: Predizione della Differenza di Ricompensa Prossimale per l'Affinamento su Larga Scala delle Ricompense nei Modelli di Diffusione

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

February 13, 2024

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

cs.AI

Abstract

Il fine-tuning basato su ricompense è emerso come un approccio promettente per allineare i modelli di base con gli obiettivi a valle. Successi notevoli sono stati ottenuti nel dominio del linguaggio utilizzando l'apprendimento per rinforzo (RL) per massimizzare ricompense che riflettono le preferenze umane. Tuttavia, nel dominio della visione, i metodi esistenti di fine-tuning basati su ricompense tramite RL sono limitati dalla loro instabilità nell'addestramento su larga scala, rendendoli incapaci di generalizzare a prompt complessi e non visti in precedenza. In questo articolo, proponiamo il Proximal Reward Difference Prediction (PRDP), che consente per la prima volta un fine-tuning stabile delle ricompense per modelli di diffusione su dataset di prompt su larga scala con oltre 100K prompt. La nostra innovazione chiave è l'obiettivo Reward Difference Prediction (RDP), che ha la stessa soluzione ottimale dell'obiettivo RL pur godendo di una maggiore stabilità nell'addestramento. Nello specifico, l'obiettivo RDP è un obiettivo di regressione supervisionata che assegna al modello di diffusione il compito di prevedere la differenza di ricompensa di coppie di immagini generate a partire dalle loro traiettorie di denoising. Dimostriamo teoricamente che il modello di diffusione che ottiene una previsione perfetta della differenza di ricompensa è esattamente il massimizzatore dell'obiettivo RL. Sviluppiamo inoltre un algoritmo online con aggiornamenti prossimali per ottimizzare in modo stabile l'obiettivo RDP. Negli esperimenti, dimostriamo che PRDP può eguagliare la capacità di massimizzazione delle ricompense dei metodi consolidati basati su RL nell'addestramento su piccola scala. Inoltre, attraverso l'addestramento su larga scala su prompt testuali provenienti dal Human Preference Dataset v2 e dal dataset Pick-a-Pic v1, PRDP raggiunge una qualità di generazione superiore su un insieme diversificato di prompt complessi e non visti in precedenza, mentre i metodi basati su RL falliscono 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.