Search papers, labs, and topics across Lattice.
The paper introduces Distributional RewArds for Generative OptimizatioN (DRAGON), a framework for fine-tuning media generation models using flexible reward functions that can evaluate individual examples or distributions. DRAGON constructs positive and negative demonstration sets based on online generations scored by the reward function, and then optimizes the model by contrasting these sets to approximate distributional reward optimization. Experiments fine-tuning a text-to-music diffusion model with 20 reward functions demonstrate DRAGON's effectiveness, achieving an 81.45% average win rate and a 60.95% human-voted music quality win rate without human preference annotations.
Forget RLHF and DPO – DRAGON lets you fine-tune generative models with rewards that compare entire *distributions* of outputs, unlocking better control and quality without human preference data.
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-modal encoders such as CLAP are used, the reference may be of a different modality (text versus audio). Then, DRAGON gathers online and on-policy generations, scores them with the reward function to construct a positive demonstration set and a negative set, and leverages the contrast between the two finite sets to approximate distributional reward optimization. For evaluation, we fine-tune an audio-domain text-to-music diffusion model with 20 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 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. DRAGON is a new approach to designing and optimizing reward functions for improving human-perceived quality. Demos at https://ml-dragon.github.io/web
