Задача персонализации text-to-image модели состоит в донастройке предобученной нейронной сети так, чтобы она могла генерировать изображения заданного объекта в выбранных сценах. Несмотря на то, что подходы к решению этой задачи существуют, для их применения в высоконагруженных системах необходимо разобраться с рядом проблем: большим временем дообучения, высокими требованиями к видеопамяти, неспособностью точно захватывать детали целевого объекта и др.

Меня зовут Сергей Михайлин. Я разработчик группы машинного обучения в ОК. В данной статье я привожу обзор современных подходов к персонализации text-to-image моделей на базе открытой архитектуры Stable Diffision, делюсь техническими подробностями каждого подхода и анализирую их применимость в реальных высоконагруженных системах. Также, на основании собственных экспериментов по персонализации text-to-image моделей, я дам список возникающих при решении этой задачи проблем и перспективных способов их решения.

Введение

За счет применения механизма диффузии большие text-to-image модели [1, 2, 3] совершили резкий скачок в развитии, чем спровоцировали революцию в создании нового контента. Модели способны генерировать изображения в разных стилях и с разными объектами, основываясь на описании. Естественным образом возникла задача персонализации таких моделей, цель которой – донастройка модели для генерации изображений, на которых представлены объекты, которых нет в обучающей выборке исходной модели. Например, можно заставить нейронную сеть сгенерировать фотографию конкретного человека в костюме космонавта. Чтобы это сделать, необходимо дообучить модель на фотографиях этого человека.  Но на практике количество фотографий целевых объектов невелико, поэтому провести полноценное дообучение невозможно.

Я рассмотрю наиболее известные подходы к персонализации, а также их улучшенные версии: Textual Inversion [4], DreamBooth [5], Low-Rank Adaptation (LoRA) [6], Encoder for Tuning (E4T) [7], HyperDreamBooth [8], чтобы определить применимость каждого метода на практике, их плюсы и минусы, а также составить общий список известных проблем и способов их решения.

Диффузионный метод генерации изображений

Начнем с диффузионного подхода к генерации изображений и модели Stable Diffusion, персонализацию которой реализуют разные методы.

Диффузионные вероятностные модели

Диффузионные нейронные сети D_θ – это вероятностные генеративные модели, которые приобрели популярность благодаря возможности генерировать качественные изображения. Они учатся приближать распределение исходных данных, итеративно очищая изображения от гауссовского шума. Для этого на исходное изображение x постепенно накладывается шум, после чего за T шагов происходит обратный процесс очистки:

Исходное изображение постепенно зашумляется x→z_x,t согласно планировщику шума (например, DDIM [9]), управляющим воздействием на который является случайная величина t∼U(0,T). При этом, чем больше t – тем больше наложено шума:

Во время генерации модель сэмплирует случайный вектор z_t∼N(0,I) и, итеративно “очищая” его от шума, производит новое изображение.

Stable Diffusion

Архитектура Stable Diffusion – это реализация диффузионной модели с открытым исходным кодом и предобученными весами. Она состоит из трех блоков: Pixel Space, Latent Space, Conditioning. Рассмотрим каждый из них.

• Pixel Space. На вход модели во время обучения поступает целевое изображение x большого размера (512×512). Оперировать в таком большом пространстве вычислительно дорого, поэтому исходное изображение сжимается кодировщиком (encoder) до меньшего размера (64×64) в карту признаков, в пространстве которой происходят все последующие манипуляции. Для получения финального результата генерации применяется декодер D, который переводит сжатое представление в исходное x: 64×64→512×512.

• Latent Space. Сжатое представление изображения зашумляется за T последовательных шагов. После этого модель UNet [10] итеративно удаляет этот шум, пытаясь восстановить из полного шума исходное изображение. Для этого ей сообщается текущая степень зашумления timestamp и вспомогательное описание prompt исходного изображения. На стадии генерации изображения шум проходит через одну и ту же модель T раз. Результатом постепенной очистки “шума от шума” является финальное изображение.

