DRAGON: Recompensas Distribucionales Optimizan Modelos Generativos de Difusión

DRAGON: Distributional Rewards Optimize Diffusion Generative Models

April 21, 2025

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

cs.AI

Resumen

Presentamos Distributional RewArds for Generative OptimizatioN (DRAGON), un marco versátil para el ajuste fino de modelos de generación de medios hacia un resultado deseado. En comparación con el aprendizaje por refuerzo tradicional con retroalimentación humana (RLHF) o enfoques de preferencias por pares como la optimización directa de preferencias (DPO), DRAGON es más flexible. Puede optimizar funciones de recompensa que evalúan tanto ejemplos individuales como distribuciones de ellos, lo que lo hace compatible con un amplio espectro de recompensas a nivel de instancia, de instancia a distribución y de distribución a distribución. Aprovechando esta versatilidad, construimos funciones de recompensa novedosas seleccionando un codificador y un conjunto de ejemplos de referencia para crear una distribución ejemplar. Cuando se utilizan codificadores de modalidad cruzada como CLAP, los ejemplos de referencia pueden ser de una modalidad diferente (por ejemplo, texto versus audio). Luego, DRAGON recopila generaciones en línea y en política, las puntúa para construir un conjunto de demostraciones positivas y un conjunto negativo, y aprovecha el contraste entre ambos para maximizar la recompensa. Para la evaluación, ajustamos finamente un modelo de difusión de texto a música en el dominio de audio con 20 funciones de recompensa diferentes, incluyendo un modelo personalizado de estética musical, puntuación CLAP, diversidad Vendi y distancia de Frechet en audio (FAD). Además, comparamos configuraciones de FAD a nivel de instancia (por canción) y de conjunto completo de datos, mientras realizamos ablaciones de múltiples codificadores y conjuntos de referencia de FAD. En todas las 20 recompensas objetivo, DRAGON logra una tasa promedio de éxito del 81.45%. Además, las funciones de recompensa basadas en conjuntos ejemplares mejoran efectivamente las generaciones y son comparables a las recompensas basadas en modelos. Con un conjunto ejemplar adecuado, DRAGON alcanza una tasa de éxito en calidad musical votada por humanos del 60.95% sin entrenamiento en anotaciones de preferencias humanas. De esta manera, DRAGON exhibe un nuevo enfoque para diseñar y optimizar funciones de recompensa para mejorar la calidad percibida por los humanos. Ejemplos de sonido en 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.