PRDP: Прогнозирование разницы проксимальных наград для масштабной тонкой настройки наград в диффузионных моделях

PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models

February 13, 2024

Авторы: Fei Deng, Qifei Wang, Wei Wei, Matthias Grundmann, Tingbo Hou

cs.AI

Аннотация

Тонкая настройка с использованием вознаграждений стала перспективным подходом для согласования базовых моделей с целевыми задачами. В области обработки языка значительные успехи были достигнуты за счет применения обучения с подкреплением (RL) для максимизации вознаграждений, отражающих человеческие предпочтения. Однако в области компьютерного зрения существующие методы тонкой настройки на основе RL ограничены своей нестабильностью при обучении в крупных масштабах, что делает их неспособными обобщать сложные, ранее не встречавшиеся запросы. В данной статье мы предлагаем метод Proximal Reward Difference Prediction (PRDP), который впервые позволяет стабильно выполнять тонкую настройку с использованием вознаграждений для диффузионных моделей на крупных наборах данных с более чем 100 тыс. запросов. Наше ключевое нововведение — это целевая функция Reward Difference Prediction (RDP), которая имеет то же оптимальное решение, что и целевая функция RL, но при этом обеспечивает лучшую стабильность обучения. В частности, целевая функция RDP представляет собой задачу регрессии с учителем, в которой диффузионная модель должна предсказывать разницу в вознаграждениях для пар сгенерированных изображений на основе их траекторий денойзинга. Мы теоретически доказываем, что диффузионная модель, идеально предсказывающая разницу в вознаграждениях, является точным максимизатором целевой функции RL. Мы также разрабатываем онлайн-алгоритм с проксимальными обновлениями для стабильной оптимизации целевой функции RDP. В экспериментах мы показываем, что PRDP может соответствовать способности максимизации вознаграждения хорошо зарекомендовавших себя методов на основе RL при обучении в малых масштабах. Более того, при обучении на крупных наборах текстовых запросов из Human Preference Dataset v2 и Pick-a-Pic v1, PRDP демонстрирует превосходное качество генерации на разнообразных сложных, ранее не встречавшихся запросах, тогда как методы на основе RL полностью терпят неудачу.

English

Reward finetuning has emerged as a promising approach to aligning foundation models with downstream objectives. Remarkable success has been achieved in the language domain by using reinforcement learning (RL) to maximize rewards that reflect human preference. However, in the vision domain, existing RL-based reward finetuning methods are limited by their instability in large-scale training, rendering them incapable of generalizing to complex, unseen prompts. In this paper, we propose Proximal Reward Difference Prediction (PRDP), enabling stable black-box reward finetuning for diffusion models for the first time on large-scale prompt datasets with over 100K prompts. Our key innovation is the Reward Difference Prediction (RDP) objective that has the same optimal solution as the RL objective while enjoying better training stability. Specifically, the RDP objective is a supervised regression objective that tasks the diffusion model with predicting the reward difference of generated image pairs from their denoising trajectories. We theoretically prove that the diffusion model that obtains perfect reward difference prediction is exactly the maximizer of the RL objective. We further develop an online algorithm with proximal updates to stably optimize the RDP objective. In experiments, we demonstrate that PRDP can match the reward maximization ability of well-established RL-based methods in small-scale training. Furthermore, through large-scale training on text prompts from the Human Preference Dataset v2 and the Pick-a-Pic v1 dataset, PRDP achieves superior generation quality on a diverse set of complex, unseen prompts whereas RL-based methods completely fail.