• Conditioning. Блок кодирования вспомогательной информации. Text-To-Image подход использует для этого текст, который далее передается в cross-attention [11] слои UNet модели для поиска наложенного шума. Здесь может кодироваться любая информация – главное, чтобы она была связана с исходным изображением.

Textual Inversion

Единственный вход, с помощью которого можно управлять генерацией – это текстовые данные (prompt). Наивный подход заключается в том, чтобы очень качественно описать объект, который нас интересует. Например, “Высокий мужчина спортивного телосложения, кудрявые волосы, европейской внешности, ...”. Проблема такого подхода в том, что описать до мелких деталей каждый объект почти невозможно. Может понадобиться очень длинное описание, но архитектурой Stable Diffusion оно ограничено 77 токенами. Кроме того, некоторое количество токенов нужно, чтобы описать сцену, в которой хотим видеть объект.

Поэтому предлагается [4] свести описание объекта до минимально возможного размера, но сохранить его максимально точным. Эту идею можно реализовать, найдя векторное представление (embedding) в существующем текстовом пространстве, которое максимально точно соответствует целевому объекту.

Для поиска такого представления требуется несколько предварительных шагов:

1. Выберем некоторую последовательность символов (например, [tag]) и заведем для нее отдельный токен в Tokenizer. Этот токен будет соответствовать целевому объекту. Важно, чтобы он не был представлен в словаре модели.

2. Для нового токена создадим embedding в пространстве существующего Text Encoder. Этот вектор будет хранить выученное описание целевого объекта.

Далее запускается несколько итераций обучения на изображениях целевого объекта. При этом в Text Encoder подается фраза, содержащая созданный токен, например, a photo of [tag]. Чтобы предотвратить переобучение на конкретную фразу, рекомендуется использовать несколько похожих по смыслу, но разных предложений: a closeup photo of [tag], a picture of [tag] и так далее. Важно следить, чтобы веса остальных слоев и векторы других слов не менялись, так как обучается только один embedding, соответствующий объекту генерации.

В результате обучения получается искомый вектор, который кодирует конкретный объект. Теперь для генерации можно добавить выбранный тег в тестовое описание и получить изображение с желаемым объектом. Описание желательно составлять в следующей форме: a photo of [tag] <описание сцены>.

К достоинствам textual inversion можно отнести простоту подхода и низкий шанс переобучения. Размер обучаемых весов мал, поэтому для дообучения достаточно 8 гигабайт видеопамяти. Кроме того, можно обучать сразу несколько тегов для разных объектов. Недостатки этого подхода в том, что он требует много времени для дообучения и часто “теряет” детали целевого объекта на сгенерированном изображении.

DreamBooth

Основная проблема подхода Textual Inversion – плохая сохраняемость деталей. Возможная причина в том, что полученное текстовое пространство недостаточно гибкое и вместимое для описания целевого объекта. Так происходит из-за того, что обучается единственное представление целевого объекта, а представления всех остальных токенов остаются замороженными.

Чтобы расширить пространство до необходимого размера, авторы подхода DreamBooth [5] предлагают обучать сразу всю сеть, за исключением части, которая кодирует входное изображение, и части, которая декодирует выход UNet. При этом авторы, так же как и в Textual Inversion, используют конкретную форму текстового описания объекта: a photo of [tag] [class], где [tag] – последовательность символов, обозначающая объект, [class] – название класса, описывающего наш объект, например, для человека – person, для кошки – cat. Добавленное значение [class] “подсказывает” сети тип генерируемого объекта – так быстрее достигается сходимость сети и изображения получаются более качественными.

Важно, какой тег используется для целевого объекта. Если попытаться обучаться на объект, описывая его существующим в словаре модели токеном, то изменить предобученное значение вектора может быть сложно. Причем, чем “сильнее” слово (эмоциональнее, значимее), тем больше будет проблем. Так, используя для целевого объекта токен blue (a photo of blue person), можно получить нужный объект, но синего цвета. Поэтому значение тега важно выбирать аккуратно: можно попробовать инициализацию нового вектора – как в Textual Inversion.

