Путеводитель для диффузионок. Как заставить нейросети качественно редактировать изображения

Привет, Хабр! Меня зовут Вадим, я — младший научный сотрудник группы Controllable Generative AI лаборатории FusionBrain в AIRI. Последние несколько лет я занимаюсь исследованием генеративных моделей в контексте задачи редактирования фотографий. Мы с командой накопили большую экспертизу в этом и хотели бы поделиться ей.

Совсем недавно мы выложили препринт статьи, которую мы представим на ECCV этой осенью (сама статья, её код, demo на HuggingFace). Там мы предложили метод редактирования реальных изображений с помощью диффузионных моделей, который достигает лучшего среди всех методов компромисса между качеством редактирования и сохранением структуры исходного изображения, а также эффективен с вычислительной точки зрения. В данной статье я хотел бы рассказать о том, почему приходится делать такой выбор, и как мы эту проблему обошли. Приятного чтения!

Введение

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

Основных архитектур не так много: вариацонные автокодировщики (VAE), модели на основе потоков (Flow-based models), генеративно-состязательные сети (GAN) и, наконец, диффузионные модели. Последние демонстрируют высокое качество генерации, благодаря чему их активно используют в text-to-image задачах. Dall-E, Stable Diffusion, Midjourney — одни из самых выдающихся представителей диффузионок.

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

Для каждой архитектуры задача о редактировании решается по-своему. Не так давно мой коллега Денис Бобков выпустил большой и подробный разбор того, как мы это делаем для GAN на примере модификации модели StyleGAN. Здесь же я поделюсь достижениями нашей группы на ниве диффузионных моделей. Но для начала пару слов о том, что это вообще такое.

Диффузия

О принципе работы диффузионных моделей можно догадаться по их названию. Вспомним физику: если мы откроем флакон с духами, запах через какое-то время начнёт распространяться по воздуху. Это происходит из-за броуновского движения пахучих молекул. Хотя это движение хаотичное, частицы в среднем стремятся покинуть область с высокой концентрацией, то есть флакон — происходит диффузия.

Обратите внимание на то, как меняется вероятностное распределение молекул: сначала оно отражает некоторую пространственную структуру (внутренность флакона), но после вскрытия размывается, в пределе сходясь к константе по всему объёму. Это можно представить, как размытие информации в шум.

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

Теперь перенесёмся в компьютерные науки. Представьте, что с самого начала у нас есть какое-то осмысленное изображение. Распределение пикселей при этом демонстрирует некоторую структуру. Затем мы проводим последовательность размытий с помощью гауссовых функций (молекулы в газе также диффундируют «по гауссу»), в результате чего картинка превращается в шум, а распределение становится случайным (как белый шум в телевизоре) — такую цепочку преобразований я буду называть forward траекторией.

С первого взгляда такое преобразование кажется бессмысленным, зачем из хорошей фотографии делать белый шум? Ответ кроется в том, что диффузионная модель способна обращать вспять данный процесс, то есть преобразовывать белый шум малыми шагами до высококачественного изображения (backward траектория).

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

Теперь, получая на вход случайный шум, такая модель начнет поэтапно «извлекать» из чистого шума осмысленные изображения в духе обучающего датасета.

Источник: Vaibhav Singh et al. / learnopencv.com

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

Постановка задачи и математика

Итак, теперь нам бы хотелось решать вот такую задачу: подать на вход модели картинку и текст, в котором сказано, что именно там нужно изменить. Существующие подходы к этой проблеме можно разделить на три основные группы.

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

Следующие две группы методов основываются на корректировке forward и backward траекторий. В одной упор делается на кэшировании промежуточных представлений во время этапа зашумления, далее данные представление во время процесса расшумления подставляются в промежуточные слои нейросети. Такой способ позволяет достаточно дешевым способом сохранить бóльшую часть изображения без дополнительных оптимизаций, откорректировав backward процесс. Главная беда этого подхода — отсутствие универсальности и высокая чувствительность к подбору гиперпараметров, из-за этого данный метод не позволяет получать стабильные результаты (примеры тут и тут).

Во последней же группе методов основная идея заключается в построении высококачественной реконструкции исходного изображением путем модификации алгоритма семплирования для backward и forward траекторий. Но основным недостатком все же является слабая способность к редактируемости (примеры тут и тут).

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

Вернемся к началу и разберёмся, как работает диффузионная модель, более подробно.

Процесс зашумления

Пусть наши данные (картинки) поступают из распределения $x_0 \sim q(x)$ . Определим процесс зашумления, в котором каждый шаг добавляется небольшое количество гауссовского шума к изначальному изображению для каждого из шагов. Формально это выражается так:

