Персонализированная генерация: Textual Inversion и DreamBooth
Введение
Персонализированная генерация — подход в генеративном моделировании, который позволяет «встроить» новые концепции в обученную модель. Он помогает адаптировать предобученные генеративные модели (Stable Diffusion или другие диффузионные модели) для создания изображений на основе пользовательских концепций — например, изображений определённого человека, объекта или стиля искусства.
Проблемы наивных подходов к персонализированной генерации:
- Обучение диффузионной модели с нуля для каждого нового концепта, например, для одного объекта или стиля — это вычислительно дорого.
- Наивное дообучение существующей модели на новый концепт может привести к проблеме «забывания» ранее полученных знаний. То есть модель, обученная рисовать животных и пейзажи, может после дообучения на один конкретный объект перестать адекватно генерировать другие объекты.
Таким образом, главная задача — найти способ, который позволит обучить модель новым концептам и избежать проблемы «забывания», снизив требования к вычислительным ресурсам.
В обзоре рассмотрим два последовательно появившихся подхода для решения этой задачи — Textual Inversion и Dreambooth.
Textual Inversion
Основная идея
При создании изображений на основе текстовых описаний применяются специальные энкодеры для обработки текста. В них слова сначала преобразуются в токены — элементы из фиксированного словаря.
Затем эти токены переводятся в векторные представления (эмбеддинги), которые передаются в диффузионную модель для дальнейшей работы.
Для обозначения нового, уникального концепта (например, конкретного объекта, домашнего животного или художественного стиля), который отсутствует в словаре модели, вводится специальное псевдослово, обозначаемое как
В дальнейшем для генерации данное псевдослово трактуется как обычное слово и может комбинироваться с другими для text-to-image генерации. Важно отметить: этот процесс не затрагивает саму генеративную модель, её веса остаются замороженными.
Поиск эмбеддинга для псевдослова
Обычно для элементов из словаря (токенов) уже есть соответствующий эмбеддинг. Теперь без замены текстового энкодера мы должны научить его новому концепту, то есть добавить в словарь ещё один токен
Для этого решается оптимизационная задача:
- Сначала эмбеддинг инициализируется либо случайно, либо эмбеддингом близкого слова (например, для картинки выше — «кот»).
- Затем берётся небольшой датасет из 3-5 изображений концепта, который нужно закодировать, и каждому изображению ставится шаблонный промпт, например, «A photo of \(S^*\)».
- Текстовый энкодер преобразует этот промпт в последовательность эмбеддингов, где новому токену \(S^*\) соответствует эмбеддинг \(v^*\). В последовательности все эмбеддинги, кроме \(v^*\), фиксированы, а \(v^*\) будет обучаться.
- Пара из зашумлённого изображения \(z_t\) и эмбеддинг текста \(c_{\theta}(y)\) передаются в замороженную диффузионную модель, а затем минимизируется стандартный диффузионный лосс для поиска оптимального \(v^*\):
\(v_* = \arg\min_{v} \mathbb{E}{z,\, y,\, \epsilon,\, t} \left[ \left\| \epsilon — \epsilon\theta(z_t, t, c_\theta(y)) \right\|_2^2 \right]\)
где \(z \sim \mathcal{E}(x)\) — латент, получаемый путём кодирования входного изображения \(x\), \(\epsilon \sim \mathcal{N}(0,1)\) — шум, \(\epsilon_\theta(z_t, t, c_\theta(y))\) — предсказание диффузионной модели на основе зашумлённого латента \(z_t\), шага зашумления t и закодированного промпта \(c_\theta(y)\).
Важно отметить: в этом сетапе обучаетcя только векторный эмбеддинг \(v^*\), а дифузионная модель и текстовый энкодер заморожены. Обученный эмбеддинг \(v^*\) сохраняется и используется в процессе инференса.
Весь пайплайн представлен на рисунке ниже:
Метрики
В качестве основных метрик авторы приводят результаты оценки CLIP и user study. CLIP оценивает косинусное расстояние с таргетным концептом (Image similarity) и с желаемым промптом (text similarity).
Участникам user study задавали два вопроса:
- Оценка Визуальной Точности. Пользователям показывали четыре изображения, использованных для обучения конкретного концепта. Им нужно было ранжировать результаты пяти различных моделей по тому, насколько они похожи на исходные обучающие изображения.
- Оценка Соответствия Тексту. Пользователям представляли текстовый запрос, задающий контекст (например, «Фотография на пляже»), и просили ранжировать результаты тех же моделей по степени соответствия заданному текстовому описанию.
Результаты приведены на рисунке ниже:
Визуализация результатов
Примеры работы подхода представлены авторами статьи на рисунке ниже. Хоть и не идеально, но подход сохраняет предложенный концепт.
Примеры работы метода для задач стилизации:
Также подход позволяет комбинировать несколько объектов вместе:
DreamBooth
DreamBooth является улучшением Textual Inversion и до сих пор активно используется в задачах.
Textual Inversion ищет новый вектор в пространстве текста, а DreamBooth идёт путём дообучения самой генеративной модели. Однако авторы предлагают подход, позволяющий уменьшать вероятность «забывания» моделью того, что она выучила ранее.
Обучение
Как и в предыдущем сэтапе, у нас есть небольшой набор 3-5 изображений уникального концепта, и мы также хотим ввести в промт уникальное псевдослово, например, a photo of a [V] dog.
Авторы говорят внедрять знание о новом концепте через дообучение самой диффузионной модели, однако:
- возможно забывание других объектов;
- нужно много вычислительных ресурсов.
Для решения этих проблем они предлагают prior preservation loss.
В подходе используется два экземпляра одной диффузионной модели: первая дообучается на новый концепт, а другая замораживается. На каждой итерации с помощью замороженной предобученной модели генерируются примеры, соответствующие только классу (например, a dog), без уникального псевдослова. Обучаемая модель теперь одновременно учится генерировать изображения с заданным концептом, используя промпт с псевдословом (a [V] dog) и изображения из предобученной замороженной модели с обычным промптом (a dog).
На рисунке ниже представлен пайплайн обучения:
Обозначим предобученную диффузионную text-to-image модель как \(\hat{\mathbf{x}}_\theta\), которую будем дообучать, а \(\hat{x}\) — как её замороженную версию. Общий лосс состоит из двух частей.
- Лосс для обучения нового концепта
где \(с\) — промпт с псевдословом, а \(x\) — изображение из обучающего датасета.
- Лосс для сохранения знаний исходной диффузионной модели
\(\mathbb{E}{\boldsymbol{\epsilon}’, t} w{t’} \left\| \hat{\mathbf{x}}\theta(\alpha{t’} \mathbf{x}{\mathrm{pr}} + \sigma{t’} \boldsymbol{\epsilon}’, \mathbf{c}{\mathrm{pr}}) — \mathbf{x}{\mathrm{pr}} \right\|_2^2\) ,
где \(x_{pr} = \hat{x}(z_{t_1}, x_{pr})\) — изображение исходной диффузионной модели из случайного шума \(z_{t_1} \sim \mathcal{N}(0,I)\) с использованием промпта без псевдослова \(c_{pr}\).
В качестве итогового лосса берётся взвешенная сумма.
\(\mathbb{E}{\mathbf{x}, \mathbf{c}, \boldsymbol{\epsilon}, \boldsymbol{\epsilon}’, t} \left[ w_t \left\| \hat{\mathbf{x}}\theta(\alpha_t \mathbf{x} + \sigma_t \boldsymbol{\epsilon}, \mathbf{c}) — \mathbf{x} \right\|2^2 + \lambda w{t’} \left\| \hat{\mathbf{x}}\theta(\alpha{t’} \mathbf{x}{\mathrm{pr}} + \sigma_{t’} \boldsymbol{\epsilon}’, \mathbf{c}{\mathrm{pr}}) — \mathbf{x}{\mathrm{pr}} \right\|_2^2 \right],\)
По словам авторов, достаточно всего 1000 итераций с \(\lambda = 1\), чтобы модель сошлась для Imagen и Stable Diffusion.
Метрики
Для валидации модели авторы собрали датасет из 30 объектов, разделив их на две категории: объекты (21) и живые существа, животные (9).
Примеры изображений из датасета:
Для валидации использовали несколько основных метрик:
- CLIP-I и DINO — для сохранения деталей нового концепта. CLIP-I считает косинусное расстояние между эмбеддингами CLIP-сгенерированной и реальной картинками, DINO — тоже косинусное расстояние, но между фичами s/16 DINO-сгенерированной и реальной картинками.
- CLIP-T — для проверки соответствия промпту, косинусное расстояние между эмбеддингами CLIP для картинки и текста.
Также авторы провели user-study. Каждому участнику показали набор реальных изображений (обучающий сет) для определённого объекта и два сгенерированных изображения — по одному от каждого метода (Textual Inversion и DreamBooth). Пользователи отвечали на два вопроса:
- Точность воспроизведения объекта (Subject Fidelity): «Какое из двух изображений лучше воспроизводит идентичность объекта (например, тип предмета и детали) из исходных примеров?»;
- Соответствие промпту (Prompt Fidelity): «Какое из двух изображений лучше всего описывается заданным текстовым запросом?».
Это соотносится с user-study, проводимым для работы Textual Inversion.
Основные метрики приведены на рисунке 15:
Можно заметить, что метод DreamBooth значительно превосходит Textual Inversion как по метрикам, так и по предпочтениям пользователей, и лучше сохраняет соответствие заданному концепту и таргетному промпту.
Кроме того, авторы говорят, что использование prior preservation loss очень важно в контексте сохранения разнообразия генерации и избежания «забывания».
Для демонстрации авторы проводят обучение модели без лосса и с ним и сравнивают несколько основных метрик:
- DINO, CLIP-I, CLIP-T;
- PRES (prior preservation metric) — среднее попарное сходство DINO-эмбеддингов между двумя наборами изображений: изображения, сгенерированные по запросу класса (например, «собака»), и реальные изображения уникального объекта. Если среднее сходство оказывалось высоким, это сигнализировало о коллапсе: модель, потеряв общие знания, начинала генерировать случайные объекты класса, которые были слишком похожи на уникальный обучаемый объект. Чем выше значение этой метрики, тем сильнее языковой дрейф, что подтверждает необходимость введения Prior Preservation Loss для сохранения разнообразия и общих знаний модели о классе;
- DIV (diversity metric) — среднее попарное LPIPS расстояние. Высокое значение DIV говорит о том, что модель генерирует визуально разнообразные изображения при одном и том же промпте.
Также влияние этого лосса представлено на рисунке 17:
Визуализация результатов
Также можно визуально оценить качество работы предложенной модели в различных задачах.
- Генерация объектов в разном окружении
- Art Rendition
- Property Modification
- Accessorization
И напоследок:
Заключение
В статье мы рассмотрели два подхода для персонализированной генерации:
- Textual Inversion — добавление нового концепта в словарь и обучение специального эмбеддинга для этого токена.
- DreamBooth — дообучение самой диффузионной модели с помощью Prior Preservation Loss.
На практике сегодня часто используются методы, сочетающие в себе оба подхода: сначала выучиваем эмбеддинг для нового псевдослова, а затем уже проводим дообучение всей модели.
Кроме того, необязательно дообучать большую модель полностью — можно дообучать LoRA-добавки к весам, оставляя большую часть параметров нетронутыми.
После рассмотренных работ вышло много различных подходов для персонализированной генерации, однако эти два метода всё ещё активно используются в работе.
Более современные подходы можно найти в репозитории на гитхабе 🙂
Полезные ссылки
Generative Computer Vision
