DRAGON : Les récompenses distributionnelles optimisent les modèles génératifs de diffusion

DRAGON: Distributional Rewards Optimize Diffusion Generative Models

April 21, 2025

Auteurs: Yatong Bai, Jonah Casebeer, Somayeh Sojoudi, Nicholas J. Bryan

cs.AI

Résumé

Nous présentons Distributional RewArds for Generative OptimizatioN (DRAGON), un cadre polyvalent pour le fine-tuning de modèles de génération de médias en vue d'un résultat souhaité. Comparé aux approches traditionnelles de renforcement avec feedback humain (RLHF) ou de préférences par paires telles que l'optimisation directe des préférences (DPO), DRAGON offre une plus grande flexibilité. Il permet d'optimiser des fonctions de récompense qui évaluent soit des exemples individuels, soit des distributions d'exemples, le rendant compatible avec un large éventail de récompenses basées sur des instances, des instances à distribution, et des distributions à distributions. En exploitant cette polyvalence, nous construisons des fonctions de récompense innovantes en sélectionnant un encodeur et un ensemble d'exemples de référence pour créer une distribution d'exemplaires. Lorsque des encodeurs intermodaux comme CLAP sont utilisés, les exemples de référence peuvent appartenir à une modalité différente (par exemple, texte versus audio). Ensuite, DRAGON collecte des générations en ligne et sur politique, les évalue pour construire un ensemble de démonstrations positives et un ensemble négatif, et exploite le contraste entre ces deux ensembles pour maximiser la récompense. Pour l'évaluation, nous effectuons le fine-tuning d'un modèle de diffusion texte-à-musique dans le domaine audio avec 20 fonctions de récompense différentes, incluant un modèle personnalisé d'esthétique musicale, le score CLAP, la diversité Vendi, et la distance audio de Fréchet (FAD). Nous comparons également les paramètres FAD par instance (par chanson) et sur l'ensemble du dataset tout en effectuant une ablation sur plusieurs encodeurs FAD et ensembles de référence. Sur l'ensemble des 20 récompenses cibles, DRAGON atteint un taux de réussite moyen de 81,45%. De plus, les fonctions de récompense basées sur des ensembles d'exemplaires améliorent effectivement les générations et sont comparables aux récompenses basées sur des modèles. Avec un ensemble d'exemplaires approprié, DRAGON atteint un taux de réussite de 60,95% en qualité musicale votée par des humains sans entraînement sur des annotations de préférences humaines. Ainsi, DRAGON illustre une nouvelle approche pour concevoir et optimiser des fonctions de récompense afin d'améliorer la qualité perçue par les humains. Des exemples sonores sont disponibles à l'adresse https://ml-dragon.github.io/web.

English

We present Distributional RewArds for Generative OptimizatioN (DRAGON), a versatile framework for fine-tuning media generation models towards a desired outcome. Compared with traditional reinforcement learning with human feedback (RLHF) or pairwise preference approaches such as direct preference optimization (DPO), DRAGON is more flexible. It can optimize reward functions that evaluate either individual examples or distributions of them, making it compatible with a broad spectrum of instance-wise, instance-to-distribution, and distribution-to-distribution rewards. Leveraging this versatility, we construct novel reward functions by selecting an encoder and a set of reference examples to create an exemplar distribution. When cross-modality encoders such as CLAP are used, the reference examples may be of a different modality (e.g., text versus audio). Then, DRAGON gathers online and on-policy generations, scores them to construct a positive demonstration set and a negative set, and leverages the contrast between the two sets to maximize the reward. For evaluation, we fine-tune an audio-domain text-to-music diffusion model with 20 different reward functions, including a custom music aesthetics model, CLAP score, Vendi diversity, and Frechet audio distance (FAD). We further compare instance-wise (per-song) and full-dataset FAD settings while ablating multiple FAD encoders and reference sets. Over all 20 target rewards, DRAGON achieves an 81.45% average win rate. Moreover, reward functions based on exemplar sets indeed enhance generations and are comparable to model-based rewards. With an appropriate exemplar set, DRAGON achieves a 60.95% human-voted music quality win rate without training on human preference annotations. As such, DRAGON exhibits a new approach to designing and optimizing reward functions for improving human-perceived quality. Sound examples at https://ml-dragon.github.io/web.