Чтобы исключить переобучение новой модели D_{θ ̂}, вводится Prior Preservation Loss, который предотвращает катастрофическое забывание [12]. Идея подхода в том, чтобы сгенерировать исходной моделью несколько изображений, соответствующих конкретному описанию, и при дообучении требовать у модели восстанавливать эти изображения как можно точнее:

Формально процедура обучения в случае DreamBoth выглядит таким образом:

1. Берем исходную Stable Diffusion с предобученными весами D_θ.

2. Замораживаем все веса.

3. Выбираем любое (из уже существующих в словаре) подходящее слово, обозначающее класс целевого объекта.

4. Подаем фразу c_p вида a photo of [class] в нашу сеть и запоминаем сгенерированные результаты x_p.

5. Размораживаем все веса.

6. Подаем на вход сети изображение объекта x_t c описанием c_t вида a photo of [tag] [class] и рассчитываем значение loss-функции.

7. Подаем на вход сети данные из п.4 (x_p, c_p) и рассчитываем значение loss-функции.

8. Используя уравнение (1), суммируем значения функции потерь из п.6 и п.7, рассчитываем градиенты и обновляем веса.

Повторяем шаги 6-8 до сходимости сети. На практике такое дообучение требует 1300 итераций.

Главным достоинством метода является высокоточный перенос объекта в генерируемое изображение. Среди недостатков: опасность переобучить веса, долгое время дообучения модели, большой размер хранимых весов, высокие требования к свободной видеопамяти (надо до 30 Гб).

Low-Rank Adaptation

В описанных выше подходах есть принципиальное “противостояние” – выбор между способностью переноса деталей и размером обновляемых весов модели. Подход LoRA (Low Rank Adaptation) [6] предлагает компромисс –– низкий размер весов, но некоторые детали могут быть потеряны.

Вместо того, чтобы модифицировать описание генерируемого изображения, предлагается обучить матрицу поправок для cross- и self-attention весов модели UNet. Эта поправка будет “сдвигать” промежуточные карты признаков для генерации необходимого изображения. Изначальные веса матрицы W₀ замораживаются, а в качестве поправки W_δ обучаются две матрицы меньшего ранга, перемножение которых дает матрицу размерности, как у W₀. Выходы слоев, полученных с помощью параметров W₀ и W_δ, складываются.

Формально решение выглядит так:

Таким образом, генерация менее требовательна к видеопамяти. Из-за того, что матрицы A и B низкого ранга – в памяти нужно держать гораздо меньше информации о градиентах. Поскольку поправки к изначальной матрице весов W₀ применяются простым суммированием, при правильной реализации заменять W_δ можно “на ходу”. Это может быть полезно в системах, где необходим механизм переключения дообученных моделей, соответствующих разным целевым объектам.

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

К плюсам метода LoRA можно отнести низкие требования к видеопамяти, низкий размер хранимых весов, возможность обучения для наложения различных эффектов на изображение и совместимость с другими методами дообучения. Недостатками подхода является неудовлетворительная переносимость деталей целевого объекта и продолжительное время дообучения (около 20 минут на каждый целевой объект).

Encoder for Tuning

Авторы оригинальной Textual Inversion предложили улучшенный метод персонализации диффузии – Encoder for Tuning (E4T) [7]. Говоря о качестве, стоит уточнить – официальная реализация не была опубликована на момент написания статьи, но результаты, представленные в статье, превосходят предыдущие подходы в способности сохранения деталей изначального объекта.

В этом подходе генерация состоит из двух этапов:

1. Манипуляции с входными данными.

2. Обучение сдвигов весов.

Манипуляции с входными данными. Авторы не просто обучают модель на конкретных изображениях целевого объекта, а предлагают подход, который сможет работать для любого объекта с минимальным дообучением. Для этого представление объекта специальным образом готовится и подается в модель вместе с текстовым представлением описания изображения.

Создание вектора объекта происходит в Word Encoder – дополнительном новом модуле сети – и подразумевает несколько этапов:

