DRAGON: Distributiebeloningen Optimaliseren Diffusie Generatieve Modellen

DRAGON: Distributional Rewards Optimize Diffusion Generative Models

April 21, 2025

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

cs.AI

Samenvatting

We presenteren Distributional RewArds for Generative OptimizatioN (DRAGON), een veelzijdig raamwerk voor het finetunen van mediageneratiemodellen naar een gewenst resultaat. In vergelijking met traditionele reinforcement learning met menselijke feedback (RLHF) of paarsgewijze voorkeursbenaderingen zoals direct preference optimization (DPO), is DRAGON flexibeler. Het kan beloningsfuncties optimaliseren die individuele voorbeelden of distributies daarvan evalueren, waardoor het compatibel is met een breed scala aan instance-wise, instance-to-distribution en distribution-to-distribution beloningen. Door gebruik te maken van deze veelzijdigheid, construeren we nieuwe beloningsfuncties door een encoder en een set referentievoorbeelden te selecteren om een voorbeelddistributie te creëren. Wanneer cross-modality encoders zoals CLAP worden gebruikt, kunnen de referentievoorbeelden van een andere modaliteit zijn (bijvoorbeeld tekst versus audio). Vervolgens verzamelt DRAGON online en on-policy generaties, scoort deze om een positieve demonstratieset en een negatieve set te construeren, en benut het contrast tussen de twee sets om de beloning te maximaliseren. Voor evaluatie finetunen we een audio-domein tekst-naar-muziek diffusiemodel met 20 verschillende beloningsfuncties, waaronder een aangepast muziekaesthetisch model, CLAP-score, Vendi-diversiteit en Frechet audio distance (FAD). We vergelijken verder instance-wise (per liedje) en full-dataset FAD-instellingen terwijl we meerdere FAD-encoders en referentiesets ablaten. Over alle 20 doelbeloningen behaalt DRAGON een gemiddelde winstratio van 81,45%. Bovendien verbeteren beloningsfuncties gebaseerd op voorbeeldsets inderdaad de generaties en zijn ze vergelijkbaar met modelgebaseerde beloningen. Met een geschikte voorbeeldset behaalt DRAGON een 60,95% door mensen gestemde muziekkwaliteit winstratio zonder training op menselijke voorkeursannotaties. Zo toont DRAGON een nieuwe benadering voor het ontwerpen en optimaliseren van beloningsfuncties om de door mensen waargenomen kwaliteit te verbeteren. Geluidsvoorbeelden op 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.