Ускорение диффузионных моделей за счёт кэширования
- Введение
- Ускорение диффузионных моделей
- DeepCache
- Мотивация
- Основной пайплайн
- Кэширование для всего процесса инференса
- Non-uniform caching
- Резюме
- Результаты
- Сравнение с бейзлайнами
- Unconditional generation
- Class-conditioning results
- Stable Diffusion test
- Выводы
- Cache Me if You Can
- Мотивация
- Основной пайплайн
- Scale-shift
- Результаты
- Выводы
- Общие выводы
Введение
За последние пару лет диффузионные модели прочно закрепились в мире генеративных (и не только) задач.
С их помощью мы получаем качественный и разнообразный материал при генерации:
- изображений;
- видео;
- музыки.
И это ещё не весь список 🙂
Ускорение диффузионных моделей
К сожалению, этот процесс остаётся достаточно медленным. В предыдущих постах мы рассмотрели различные способы ускорения диффузионных моделей за счёт уменьшения количества шагов, включая consistency models и diffusion distillation (part 1, part 2).
А сегодня мы познакомимся с ещё одной парадигмой ускорения диффузионных моделей — ускорения самого шага сэмплирования.
В нашем обзоре мы рассмотрим две работы, которые вышли примерно в одно время. Их основная идея — кэширование информации при сэмплировании для свёрточных диффузионных моделей.
DeepCache (arxiv)
Мотивация
Как было сказано ранее, ускорять диффузионные модели можно не только за счёт уменьшения количества шагов, но и с помощью ускорения каждого шага. А как можно ускорить шаг? Например, переиспользовать что-то с предыдущих шагов.
Сначала авторы выдвигают предположение о том, что на соседних шагах расшумления генерации имеют схожие высокоуровневые фичи. В этом можно убедиться, например, посмотрев на рисунок ниже. Мы видим, что на соседних шагах извлекаемые из up-sampling block U2 Stable diffusion фичи очень похожи.
Помимо визуального сходства, авторы анализируют три различные диффузионные модели и обнаруживают следующее: независимо от модели хотя бы 1% соседних шагов имеет очень высокое сходство (> 0.95).
А в некоторых случаях (например, DDPM for LSUN-Church и LSUN-bedroom) одни шаги демонстрируют высокую степень сходства с 80% других шагов, как указано на рисунке ниже (с).
Как мы видим, за счёт явного сходства шагов можно сокращать вычислительные расходы с помощью переиспользования результатов прошлых шагов. Давайте разберёмся с этим подробнее.
Основной пайплайн
Архитектура UNet для диффузионной модели состоит из Down, Mid и Up блоков и, соответственно, skip-connections, на которых и сфокусировались авторы.
Пусть у нас есть два соседних шага: \(t\) и \(t-1\). Сгенерируем сначала \(x_t\), затем для генерации \(x_{t-1}\) извлечем high-level features, полученные на предыдущем шаге. Это значит, что мы кэшируем фичи с up-sampling blocks прошлого шага для каждой skip-branch (обведены красным цветом на рисунке 5) ветки \(m\):
\(F_{\text{cache}}^t \leftarrow U_{m+1}^t(\cdot)\)
\(U_m^{t-1} \leftarrow \text{Concat}(D_m^{t-1}(\cdot), F_{\text{cache}}^t)\)
Этот процесс проиллюстрирован на рисунке ниже:
Мы видим, как после кэширования на первом шаге мы затем можем захватывать вычисления одного down-sampling блока \(D_1^{t-1}\).
Кэширование для всего процесса инференса
Описанный выше алгоритм теперь надо каким-то образом применить не только к двум соседним шагам, но и ко всему процессу инференса. Как мы уже видели выше, сходство фичей заметно на нескольких последовательных шагах денойзинга. Поэтому самый простой способ применить этот алгоритм — сделать кэширование просто один раз в \(N\) шагов. То есть полный инференс мы будем проводить на шагах \(\text{Steps} = \{\text{step } \in \mathbb{N}| \text{x} = iN, 0 \leq i \leq k\}\) , \(k = \lceil T/N \rceil\).
Non-uniform caching
Однако такая стратегия не является оптимальной, так как в зависимости от шага, фичи на up-блоках имеют разную степень сходства с фичами соседних шагов. Таким образом, нелинейный выбор шагов для полного инференса должен дать лучший результат генерации. Авторы предлагают сэмплировать на тех шагах, которые имеют наименьшее сходство с соседними шагами. Их предлагается определить следующим образом:
\(\mathcal{L} = \{l_i| l_i \in \text{linear space } \Big( (-c)^{\frac{1}{p}}, (T-c)^{\frac{1}{p}}, k\Big)\}\)
\(\text{Steps} = \text{unique int} (\{i_k |i_k = (l_k)^p + c, \text{where } l_k \in \mathcal{L} \})\)
где \(\text{linear space}(s, e, n)\) — равномерное распределение \(n\) чисел на отрезке от \(s\) до \(e\), а \(\text{unique int}(\cdot)\) — преобразование числа в int, которое следит, чтобы шаги не повторялись. Гиперпараметр \(c\) выбирает «центаральный» шаг для заданного linear space.
В этом уравнении частота изменения индексов изменяется квадратично по мере удаления от центрального временного шага \(c.\) Важно отметить: описанная стратегия представляет собой лишь один из возможных подходов. Альтернативные последовательности, особенно — сосредоточенные на конкретном временном шаге, также могут приводить к улучшению качества изображения.
Авторы дополнительно проводят исследования, выбирая различные значения параметра \(c\):
И параметра \(p\):
Резюме
Таким образом, можно записать основной алгоритм предложенного метода:
Результаты
Сравнение с бейзлайнами
Авторы протестировали свою гипотезу на нескольких диффузионных моделях:
- DDPM;
- LDM;
- Stable Diffusion.
А также для различных сэмплеров шагов:
- 100-step DDIM для DDPM;
- 250-step DDIM для LDM;
- 50-step и 25-step PLMS для Stable Diffusion.
Также авторы для сравнения берут BK-SDM — модель, где они убирают часть слоёв архитектуры для ускорения.
Unconditional generation
Class-conditioning results
Stable Diffusion test
Выводы
Таким образом, представленная идея является альтернативным способом ускорения диффузионок. Она не очень сложная, при этом помогает ускорить процесс без особой потери в качестве.
Ещё картинки
Cache Me if You Can (arxiv)
Однако у предыдущего метода имеется ряд недостатков. Например, шаги для кэширования одинаковые и предопределённые для всех картинок. Более полной и подробной работой со схожей идеей является статья «Cache me if you can» (правда, у неё нет кода 😞).
Давайте разберём её подробнее.
Мотивация
Мотивация здесь такая же, как и в предыдущем методе.
Мы также рассматриваем UNet-архитектуру, состоящую из блоков \(B_i(x_i, s_i)\), \(i \in [0, N-1]\), N — число блоков, s — дополнительная информация.
В предложенных архитектурах диффузионных моделей имеются skip-connection блоки. Таким образом, \(B_i(x,s)\) записывается следующим образом:
\(B_i(x,s) = C_i(x, s) + \text{concat}(x, s)\)
\(C_i(x, s) = \text{layers}_i(\text{concat}(x, s))\)
Авторы проводят аналогичное исследование и смотрят, как сильно изменяется результат в зависимости от шага. Для этого они вводят метрику (обратите внимание, что метрика является относительной величиной!):
\(L1_{rel}(i, t) = \frac{\|C_i(x_t, s_t) — C_i(x_{t-1}, s_{t-1})\|_1}{\|C_i(x_t, s_t)\|_1}\)
Далее авторы анализируют изменение фичей.
Во-первых, они визуализируют результат:
А также считают метрику на 32-х различных картинках с 2-мя разными сидами, и при этом учитывают стандартное отклонение.
Авторы делают следующие выводы:
- Плавное изменение результатов по шагам.
- Разные блоки ведут себя по-разному.
В целом очень похоже на то, что утверждают авторы DeepCache.
Однако авторы улучшают предложенную выше идею, добавляя некоторые модификации:
- динамическое кэширование: автоматический выбор, когда и что кэшировать.
- Shift and Scale: дополнительный скейлинг для избавления от артефактов кэширования.
Давайте разберёмся с этими пунктами подробнее.
Основной пайплайн
Предполагается, что нам не надо вычислять все блоки на всех шагах с учётом слабого изменения, что приводит к идее кэширования.
Однако не каждый блок должен кэшироваться каждый шаг. Чтобы правильно выбрать алгоритм кэширования, авторы предлагают дополнительно оценивать метрику:
\(L1_{rel}(i, t) = \frac{\|C_i(x_t, s_t) — C_i(x_{t-1}, s_{t-1})\|_1}{\|C_i(x_t, s_t)\|_1}\)
А затем сравнивают её с определённым трешхолдом. То есть интуиция следующая: для блока \(i\) мы сохраняем кэшированное значение, посчитанное на шаге \(t_a\), пока накопленное изменение не достигнет трешхолда \(\delta\). Как только мы доходим до трешхолда — значение пересчитывается:
\(\sum_{t=t_a}^{t_{b-1}} \text{L1}_{\text{rel}}(i,t) \leq \delta < \sum_{t=t_a}^{t_b} \text{L1}_{\text{rel}}(i,t)\)
Трешхолд подбирается имперически, ниже приведены результаты для различных значений \(\delta\):
Увеличение трешхолда действительно ведёт к большему ускорению, однако в то же время может заметно снизить качество генерации, добавляя артефактов и в целом ухудшая генерацию.
Авторы приводят пример расписания кэширования для LDM-512 на 20 шагах DPM.
Итак, можно активно кэшировать первые и последние блоки. В то же время более глубокие слои, особенно на ранних шагах генерации, меняются достаточно быстро, и их нельзя так активно кэшировать.
Scale-shift
Этот алгоритм уже работает неплохо, но при таком кэшировании могут возникать артефакты на итоговых картинках.
Для избавления от подобных артефактов авторы предлагают scale-shift adjustment mechanism. Он включает в себя дополнительные, зависящие от шага времени \(t\) scale и shift параметры для каждого слоя, который принимает кэш.
Такие величины получаются при оптимизации на выбранном трейн-сете. Алгоритм проиллюстрирован на рисунке ниже. Модель, работающая с кэшированием, представлена как студент, в то время как обычная модель — учитель. Сначала мы запускаем процесс расшумления студентом из шума, а затем аналогично учителем. Обратите внимание: учитель стартует каждый раз из точки траектории ученика. Затем считается лосс между двумя полученными траекториями (предсказаниями ученика и учителя на каждом шаге).
Результаты
Основные эксперименты авторы проводят с моделями LDM-512 и EMU-768. Все эксперименты проводятся на A100 GPU.
Сравнивают они с двух точек зрения:
- с фиксированным количеством шагов можно провести ускорение, не потеряв в качестве;
- имея фиксированные вычислительные ресурсы можно проделать больше шагов и получить лучшее качество.
Ниже приведены результаты для кэширования с / без shift-scale техникой:
Также ниже приведено сравнение базовых моделей, моделей с кэшированием и с дополнительной shift-scale корректировкой. Отметим, что shift-scale корректировка помогает улучшить метрику FID, однако совсем немного даёт проседание в скорости:
Также авторы приводят анализ картинок людьми, где картинки с / без кэширования занимают одинаковое время. Ниже приведены результаты.
Выводы
В отличие от DeepCache, эта работа предлагает более сложный подход к решению задачи ускорения за счёт кэширования. Авторы предлагают делать динамический кэшинг и дополнительную корректировку на shift и scale. В целом, выглядит более интересно, однако без кода, что усложняет возможность тестирования данной модели.
Общие выводы
Итак, сегодня мы познакомились с двумя интересными статьями, которые предлагают ускорять свёрточные диффузионные модели за счёт кэширования результатов на разных шагах, но не уменьшать количество шагов. Это альтернативный метод ускорения диффузионных моделей.
Тем не менее, упомянутые статьи предлагались в первую очередь для моделей архитектуры UNet, и многие предположения и правила кэширования строились на особенностях этой архитектуры.
Для трансформерной диффузии всё может выглядеть немного иначе, хотя некоторые идеи применимы и там 🙂
Generative Computer Vision