1. Обучение на большом датасете “общего представления”, которое описывает класс целевого объекта.

2. Извлечение признаков целевого объекта и преобразование к размерности “общего представления”.

3. Суммирование значений из п.1 и п.2 с весом λ, для реализации необходимого сдвига от общего описания к конкретному.

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

1. Фото целевого объекта пропускается через энкодер OpenCLIP ViT-H/14 [13]. Сохраняются промежуточные карты признаков каждого второго слоя.

2. Эти признаки подаются в линейные слои.

3. Зашумленное изображение объекта пропускается через кодировщик UNet Stable Diffusion.

4. Признаки из п.2 и п.3 конкатенируются и пропускаются через финальный обучаемый линейный слой.

Таким образом, авторы комбинируют представления двух сетей –– CLIP и UNet –– с целью сохранения семантики объекта (общий вид, национальность, раса, цвет волос) и его “структурных” свойств (форма волос, шрамы, расположение бровей, глаз, губ и другие отличительные черты). Полученный вектор подается в Stable Diffusion вместе с остальным текстовым описанием по аналогии с Textual Inversion.

Обучение сдвигов весов. Авторы провели 50 обучений Stable Diffusion на разных входных данных и заметили, что сильнее всего меняются матрицы W_Q, W_K, W_O в cross-attention слоях. Поэтому они предлагают дообучать только эти части сети для исключения рисков переобучения и быстрой сходимости. Предлагается обучить “сдвиг” W_δ для каждой из матриц, который вычисляется следующим образом:

Генерация происходит в три шага:

1. Предобучение сдвига весов на большом наборе данных “общего представления” и word encoder. Выполняется заранее, до генерации.

2. Дообучение word encoder. Выполняется в момент генерации. Необходимо всего 3-5 шагов.

3. Генерация с полученными весами.

Плюсами метода является низкое время дообучения (∼1 минута) и лучшие заявленные результаты. Среди недостатков стоит отметить отсутствие официальной реализации, высокие требования к видеопамяти (около 50 Гб) и необходимость предварительного обучения полноценной модели Word Encoder.

HyperDreamBooth

В новой статье [8] об усовершенствованном подходе DreamBooth авторы решили проблему долгого дообучения на конкретном объекте и накладных расходов на размер весов. Архитектура состоит из трех компонентов: LiDB (Lightweight DreamBooth), HyperNetwork и Relaxed Fine-tuning.

Lightweight DreamBooth

При обучении DreamBooth приходится сохранять большой объем весов: от 1 до 3 Гб на каждый объект. Авторы решают эту проблему с помощью модифицированного подхода LoRA. Так, изначально разложенная матрица весов W=AB, разбивается на еще более “маленькие части”:

При этом размерности этих матриц определяются гиперпараметрами a, b, r:

Параметры A_train и B_train обучаются. Параметры A_aux и B_aux инициализируются случайными значениями в виде полуортогональных матриц и замораживаются. В силу свойств полуортогональных матриц норма матриц A и B зависит только от обучаемых частей A_train и B_train. Cледовательно, размер значений входной карты признаков меняется только в зависимости от обучаемых частей.

Рассмотрим на примере, за счет чего получается сократить вес. Предположим, изначальная матрица W∈R^5×5 содержит 25 параметров. Тогда с помощью подхода LoRA можно сократить количество обучаемых весов до 10 (все веса в A∈R^5×1 и B∈R^1×5 являются обучаемыми):

А теперь посмотрим, сколько обучаемых весов получается в рассматриваемом подходе. Пусть обучаемые матрицы имеют размерность A_train∈R^3×1,B_train∈R^1×3. Веса остальных матриц замораживаются:

В итоге, в примере предлагаемый подход требует 6 обучаемых параметров против 10 для LoRa и 25 для полной матрицы. В своей модели авторы используют a=100,b=50,r=1: в такой конфигурации веса обучаемых параметров требуют 120 килобайт памяти. Это позволяет обучать и хранить необходимые веса для нескольких целевых объектов.