$q(x_t | x_{t - 1}) = \mathcal{N}(x_t | \sqrt{1 - \beta_t}x_{t - 1}, \beta_t \mathbf{I}) \\ x_t = \sqrt{1 - \beta} x_{t - 1} + \sqrt{\beta} \epsilon \\ \epsilon \sim \mathcal{N}(0, \mathbf{I}), ~~~ \beta \in (0, 1).$

Зависимость x_t можно выразить иначе (это понадобится нам в дальнейшем):

$x_t \sim q(x_t | x_0) = \mathcal{N}(x_t | \sqrt{\bar{\alpha_t}}x_0, (1 - \bar{\alpha_t})\mathbf{I})\\ \bar{\alpha_t} = \prod\limits_{i=1}^t (1 - \beta_i).$

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

$q(x_{1:T} | x_0) = \prod\limits_{t = 1}^T q(x_t|x_{t - 1}).$

Из формулы шага видно, что по мере увеличения шага коэффициент $\alpha_t$ стремится к 0, и, следовательно, стремится к нулю среднее . При достаточно больших распределение превращается в распределение чистого шума. Теперь мы поняли, как математически можно описать этап зашумления, но что делать с расшумлением?

Цепочка преобразований данных в прямом и обратном диффузионном процессе. Источник: Ho et al. / arXiv, 2020 с дополнительными обозначениями.

Процесс расшумления

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

$x_{t-1} \sim q(x_{t-1} | x_t)$

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

$p_\theta(x_{0:T}) = p(x_T) \prod\limits_{t=1}^T p_\theta(x_{t-1} | x_t), \\ p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t - 1} | \mu_\theta(x_t, t), \mathbf\Sigma_\theta(x_t, t))$

Такая модель обучается следующим способом — каждую итерацию мы случайно выбираем из наших данных изображение (), генерируем случайный шум ( $\epsilon_t$ ), подаем в модель и делаем шаг оптимизации по обычной функции потерь — -лоссу.

$L_t = || \epsilon_t - \epsilon_\theta(x_t, t) ||_2 = \\ || \epsilon_t - \epsilon_\theta(\sqrt{\bar{\alpha_t}}x_0 + (1 - \bar{\alpha_t})\epsilon_t,~ t) ||_2$

Таким образом мы научим нейросеть правильно расшумлять изображения.

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

Про наш метод

Начну издалека. Представим, что у нас есть такая функция $\nabla_x\log p(x)$ , и мы хотим обусловить распределение на . В этом нам поможет байесовский вывод:

$\nabla_x\log p(x | y) = \nabla_x\log\dfrac{p(y|x)p(x)}{p(y)} \\ = \nabla_x\log p(y|x) + \nabla_x\log p(x) - \nabla_x\log p(y) \\ = | \nabla_x\log p(y) = 0 | = \nabla_x\log p(x) + \nabla_x\log p(y|x)$

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

Вспомним диффузионные модели и редактирование на их основе. После того, как обучение модели завершено, нейросеть начинает предсказывать нам не что иное, как градиент логарифма правдоподобия данных $\nabla_x\log p(x)$ . Мы решили использовать это свойство и возможность обуславливать распределение через прием сверху, который называется guidance.

Вернусь на главу раньше и напомню, что изначально наше исследование было направлено на методы редактирования, основанные на кешировании внутренних представлений исходных изображений при forward процессе и подстановке их при редактировании на backward. А что, если не подставлять их, а использовать в качестве гайденса на backward траектории, подумали мы? И сделали!

План такой:

на forward этапе мы кэшируем внутренние представления диффузионной модели
на backward этапе (генерации) при каждом шаге расшумления мы модифицируем предсказание модели с помощью дополнительного слагаемого:

$\nabla_{x_t} \log \mathcal{N} (u_t |\hat{u}_t, x_t) = \nabla_{x_t} ||u_t(x_t) - \hat{u}_t||_2,$

где — внутренние представления в момент генерации. Таким образом, вместо того, чтобы явно заменить на $\hat{u}_t$ , как делали предшествующие методы, мы модифицируем траекторию генерации так, чтобы эти представления менее болезненным образом приближались друг к другу. Мы назвали этот метод Guide-and-Rescale.

Теперь перейдем к экспериментам.

Эксперименты

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

Локальное изменение внешности и одежды человека
Изменение эмоций человека
Замена одного животного на другое
Глобальная стилизация изображения

Для каждого типа мы собрали 20 примеров редактирования, на которых проводили тесты.

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

