Обучение диффузионных моделей с помощью обучения с подкреплением | GPTMain News

function show() { const replay = document.querySelector(‘.ddpo-replay’); replay.style.display = ‘гибкий’; } window.onload = () => { const replay = document.querySelector(‘.ddpo-replay’); replay.addEventListener(‘click’, () => { const video = document.querySelector(‘.ddpo-video’); video.currentTime = 0; video.play(); replay.style.display = ‘none’; }); }

Диффузионные модели недавно стали фактическим стандартом для получения сложных, многомерных результатов. Возможно, вы знаете их за их способность создавать потрясающие произведения искусства с использованием искусственного интеллекта и гиперреалистичные синтетические изображения, но они также добились успеха в других приложениях, таких как разработка лекарств и непрерывный контроль. Ключевая идея диффузионных моделей заключается в итеративном преобразовании случайного шума в образец, например изображение или структуру белка. Обычно это мотивируется проблемой оценки максимального правдоподобия, когда модель обучается генерировать выборки, которые максимально точно соответствуют обучающим данным.

Однако большинство случаев использования моделей диффузии напрямую связаны не с сопоставлением обучающих данных, а с последующей целью. Нам нужно не просто изображение, похожее на существующие изображения, а изображение, имеющее определенный тип внешнего вида; нам нужна не просто физически правдоподобная молекула лекарства, но и максимально эффективная. В этом посте мы покажем, как диффузионные модели могут быть обучены для достижения этих последующих целей напрямую с помощью обучения с подкреплением (RL). Для этого мы настраиваем стабильную диффузию с учетом различных целей, включая сжимаемость изображения, эстетическое качество, воспринимаемое человеком, и быстрое выравнивание изображения. Последняя из этих целей использует обратную связь от большой модели языка видения для улучшения производительности модели при выполнении необычных подсказок, демонстрируя, как мощные модели искусственного интеллекта могут использоваться для улучшения друг друга без участия человека.

Диаграмма, иллюстрирующая цель выравнивания подсказки-изображения. Он использует LLaVA, большую модель языка видения, для оценки сгенерированных изображений.

Оптимизация политики шумоподавления и диффузии

Превращая диффузию в проблему RL, мы делаем только самое основное предположение: учитывая образец (например, изображение), у нас есть доступ к функции вознаграждения, которую мы можем оценить, чтобы сказать нам, насколько «хорош» этот образец. Наша цель состоит в том, чтобы модель диффузии генерировала образцы, которые максимизируют эту функцию вознаграждения.

Модели диффузии обычно обучаются с использованием функции потерь, полученной на основе оценки максимального правдоподобия (MLE), что означает, что им рекомендуется генерировать выборки, которые делают обучающие данные более правдоподобными. В настройке RL у нас больше нет обучающих данных, есть только образцы из модели диффузии и связанные с ними вознаграждения. Один из способов, которым мы все еще можем использовать ту же функцию потерь, мотивированную MLE, — это рассматривать выборки как обучающие данные и включать вознаграждения, взвешивая потери для каждой выборки по ее вознаграждению. Это дает нам алгоритм, который мы называем регрессией, взвешенной по вознаграждению (RWR), в честь существующих алгоритмов из литературы по RL.

Однако у этого подхода есть несколько проблем. Во-первых, RWR не является особенно точным алгоритмом — он лишь приблизительно максимизирует вознаграждение (см. Наир и др., Приложение A). Потери на диффузию, вызванные MLE, также не являются точными и вместо этого рассчитываются с использованием вариационной границы истинной вероятности каждой выборки. Это означает, что RWR максимизирует вознаграждение за счет двух уровней аппроксимации, что, по нашему мнению, существенно ухудшает его производительность.

Мы оцениваем два варианта DDPO и два варианта RWR по трем функциям вознаграждения и обнаруживаем, что DDPO стабильно достигает наилучшей производительности.

Ключевая идея нашего алгоритма, который мы называем оптимизацией политики диффузии шумоподавления (DDPO), заключается в том, что мы можем лучше максимизировать вознаграждение за окончательную выборку, если обратим внимание на всю последовательность шагов шумоподавления, которые привели нас к этому. Для этого мы переосмыслим процесс диффузии как многоэтапный марковский процесс принятия решений (MDP). В терминологии MDP: каждый шаг шумоподавления является действием, и агент получает вознаграждение только на последнем этапе каждой траектории шумоподавления, когда создается окончательная выборка. Эта структура позволяет нам применять множество мощных алгоритмов из литературы по RL, разработанных специально для многошаговых MDP. Вместо использования приблизительной вероятности окончательной выборки эти алгоритмы используют точную вероятность каждого шага шумоподавления, которую чрезвычайно легко вычислить.