HyperNetwork

Применим LiDB подход для всех self- и cross-attention слоев и обозначим параметры как θ ̂. Введем дополнительную модель-гиперсеть, которая будет предсказывать A_train и B_train часть этих весов: θ ̂=H(x). Чтобы такая модель могла выдавать качественные веса независимо от входной фотографии, обучим ее на специфичном для домена наборе изображений – например, CelebAHQ [14]. Авторы статьи утверждают, что для этого достаточно порядка 15000 уникальных лиц.

Архитектура HyperNetwork реализована следующим образом:

1. Сначала кодируется информация о лице с помощью Visual transformer Encoder (ViT-H) [13]. Этот шаг позволяет получить представление лица (face embedding).

2. Полученное представление конкатенируется с набором весов, которые изначально инициализированы нулями.

3. Результат пропускается через Transormer Decoder [11] рекурсивно K раз.

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

При обучении HyperNetwork принимает на вход картинку и подставляет предсказанные веса в слои внимания основной модели. При этом Text Encoder принимает фразу вида a [tag] face. Подсчитывается функция потерь обычной Stable Diffusion с дополнительным членом регуляризации, который не позволяет предсказанным параметрам сильно отклоняться от оригинальных значений:

где α и β – гиперпараметры, а θ – оригинальные веса внимания.

При таком подходе модель учится рисовать исходное изображение только за счет изменяющихся параметров θ ̂. Обучение сети происходит один раз и навсегда.

Relaxed Fine-tuning

Авторы заметили, что с HyperNetwork захватывается общая семантика объекта. Для человека это пол, форма и цвет волос, цвет кожи и т.д.. Одновременно с этим, мелкие черты лица теряются – это объясняется тем, что снижение размерности обучаемых матриц A_train,B_train до r=1 приводит к потере выразительной силы модели.

Для выхода из ситуации авторы предлагают интерпретировать и использовать выход гиперсети, как инициализацию слоев внимания с последующим дополнительным дообучением с помощью классической версии LoRA, но уже с r>1. Причем достаточно 40 итераций (20-30 секунд) дообучения для получения качественных результатов. Это на порядок быстрее всех предыдущих методов персонализации.

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

Сравнение методов

Проведем наглядное сравнение результатов трех известных подходов: Textual Inversion, DreamBooth, LoRA и их комбинаций.

Генерация изображений с помощью подходов Encoder for Tuning и HyperDreamBooth на момент написания статьи невозможна из-за отсутствия реализации в открытом доступе. Заметно, что наилучшие результаты по похожести на изначального человека демонстрирует метод DreamBooth и его комбинации с другими подходами.

В таблице дано сравнение времени дообучения, необходимого объема видеопамяти и размера весов моделей, описанных в статье. Textual Inversion в комбинации с другими подходами подразумевает последовательный запуск, для таких случаев указано суммарное время обучения и максимальная затрачиваемая видеопамять. Видно, что подход LoRA позволяет ускорить обучение при минимальном объеме видеопамяти.

Наилучшей по похожести, скорости обучения и потреблению видеопамяти является комбинация подходов DreamBooth и LoRA. Добавление Textual Inversion улучшает качество результата, но заметно снижает скорость дообучения.

Выявленные проблемы и их решения

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

Накладные расходы на хранение весов после дообучения

Решением может служить представление объекта в виде вектора в существующем текстовом пространстве вместе с разложением промежуточных матриц весов на матрицы меньшего размера.

Проблема переобучения

Большинство подходов стараются не менять изначальную архитектуру Stable Diffusion и минимально изменять существующие веса – поскольку изменение существующих параметров может негативно сказываться на обобщающей способности модели. Управление генерацией происходит за счет внедрения дополнительных весов, которые вносят правки в промежуточные карты признаков.

DreamBooth в этом плане сильно выделяется, поскольку меняет веса двух из трех основных частей модели – UNet и Text Encoder. Но риск переобучения гораздо выше, чем у других подходов.

Долгое дообучение на каждом объекте