В качестве конкурентов мы выбрали подходы, лежащие в отрытом доступе: NPI, NPI Prox и NTI, совмещенные с P2P, MasaCtrl, ProxMasaCtrl, PnP и EDICT. Мы использовали оригинальный код авторов с параметрами по умолчанию, рекомендованными в описании каждого метода. Для нашего метода мы также фиксировали гиперпараметры во время оценки.

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

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

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

Первое — это степень, в которой результат метода содержит желаемую правку. Для измерения этой характеристики мы используем метрику CLIP, которая рассчитывает соответствие между отредактированным изображением и инструкцией для редактирования. Второе свойство — это степень, в которой метод сохраняет структуру и детали исходного изображения, которые не должны меняться в процессе редактирования. Для этого мы используем метрику LPIPS между исходным изображением и его отредактированной версией.

Результаты этих метрик для нашего метода и других приведены в Таблице 1. Они показывают, что наш метод достигает наилучшего баланса между сохранением структуры исходного изображения и степенью редактирования. По метрике LPIPS мы показали наилучшие результаты, за исключением метода EDICT, который мы значительно превзошли по показателю CLIP. По показателю CLIP мы также лучше всех справляемся с задачей, кроме метода PnP, который сильно уступает в способности сохранять структуру исходного изображения по метрике LPIPS. Остальные подходы показывают одинаково худшие результаты по этим двум метрикам.

По сравнению с наиболее продвинутыми методами EDICT и PnP наш метод более эффективен с вычислительной точки зрения, что видно по времени работы при редактировании изображения. Для количественной оценки расстояния между распределениями отредактированных и реальных изображений мы рассчитали FID Score на задаче преобразования собаки в кошку. Как показано в Таблице 1, наш метод превосходит остальные по FID Score. MasaCtrl и ProxMasaCtrl показывают самый высокий FID. Наиболее близкий к нашему методу показатель FID Score имеет PnP, но он работает в несколько раз медленнее нашего метода.

Таблица 1. Сравнение методов по метрикам.

Провели мы и опыты на людях (*зловещий смех*)! Для пользователей мы подготовили всё те же 4 типа заданий, в которых люди оценивали, насколько наш метод предпочтительнее других, в форме вопроса о парном сравнении «наш методов против другого». Мы опросили 62 пользователя, которые ответили в общей сложности на 960 вопросов. В каждом вопросе было показано исходное изображение, инструкция к редактированию и два результата. Пользователи должны были ответить на два вопроса. Первый вопрос (Q1) — «Какое изображение лучше всего соответствует желаемому описанию редактирования?». Второй вопрос (Q2) — «Какое изображение лучше всего сохраняет общую структуру „Оригинального изображения“?».

Результаты исследования пользователей представлены в Таблице 2. Каждое значение представляет собой процент пользователей, которые предпочли наш метод по сравнению с соответствующим бейзлайном в первом столбце. По качеству редактирования наш метод предпочтительнее почти всех бейзлайнов и субъективно схож с P2P + NPI Prox. Однако по качеству сохранения наш метод лучше, чем P2P + NPI Prox. Подводя итог, можно сказать, что наш метод устойчив в широком диапазоне различных редактирований и субъективно дает лучшие результаты по качеству редактирования при сохранении областей исходного изображения, которые не должны быть отредактированы.

Таблица 2. Исследование на пользователях.

Заключение

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

Надеюсь, данное чтение было понятным, интересным и познавательным для вас! Если появились вопросы, то готов ответить на них в телеграмме @MACderRu!

Статья: https://arxiv.org/abs/2409.01322

Код статьи на Github: https://github.com/FusionBrainLab/Guide-and-Rescale

HuggingFace demo: https://huggingface.co/spaces/AIRI-Institute/Guide-and-Rescale

Сайт конференции ECCV: https://eccv2024.ecva.net

Список авторов: Вадим Титов, Мадина Халматова, Александра Иванова, Дмитрий Ветров, Айбек Аланов. Анализ смежных работ в разделах 1 и 2 был получен Айбеком Алановым при поддержке гранта для исследовательских центров в области ИИ, предоставленного Аналитическим центром при Правительстве РФ Центру ИИ НИУ ВШЭ.

В качестве дополнительной литературы о text-to-image задачах порекомендую статью на Хабре, посвященную выходу модели Kandinsky 3.0.

Для сотрудничества можно обращаться к руководителю лаборатории Fusion Brain Андрею Кузнецову и руководителю группы «Контролируемый генеративный ИИ» Айбеку Аланову.

Следите за новостями в каналах: AI Brain, CompleteAI, AbstractDL, Dendi Math&AI.

ссылка на оригинал статьи https://habr.com/ru/articles/837268/