DRAGON: Оптимизация генеративных диффузионных моделей с помощью распределенных наград

DRAGON: Distributional Rewards Optimize Diffusion Generative Models

April 21, 2025

Авторы: Yatong Bai, Jonah Casebeer, Somayeh Sojoudi, Nicholas J. Bryan

cs.AI

Аннотация

Мы представляем Distributional RewArds for Generative OptimizatioN (DRAGON) — универсальную структуру для тонкой настройки моделей генерации медиа с целью достижения желаемого результата. По сравнению с традиционным обучением с подкреплением на основе человеческой обратной связи (RLHF) или подходами, основанными на парных предпочтениях, такими как оптимизация прямых предпочтений (DPO), DRAGON обладает большей гибкостью. Он способен оптимизировать функции вознаграждения, которые оценивают как отдельные примеры, так и их распределения, что делает его совместимым с широким спектром вознаграждений: на уровне отдельных экземпляров, между экземпляром и распределением, а также между распределениями. Используя эту универсальность, мы создаем новые функции вознаграждения, выбирая кодировщик и набор эталонных примеров для формирования распределения образцов. Когда используются кросс-модальные кодировщики, такие как CLAP, эталонные примеры могут принадлежать другой модальности (например, текст против аудио). Затем DRAGON собирает онлайн-генерации и генерации на основе текущей политики, оценивает их для формирования набора положительных демонстраций и отрицательного набора, а также использует контраст между ними для максимизации вознаграждения. Для оценки мы тонко настраиваем диффузионную модель преобразования текста в музыку в аудиодомен с 20 различными функциями вознаграждения, включая пользовательскую модель эстетики музыки, оценку CLAP, разнообразие Vendi и расстояние Фреше для аудио (FAD). Мы также сравниваем настройки FAD на уровне отдельных экземпляров (на одну песню) и на уровне всего набора данных, одновременно исследуя различные кодировщики FAD и эталонные наборы. По всем 20 целевым вознаграждениям DRAGON достигает средней доли побед в 81,45%. Более того, функции вознаграждения, основанные на наборах образцов, действительно улучшают генерацию и сопоставимы с вознаграждениями, основанными на моделях. С подходящим набором образцов DRAGON достигает 60,95% доли побед в качестве музыки по оценке людей, не обучаясь на аннотациях человеческих предпочтений. Таким образом, DRAGON демонстрирует новый подход к проектированию и оптимизации функций вознаграждения для улучшения качества, воспринимаемого человеком. Примеры звуков доступны на 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.