Мы решили применить алгоритмы градиента политики из-за их простоты реализации и прошлых успехов в точной настройке языковой модели. Это привело к двум вариантам DDPO: DDPO_{Сан-Франциско}, который использует простую функцию оценки градиента политики, также известную как REINFORCE; и ДДПО_{ЯВЛЯЕТСЯ}, который использует более мощный выборочный оценщик важности. ДДПО_{ЯВЛЯЕТСЯ} — это наш наиболее эффективный алгоритм, и его реализация очень похожа на реализацию проксимальной политики оптимизации (PPO).

Точная настройка стабильной диффузии с использованием DDPO

Для наших основных результатов мы настраиваем Stable Diffusion v1-4 с использованием DDPO._{ЯВЛЯЕТСЯ}. У нас есть четыре задачи, каждая из которых определяется отдельной функцией вознаграждения:

• Сжимаемость: насколько легко сжать изображение с помощью алгоритма JPEG? Наградой является отрицательный размер файла изображения (в КБ) при сохранении в формате JPEG.
• Несжимаемость: Насколько сложно сжать изображение с помощью алгоритма JPEG? Наградой является положительный размер файла изображения (в КБ) при сохранении в формате JPEG.
• Эстетическое качество: насколько эстетично изображение для человеческого глаза? Наградой становятся результаты эстетического предсказателя LAION — нейронной сети, обученной на человеческих предпочтениях.
• Выравнивание подсказки и изображения: насколько хорошо изображение отображает то, что было запрошено в подсказке? Этот немного сложнее: мы передаем изображение в LLaVA, просим его описать изображение, а затем вычисляем сходство между этим описанием и исходным приглашением с помощью BERTScore.

Поскольку Stable Diffusion — это модель преобразования текста в изображение, нам также необходимо выбрать набор подсказок, которые будут выдаваться во время точной настройки. Для первых трех задач используем простые подсказки вида «а(н) [animal]». Для выравнивания изображения-подсказки используем подсказки вида «а(н) [animal] [activity]»где осуществляется деятельность “мыть посуду”, «играть в шахматы»и «езда на велосипеде». Мы обнаружили, что Stable Diffusion часто с трудом создает изображения, соответствующие этим необычным сценариям, оставляя много возможностей для улучшения с помощью точной настройки RL.

Во-первых, мы проиллюстрируем эффективность DDPO на простых наградах (сжимаемость, несжимаемость и эстетическое качество). Все изображения генерируются с одним и тем же случайным начальным числом. В верхнем левом квадранте мы иллюстрируем, что «ванильная» стабильная диффузия генерирует для девяти разных животных; все модели с точной настройкой RL демонстрируют явную качественную разницу. Интересно, что модель эстетического качества (вверху справа) имеет тенденцию к минималистичным черно-белым линейным рисункам, раскрывая те виды изображений, которые эстетический предсказатель LAION считает «более эстетичными».¹

Далее мы продемонстрируем DDPO на более сложной задаче выравнивания изображения-подсказки. Здесь мы показываем несколько снимков процесса обучения: каждая серия из трех изображений показывает образцы для одного и того же быстрого и случайного начального числа с течением времени, при этом первый образец получен из ванильного Stable Diffusion. Интересно, что модель переходит в сторону более мультяшного стиля, что не было намеренно. Мы предполагаем, что это связано с тем, что животные, выполняющие действия, подобные человеческим, с большей вероятностью будут отображаться в данных перед обучением в мультяшном стиле, поэтому модель смещается в сторону этого стиля, чтобы легче согласовываться с подсказкой, используя то, что она уже знает.

Неожиданное обобщение

Было обнаружено неожиданное обобщение при точной настройке больших языковых моделей с помощью RL: например, модели, точно настроенные на следование инструкциям только на английском языке. часто улучшаются на других языках. Мы обнаружили, что то же самое явление происходит с моделями диффузии текста в изображение. Например, наша модель эстетического качества была точно настроена с использованием подсказок, выбранных из списка из 45 распространенных животных. Мы обнаруживаем, что оно распространяется не только на невидимых животных, но и на повседневные предметы.

