Dream Booth — очень умное дообучение stable diffusion

Исходная публикация с описаниеим метода [Mar 2023] - DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

Что такое stable-diffusion

Stable Diffusion - новаторская модель генерации изображений на основе искусственного интеллекта. Она использует метод диффузии для постепенного улучшения качества сгенерированных изображений.

В основеStable Diffusion лежит архитектура autoencoder. Сначала модель "кодирует" текстовое описание в вектор признаков. Затем декодер поэтапно создает изображение, уточняя его с каждым шагом. Этот процесс напоминает диффузию - плавное распространение частиц.

Что такое дообучение

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

Пробелемы

1. Language drift - модель предварительно обученная на большом корпусе текста и позже точно настроенная для конкретной задачи, постепенно теряет синтаксические и семантические знания языка.

2. Уменьшения разнообразия выходных данных. Модели преобразования текста в изображение естественным образом обеспечивают большое разнообразие выходных данных. При точной настройке мы хотели бы иметь возможность создавать объект в новых ракурсах, позах. Однако из за переобучения вариативность выходных поз объекта сокращается.

Основная идея решения

Авторы статьи решают эту задачу дополнительным лоссом - class-specific prior preservation loss. Вначале обучаемая модель пытаться сгенерировать нашу собаку, и обучается по лоссу с фотографиями из выборки для дообучения. Затем она же генерирует произвольную собаку, и лосс считается с фотографиями, сгенерированными исходной моделью. Таким образом, мы наказываем модель, за то что она забывает абстрактные признаки собак, концентрируясь на обучающей выборке. Более лаконично на схеме:

Или более строго прибавку к лосу можно записать в виде:

Насколько хорошо это работает?

Как видно из картинок ниже, добавление class-specific prior preservation loss очень позитивно сказалось на возможности модели генерировать исходный объект в разнообразных вариациях

Просто сравните с тем что было выше!

Так же модель научилась видоизменять обьект в зависимости от промпта, генерируя полностью новые объекты на основе существующего:

При этом сама модель не потеряла возможности генерировать обычные объекты (произвольных собак, отличных от той, на которой происходило до обучение)

А так же такой подход очень хорошо сработал для задач style-transfer конкретного объекта. Авторы смогли сгенерировать картинки в разных стилях, но с исходной собакой.

LoRA

При обучении модели нет смысла обучать все её веса, достаточно обучения небольших адаптеров. Подробнее о том как это работает написано в моей другой статье на хабре.

Воспроизведение результатов

Я до обучил stable-diffusion-v1-5 на двух собственных датасетах - фотографии кошки (cloody-cat) и декоративного гриба. Обучения проходило на видеокарте A100 (занимая примерно 1/8 видеопамяти) в течении 2х минут. Интерактивная демонстрация работы моей модели на Hugging Face - freQuensy23/diffusion-cloody - там можно потыкать свои запросы бесплатно (за такое им огромный респект)! В таблице ниже показано как со временем обучения меняется результат генерации по промпту "A [V] cat", а более красочно вы можете посмотреть это в анимации в начале этой статьи или тут. Так же я проверил данную систему на генерации этого кота в очень разных условиях и стилях. Дальше будет просто фотогалерея

Прошу прощения за такой спам картинками в конце, чуть менее формально про эту технологию и идеи по ее монетизации описаны в моей телеге
P.S. Если есть ошибка в правописании - сообщите мне через стандартную систему хабра - я исправлю. Мне бывает сложно писать на русском.