DRAGON: Recompensas Distribucionais Otimizam Modelos Generativos de Difusão

DRAGON: Distributional Rewards Optimize Diffusion Generative Models

April 21, 2025

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

cs.AI

Resumo

Apresentamos o Distributional RewArds for Generative OptimizatioN (DRAGON), um framework versátil para ajuste fino de modelos de geração de mídia visando um resultado desejado. Em comparação com métodos tradicionais de aprendizado por reforço com feedback humano (RLHF) ou abordagens de preferência pareada, como a otimização direta de preferência (DPO), o DRAGON é mais flexível. Ele pode otimizar funções de recompensa que avaliam tanto exemplos individuais quanto distribuições deles, tornando-o compatível com uma ampla gama de recompensas baseadas em instâncias, instância-para-distribuição e distribuição-para-distribuição. Aproveitando essa versatilidade, construímos novas funções de recompensa selecionando um codificador e um conjunto de exemplos de referência para criar uma distribuição exemplar. Quando codificadores de modalidade cruzada, como o CLAP, são utilizados, os exemplos de referência podem ser de uma modalidade diferente (por exemplo, texto versus áudio). Em seguida, o DRAGON coleta gerações online e on-policy, pontua-as para construir um conjunto de demonstrações positivas e um conjunto negativo, e utiliza o contraste entre os dois conjuntos para maximizar a recompensa. Para avaliação, ajustamos finamente um modelo de difusão de texto-para-música no domínio de áudio com 20 funções de recompensa diferentes, incluindo um modelo personalizado de estética musical, pontuação CLAP, diversidade Vendi e distância de áudio Fréchet (FAD). Além disso, comparamos configurações de FAD baseadas em instâncias (por música) e em todo o conjunto de dados, enquanto realizamos ablações em múltiplos codificadores FAD e conjuntos de referência. Em todas as 20 recompensas-alvo, o DRAGON alcança uma taxa média de vitória de 81,45%. Além disso, funções de recompensa baseadas em conjuntos exemplares de fato melhoram as gerações e são comparáveis a recompensas baseadas em modelos. Com um conjunto exemplar apropriado, o DRAGON alcança uma taxa de vitória de 60,95% na qualidade musical votada por humanos, sem treinamento em anotações de preferência humana. Assim, o DRAGON apresenta uma nova abordagem para projetar e otimizar funções de recompensa visando melhorar a qualidade percebida por humanos. Exemplos de áudio estão disponíveis em 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.