На момент написания статьи генерация персонализированного изображения требует дообучения сети, которое занимает продолжительное время. Создание гиперсети, которая умеет представлять целевой объект в существующем пространстве Stable Diffusion, позволяет минимально дообучать существующую архитектуру. За счет этого можно добиться снижения времени перед генерацией.

Высокие требования к объему видеопамяти

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

Здесь я не рассматриваю другие способы снижения вычислительной сложности, такие как квантизация, кэширование неизменяемых этапов, использование xformers [15] и др., так как эти способы не являются специфичными для задачи персонализации диффузии.

Плохое сохранение деталей объектов

Методы, реализующие какую-либо модификацию только в одном из нескольких модулей, приводят к плохому захвату деталей целевого объекта. Комплексное решение в виде внесения правок в текстовое пространство и в пространство модели UNet позволяет добиться лучших результатов генерации.

Зашумление функции потерь

В наших экспериментах значение функции потерь не показало себя, как хороший критерий для определения сходимости модели. Большое количество случайных процессов делают функцию потерь похожей на “пилу”, по которой невозможно принять решение о необходимости остановить обучение. Со спецификой проблемы и возможным ее решением можно ознакомиться в статье [16].

Заключение

Несмотря на то, что направление бурно развивается, задача персонализации изображений пока не решена. Ни одна из публично доступных моделей не может быть использована в высоконагруженном сервисе. Я описал главные проблемы, возникающие на этом пути, и обозначил перспективные идеи по их устранению. Решение этих проблем в будущем позволит с высокой скоростью генерировать изображения с желаемыми объектами в нужном контексте на недорогих видеокартах.

Список использованных источников

[1] Saharia C. et al. Photorealistic text-to-image diffusion models with deep language understanding //Advances in Neural Information Processing Systems. – 2022. – Т. 35. – С. 36479-36494.

[2] Ramesh A. et al. Hierarchical text-conditional image generation with clip latents //arXiv preprint arXiv:2204.06125. – 2022.

[3] Rombach R. et al. High-resolution image synthesis with latent diffusion models //Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. – 2022. – С. 10684-10695.

[4] Gal R. et al. An image is worth one word: Personalizing text-to-image generation using textual inversion //arXiv preprint arXiv:2208.01618. – 2022

[5] Ruiz N. et al. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. – 2023. – С.22500-22510.

[6] Hu E. J. et al. Lora: Low-rank adaptation of large language models //arXiv preprint arXiv:2106.09685. – 2021.

[7] Gal R. et al. Designing an encoder for fast personalization of text-to-image models //arXiv preprint arXiv:2302.12228. – 2023.

[8] Ruiz N. et al. HyperDreamBooth: HyperNetworks for Fast Personalization of Text-to-Image Models //arXiv preprint arXiv:2307.06949. – 2023.

[9] Song J., Meng C., Ermon S. Denoising diffusion implicit models //arXiv preprint arXiv:2010.02502. – 2020.

[10] Ronneberger O., Fischer P., Brox T. U-net: Convolutional networks for biomedical image segmentation //Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. – Springer International Publishing, 2015. – С. 234-241.

[11] Vaswani A. et al. Attention is all you need //Advances in neural information processing systems. – 2017. – Т. 30.

[12] Kirkpatrick J. et al. Overcoming catastrophic forgetting in neural networks //Proceedings of the national academy of sciences. – 2017. – Т. 114. – №. 13. – С. 3521-3526.

[13] Dosovitskiy A. et al. An image is worth 16x16 words: Transformers for image recognition at scale //arXiv preprint arXiv:2010.11929. – 2020.

[14] Karras T. et al. Progressive growing of gans for improved quality, stability, and variation //arXiv preprint arXiv:1710.10196. – 2017.

[15] xFormers - Toolbox to Accelerate Research on Transformers. https://github.com/facebookresearch/xformers

[16] Voronov A. et al. Is This Loss Informative? Speeding Up Textual Inversion with Deterministic Objective Evaluation //arXiv preprint arXiv:2302.04841. – 2023