Согласование моделей диффузии текст-изображение с обратным распространением вознаграждения

Aligning Text-to-Image Diffusion Models with Reward Backpropagation

October 5, 2023

Авторы: Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki

cs.AI

Аннотация

Модели диффузии для генерации изображений по тексту недавно вышли на передний план в области создания изображений, благодаря использованию очень крупномасштабных наборов данных для обучения без учителя или с слабым контролем. Из-за их обучения без учителя управление их поведением в последующих задачах, таких как максимизация воспринимаемого человеком качества изображения, соответствия изображения и текста или этичной генерации изображений, является сложной задачей. Недавние работы адаптируют модели диффузии к целевым функциям с использованием классического обучения с подкреплением, известного высокой дисперсией оценок градиента. В данной статье мы предлагаем AlignProp — метод, который выравнивает модели диффузии по целевым функциям с помощью сквозного обратного распространения градиента награды через процесс удаления шума. Хотя наивная реализация такого обратного распространения потребовала бы непомерных ресурсов памяти для хранения частных производных современных моделей генерации изображений по тексту, AlignProp адаптирует низкоранговые модули весов и использует контрольные точки градиента, чтобы сделать использование памяти приемлемым. Мы тестируем AlignProp в адаптации моделей диффузии к различным целям, таким как семантическое соответствие изображения и текста, эстетика, сжимаемость и управляемость количества объектов, а также их комбинации. Мы показываем, что AlignProp достигает более высоких значений награды за меньшее количество шагов обучения по сравнению с альтернативами, оставаясь концептуально более простым, что делает его прямым выбором для оптимизации моделей диффузии по дифференцируемым целевым функциям. Код и визуализация результатов доступны по адресу https://align-prop.github.io/.

English

Text-to-image diffusion models have recently emerged at the forefront of image generation, powered by very large-scale unsupervised or weakly supervised text-to-image training datasets. Due to their unsupervised training, controlling their behavior in downstream tasks, such as maximizing human-perceived image quality, image-text alignment, or ethical image generation, is difficult. Recent works finetune diffusion models to downstream reward functions using vanilla reinforcement learning, notorious for the high variance of the gradient estimators. In this paper, we propose AlignProp, a method that aligns diffusion models to downstream reward functions using end-to-end backpropagation of the reward gradient through the denoising process. While naive implementation of such backpropagation would require prohibitive memory resources for storing the partial derivatives of modern text-to-image models, AlignProp finetunes low-rank adapter weight modules and uses gradient checkpointing, to render its memory usage viable. We test AlignProp in finetuning diffusion models to various objectives, such as image-text semantic alignment, aesthetics, compressibility and controllability of the number of objects present, as well as their combinations. We show AlignProp achieves higher rewards in fewer training steps than alternatives, while being conceptually simpler, making it a straightforward choice for optimizing diffusion models for differentiable reward functions of interest. Code and Visualization results are available at https://align-prop.github.io/.