В нашей модели выравнивания изображений во время обучения использовался один и тот же список из 45 обычных животных и только три вида деятельности. Мы обнаруживаем, что оно распространяется не только на невидимых животных, но и на невидимые действия и даже на новые комбинации того и другого.

Чрезмерная оптимизация

Хорошо известно, что точная настройка функции вознаграждения, особенно изученной, может привести к чрезмерной оптимизации вознаграждения, когда модель использует функцию вознаграждения для достижения высокого вознаграждения бесполезным способом. Наш сеттинг не является исключением: во всех задачах модель в конечном итоге уничтожает любой значимый контент изображения, чтобы максимизировать вознаграждение.

Мы также обнаружили, что LLaVA подвержен типографским атакам: при оптимизации выравнивания по подсказкам вида «[n] животные»DDPO удалось успешно обмануть LLaVA, вместо этого сгенерировав текст, отдаленно напоминающий правильное число.

В настоящее время не существует универсального метода предотвращения переоптимизации, и мы выделяем эту проблему как важное направление будущей работы.

Заключение

Диффузионные модели трудно превзойти, когда дело доходит до получения сложных, многомерных результатов. Однако до сих пор они в основном добились успеха в приложениях, целью которых является изучение закономерностей на основе большого количества данных (например, пар изображений и подписей). Мы нашли способ эффективного обучения моделей диффузии, выходящий за рамки сопоставления с образцом и не требующий каких-либо обучающих данных. Возможности ограничены только качеством и креативностью вашей функции вознаграждения.

То, как мы использовали DDPO в этой работе, вдохновлено недавними успехами точной настройки языковой модели. Модели OpenAI GPT, такие как Stable Diffusion, сначала обучаются на огромных объемах интернет-данных; затем они настраиваются с помощью RL для создания полезных инструментов, таких как ChatGPT. Обычно их функция вознаграждения определяется человеческими предпочтениями, но у других их больше. недавно придумали, как создавать мощных чат-ботов, используя вместо этого функции вознаграждения на основе отзывов ИИ. По сравнению с режимом чат-бота наши эксперименты невелики и ограничены по масштабу. Но, учитывая огромный успех этой парадигмы «предварительная подготовка + точная настройка» в языковом моделировании, кажется, что ее стоит развивать дальше в мире диффузионных моделей. Мы надеемся, что другие смогут использовать нашу работу для улучшения больших моделей диффузии не только для генерации текста в изображение, но и для многих интересных приложений, таких как создание видео, создание музыки, редактирование изображений, синтез белка, робототехника и многое другое.

Более того, парадигма «предварительная подготовка + точная настройка» — не единственный способ использования DDPO. Если у вас есть хорошая функция вознаграждения, ничто не мешает вам тренироваться с RL с самого начала. Хотя эта ситуация еще не исследована, это место, где сильные стороны DDPO могут действительно проявиться. Чистый RL уже давно применяется в самых разных областях: от игр до манипуляций с роботами, ядерного синтеза и проектирования чипов. Добавление к этому мощной выразительности моделей диффузии может вывести существующие приложения RL на новый уровень или даже открыть новые.

Этот пост основан на следующей статье:

• Обучение диффузионных моделей с обучением с подкреплением
  

  Кевин Блэк*, Майкл Дженнер*, Илун Ду, Илья Костриков и Сергей Левин
  
  Препринт arXiv.
  

Если вы хотите узнать больше о DDPO, вы можете просмотреть статью, веб-сайт, исходный код или узнать вес модели на Hugging Face. Если вы хотите использовать DDPO в своем собственном проекте, ознакомьтесь с моей реализацией PyTorch + LoRA, где вы можете точно настроить Stable Diffusion с менее чем 10 ГБ памяти графического процессора!

Если DDPO вдохновляет вашу работу, процитируйте ее, указав:

@misc{black2023ddpo,
      title={Training Diffusion Models with Reinforcement Learning}, 
      author={Kevin Black and Michael Janner and Yilun Du and Ilya Kostrikov and Sergey Levine},
      year={2023},
      eprint={2305.13301},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}