Введение
В этом параграфе мы снова попробуем решить задачу генерации, когда нам дана выборка объектов из распределения , и хотим научиться генерировать новые объекты из распределения , которых нет в нашей выборке.
Вероятно, вы уже знакомы с другими генеративными моделями, например VAE или GAN-ы. Здесь же мы познакомим вас с еще одним видом генеративных моделей: диффузионные модели, которые стали крайне популярны в последнее время благодаря своему высокому качеству генерации объектов из заданного распределения. В общий чертах, они работают следующим образом: берем шум из и шаг за шагом удаляем компоненты шума до тех пор, пока не получим объект из распределения, см. иллюстрацию ниже.
Более детально
Для детального понимания стоит объяснить, что такое прямой и обратный диффузионные процессы. Прямой процесс заключается в постепенном зашумлении картинки с помощью распределения , а обратный, наоборот, в расшумлении с помощью распределения . Их можно схематично изобразить следующим образом:
Прямой диффузионный процесс определяется как апостериорное распределение . Это распределение также является Марковской цепочкой, которая постепенно добавляет гауссовский шум к объекту . На каждом шаге шум добавляется с различной магнитудой, которая определяется расписанием дисперсий . При правильном выборе расписания в пределе по числу шагов мы должны сойтись к шуму из . В качестве распределений берут нормальные распределения:
Теперь перейдем к обратному процессу и к самой диффузионной модели.
Диффузионная модель - это вероятностная модель с латентными переменными вида , где промежуточные состояния соответствуют зашумленным объектам, a - объект из распределения. Совместное распределение называет обратным диффузионным процессом, который представляет собой Марковскую цепочку из гауссовских распределений :
Таким образом, обратный процесс параметризуется моделью , которая по зашумленному объекту и шагу предсказывает среднее и дисперсию .
Обучение диффузионной модели
Диффузионный модели обучаются, максимизируя вариационную нижнюю оценку (ELBO) логарифма правдоподобия . По тому же принципу обучаются VAE, с тем лишь отличием, что у диффузионных моделей другая форма модели с латентными переменными. Итак, давайте выведем ELBO для диффузии:
- Комментарий
Если вы знакомы с VAE, то вывод должен быть вам понятен, однако ниже приведен вывод с помощью неравенства Йенсена
Теперь вернемся к распределению . Для того чтобы получить , придется итеративно получать . Однако это можно сделать более эффективно благодаря нормальным распределениям. Для этого обозначим и , тогда
- Формальный вывод этого факта
(*) Пояснение ко второму переходу. У нас выходит
Тогда может быть переписано как
- Долгий вывод
Серым в скобках комментарий к последующему переходу.
Пояснение (*). Пользуемся тем, что у нас Марковский процесс, и теоремой Байеса:
Таким образом во время обучения, на каждой итерации параллельно оптимизируются случайные член с помощью градиентного спуск (сэмлируем ). Поскольку все распределения нормальные, то KL между ними можно выписать в явной форме (см. ниже).
- Формула KL между двумя нормальными
Если
Осталось только выписать . Мы знаем, поскольку у нас все распределения нормальные, то и будет нормальным.
Обозначим
- Вывод
Применим формулу Байеса и распишем. Тут мы просто пытаемся понять, как будут выглядеть среднее и дисперсия, выделяя квадратичную форму в показателе экспоненты
Далее перепишем красные и синие выражения в более красивой форме
Другой лосс. Предсказываем шум
В прошлой подсекции наша модель предсказывала среднее и дисперсию нормального распределения. Давайте зафиксируем . Обычно берут или Тогда из предыдущей секции можно переписать как
Это первый момент, как меняется функционал, если мы не хотим предсказывать , а фиксируем её.
Теперь вспомним, что , но благодаря тому, что у нас гауссовское распределение, это можно переписать в виде
Выразим отсюда и получим, что , тогда подставим это выражение в формулу для (из подсекции «Вывод ») и получим
Теперь скажем, что наша модель будет предсказывать . И просто будем «подставлять» его в выражение для выше. Обозначим предсказание модели как — предсказанный шум . Тогда лосс превратиться в
Тем не менее лосс можно еще больше упростить и просто обучать с помощью MSE на .
Итак, алгоритмы обучения и сэмплирования выглядят вот так (на картинке ).
Алгоритм обучения и сэмплирования диффузионной модели (Изображение взято из: Ho et al. 2020)
Стоит отметить, что важным недостатком диффузионных моделей является низкая скорость сэмплирования. Согласно Song et al. 2020: «Требуется 20 часов на генерацию 50 тысяч картинок размера 32х32, используя DDPM, и меньше минуты, используя GAN» (Nvidia 2080 Ti GPU). Тем не менее, в данном направлении был достигнут значительный прогресс и в целом проблема медленного сэмплирования была частично решена: Jiaming Song et al. (2021), Kong & Ping (2021), Bond-Taylor et al. (2021)
Давайте зафиксируем, какие функции потерь можно использовать. Для всех них справедлив тот факт, что мы сэмплируем шаг равномерно во время обучение и оптимизируем соответствующий .
- Оптимизируя член из суммы . Это KL дивергенция между двумя нормальными распределениями
- При фиксированной дисперсии можно оптимизировать взвешенную MSE между средними нормальных распределений
- При фиксированной дисперсии и при предсказании шума с помощью взвешенной MSE. Или просто MSE. является самым популярным вариантом, который на практике дает лучшие результаты.
Выбор расписания
Расписание является гиперпараметром, основными требованиями на который являются невозрастание и чтобы прямой процесс сходился к в пределе по . Второе может гарантироваться тем, что . Вспомним,
Однако на практике оно также проверяется, чтобы было близко к 0.
Также стоит упомянуть, что обычно берут . Но также важно помнить про требования выше, ведь расписание шума непосредственно зависит от .
Чаще всего используют линейное расписание, где . У данных констант нет никакой мотивации, кроме той, которая описана выше. Они были предложены в Ho et al. (2020).
В Nichol & Dhariwal (2021) было предложено косинусное расписание, которое помогло диффузионным моделям достичь лучшего NLL (negative loglikelihood):
Авторы обнаружили, что линейное расписание плохо работает на картинках 64х64 и меньше. А именно, последнии шаги прямого прохода были шумными и малоинформатиыными (просто зашумляем шум еще больше):
Также они обнаружили, что если обучать модель с линейным расписанием только на 80% первых шагов, то модель не становится сильно хуже, что подтверждает неиформативность последних шагов. Далее, они подобрали расписание так, чтобы убывало линейно на большей части отрезка (от 0 до ) и почти не менялось рядом с 0 и . Разницу в для разных расписаний можно увидеть на картинке ниже:
Изображение взято из Nichol & Dhariwal, 2021
- Детали
Также они ограничивают числом 0.999, чтобы в конце процесса не было проблем с численной устойчивостью. Коэффициент используется, чтобы не были слишком малы рядом с нулем. Он равен 0.008. Такое число было выбрано так, чтобы « была немного меньше, чем размер бина одного пикселя, то есть »
Classifier guidance
В Nichol & Dhariwal (2021) был предложен метод условной генерации, который повышает качество генерируемых картинок, при этом уменьшая их разнообразие. Для этого предобучается «шумный» классификатор на зашумленных картинках, то есть . Затем он используется во время сэмплирования, корректируя предсказанное среднее на . В Nichol & Dhariwal (2021) (Секция 4.1) показывают, что данная добавка позволяет превратить распределение в . Важно, что исходная диффузионная модель никак не меняется, что делает трюк еще более привлекательным. Алгоритм сэмплирования можно видеть на картинке ниже. Коэффициент отвечает за силу guidance.
- Мотивация
У генеративной модели GAN есть способ, который позволяет «балансировать» между разнообразием картинок и их качеством — truncation trick. Он заключается в сэмплировании латентного вектора truncated normal distibution. Данный трюк был хорошо описан и исследован в статье про BigGAN. Поэтому в диффузионных моделях тоже хотелось бы иметь метод, который позволяет балансировать между качеством и разнообразием. Авторы предложили classifier guidance, сравнили его с truncation trick и показали, что их метод строго лучше.
Изображение взято из Nichol & Dhariwal, 2021
Classifier-free guidance
Ho & Salimans (2021) предложили метод, в котором guidance достигается без использования дополнительной модели, поскольку это достаточно затратно. Для этого они обучали условную модель , у которой во время обучения реальная метка заменялась с какой-то фиксированной вероятностью (10%) на пустую метку (). Это по сути позволяет нам обучать безусловную модель одновременно с условной Тогда во время сэмплирования делаем так, чтобы предсказание немного менялось в сторону , а именно:
Мотивация этой формулы следовала из формулы Байеса:
Тогда мы можем просто подставить в формулу для classifier guidance из предыдущей подсекции и получить желаемое равенство с точностью до коэффициента .
Овервью ключевых работ на сегодняшний день
- Jonathan Ho et al. «Denoising diffusion probabilistic models.» arxiv Preprint arxiv:2006.11239 (2020)
Основная работа, в которой диффузионные модели (Denoising Diffusion Probabilistic Models, DDPMs) были применены для генерации картинок. Параграф в основном построен на ней.
- Jiaming Song et al. «Denoising diffusion implicit models.» arxiv Preprint arxiv:2010.02502 (2020)
Одна из первых попыток ускорить генерацию объектов. Идея следущая: давайте изменим прямой диффузионный процесс так, чтобы используя предобученную DDPM, приближать новый обратный процесс за меньшее число шагов.
Чтобы не обучать новую модель, нам нужен прямой диффузионный процесс, у которого будет такая же (суррогатная) функция потерь, а обратный процесс все еще останется Марковским. Оказалось, что существует целое семейство не-Марковских прямых процессов, удовлетворяющих этим требования. Это семейство имеет следующий вид:
где и для всех
Среднее было выбрано так, чтобы для всех . (см. Лемму 1 в Приложении B к статье). То есть важно лишь то, чтобы маргинальное распределение не менялось по сравнению с обычным Марковским случаем. Прямой процесс может быть получен с помощью теоремы Байеса:
Тут контролирует степень стохастичности прямого процесса. Можно заметить, что в отличии от исходного диффузионного процесса, предложенный прямой процесс больше не является Марковским, так как каждый теперь зависит и от и от . Схематично, это можно изобразить как на картинке справа. (Слева исходный диффузионный процесс для сравнения)
- Заметка
Авторы обращают внимание, что функция потерь в DDPM зависит от , а не от напрямую. Это означает, что нам нужно выбрать любой другой прямой диффузионный процесс, у которого остались те же.
Далее, мы можем переписать обратный процесс в данном виде:
Заметим, что при прямой процесс становится марковским, а обратный как у DDPM (обычное сэмплирование, описанное в основной секции). При процесс сэмплирования становится детерминистичным (данный способ и называется DDIM). Ускорение сэмплирования достигается засчет использования лишь какого-то подмножества шагов (). Также одним из плюсов детерминистичного сэмплирования является возможность делать семантическую интерполяцию в латентном пространстве (как у GANов).
- Alex Nichol & Prafulla Dhariwal. «Improved denoising diffusion probabilistic models» arxiv Preprint arxiv:2102.09672 (2021)
Улучшение DDPM, в котором был предложен новое расписание шума, что улучшило NLL. Также был изучен вариант, в котором дисперсия предсказывается моделью.
- Prafula Dhariwal & Alex Nichol. «Diffusion Models Beat GANs on Image Synthesis.» arxiv Preprint arxiv:2105.05233 (2021).
Статья, в которой показывается, что DDPM могут генерировать более качественные картинки по сравнению с GANами. Также был предложен метод conditional сэмплирования. Для этого предобучается классификатор на зашумленных сэмплах, а во время сэмплирования среднее нормального распределения «корректируется» на градиент классификатора.
- Jacob Austin et al. «Structured Denoising Diffusion Models in Discrete State-Spaces».arXiv:2107.03006 (2021)
Диффузионные модели на дискретных данных (например, текст). Вместо нормальных распределений используются категориальные. Также была обобщена мультиномиальная диффузия с помощью «матриц перехода», которые задают способ зашумления дискретных данных.
Более подробно: у нас есть — дискретная величина на всех шагах диффузии, тогда для каждого шага определена матрица прямого перехода такая, что . То есть строки матрицы суммируются в единицу. Тогда если обозначить через one-hot-закодированную версию , то прямой процесс можно описать через категориальные распределения:
Как и в нормальных распределениях, можем выписать
Поскольку тут нет такой хорошей параметризации через , как у нормальных распределений, то единственный способ обучать — с помощью KL дивергенции (членами).
Остается только понять, как выбирать . Помимо того, чтобы сумма в каждой строчке была один, требуется, чтобы сходилось (при ) к равномерному распределению в каждой строчке (аналог нормального шума). За конкретными примерами стоит обратиться к статье.
Опишем работу метода GLIDE. Стоит задача генерировать картинки по заданному текстовому описанию. Для этого используется classifier-free guided diffusion model или CLIP. Это два разных варианта модели, которые авторы сравнивают. В первом случае модель обуславливается на эмбеддинги текста, которые были получены из обучаемого трансформера. Во втором случае guidance осуществляется за счет (это по сути градиент лосса метода CLIP) . Тут — это картиночный энкодер (на зашумленных картинках), а — это энкодер текстового входа. В целом, авторы получили, что classifier-free guidance генерирует более качественные картинки.
Способ описать диффузионные модели через стохастические дифференциальные уравнения.
- What are Diffusion Models?. Прекрасный блог от Lilian Weng (OpenAI).