Обрезка токенов для лучшего кэширования: 9-кратное ускорение на Stable Diffusion бесплатно

Основная информация

• ID статьи: 2501.00375
• Название: Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
• Авторы: Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
• Категория: cs.CV (компьютерное зрение), cs.LG (машинное обучение)
• Дата публикации: 31 декабря 2024 г.
• Ссылка на статью: https://arxiv.org/abs/2501.00375
• Ссылка на код: github.com/EvelynZhang-epiclab/DaTo

Аннотация

Stable Diffusion достигла значительного успеха в области генерации изображений из текста, однако её итеративный механизм удаления шума приводит к высоким вычислительным затратам и медленной скорости генерации. Хотя методы, такие как кэширование признаков, привлекают внимание благодаря своей эффективности и простоте, простое повторное использование признаков, вычисленных на предыдущих временных шагах, приводит к тому, что признаки соседних временных шагов становятся похожими, снижая динамичность признаков во времени и в конечном итоге влияя на качество генерируемых изображений. В данной работе предлагается метод динамически осведомлённой обрезки токенов (DaTo) для решения ограничений кэширования признаков. DaTo выборочно обрезает токены с низкой динамичностью, позволяя участвовать в слоях самовнимания только высокодинамичным токенам, тем самым расширяя динамичность признаков между временными шагами. При применении к Stable Diffusion на ImageNet метод достигает 9-кратного ускорения при снижении FID на 0,33; на COCO-30k наблюдается 7-кратное ускорение со значительным снижением FID на 2,17.

Предпосылки и мотивация исследования

Постановка проблемы

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

Ограничения существующих методов

Текущие методы ускорения моделей диффузии включают:

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

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

Мотивация исследования

Авторы через экспериментальные наблюдения обнаружили, что по сравнению с исходной Stable Diffusion модель с кэшированием признаков показывает значительное снижение различий в признаках между соседними временными шагами. Это поднимает критический вопрос: можно ли одновременно выполнять кэширование признаков и сохранять правильную динамичность признаков?

Основные вклады

1. Предложен метод динамически осведомлённой обрезки токенов (DaTo): путём обрезки токенов, динамичность которых была снижена кэшированием признаков на разных временных шагах, и восстановления их токенами с высокой динамичностью, избегается снижение качества генерации, вызванное кэшированием признаков.
2. Разработана стратегия эволюционного поиска: предложен метод поиска оптимальной стратегии кэширования признаков и обрезки токенов через эволюционный подход, полностью раскрывая потенциал DaTo.
3. Достигнуто значительное повышение производительности: обширные эксперименты на Stable Diffusion и SDXL показывают, что без переподготовки и дополнительных данных можно достичь ускорения до 9× на Stable Diffusion без потери качества генерации.

Подробное описание метода

Определение задачи

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

Архитектура модели

1. Динамически осведомлённая обрезка токенов (DaTo)

Базовый выбор токенов:

• Оценка временной разности шума: для t-го временного шага вычисляется абсолютная разница выходов двух соседних временных шагов:

  DiffScore = (1/C) * Σ|f_up_0(x_{t+2})_c - f_up_0(x_{t+1})_c|
  

• Выбор токенов на основе патчей: изображение разбивается на неперекрывающиеся патчи размером s×s, в каждом патче выбирается токен с наивысшим DiffScore в качестве базового токена.

Выравнивание CFG: Для обработки классификатора свободного руководства (CFG) позиции базовых токенов из условной генерации копируются в безусловную генерацию:

X_base,i,j[k] = X_base,i,j[k - B/2], k ∈ {B/2, B/2+1, ..., B-1}

Выбор обрезаемых токенов: На основе косинусного сходства выбираются K токенов, наиболее похожих на базовые токены, для обрезки:

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

Восстановление обрезанных токенов: Обрезанные токены восстанавливаются путём прямого копирования их наиболее похожих базовых токенов.

2. Кэширование признаков с учётом временных шагов

Обрезка пространства поиска:

• Глубина кэширования d ограничена {0, 1, 1/2}
• Коэффициент обрезки r ограничен {0,3, 0,4, 0,5, 0,6, 0,7}

Алгоритм эволюционного поиска: Используется многоцелевой алгоритм оптимизации NSGA-II с целями оптимизации:

• Задержка вывода
• Качество генерации (FID)

Процесс поиска включает стандартные эволюционные операции: отбор, кроссовер, мутацию и т.д., в результате чего получается оптимальная стратегия, зависящая от шагов F(t).

Технические инновации

1. Механизм восстановления динамичности: путём выборочной обрезки низкодинамичных токенов и восстановления их высокодинамичными токенами успешно восстанавливается распределение динамичности признаков, нарушенное кэшированием признаков.
2. Единая структура кэширования-обрезки: объединение кэширования признаков и обрезки токенов в единую структуру, независимую от обучения, реализует повторное использование информации на уровне времени и токенов.
3. Адаптивный поиск стратегии: для различных избыточных характеристик на разных временных шагах предлагается метод автоматического поиска оптимальной глубины кэширования и коэффициента обрезки.

Экспериментальная установка

Наборы данных

• ImageNet-1k: генерация 2000 изображений размером 512×512 (2 на класс)
• COCO-30k: генерация 30000 изображений (1 на подпись)
• Набор проверки MS COCO: для оценки SDXL, генерация 5k изображений размером 1024×1024

Метрики оценки

• FID (Fréchet Inception Distance): измерение качества генерации
• CLIP Score: оценка выравнивания текста и изображения
• Inception Score: оценка качества изображения
• Задержка и коэффициент ускорения: оценка эффективности

Методы сравнения

• DDIM/DPM: быстрые семплеры
• ToMeSD: метод объединения токенов
• DeepCache: метод кэширования признаков
• DeepCache & ToMeSD: наивный комбинированный метод

Детали реализации

• Использование алгоритма NSGA-II с размером популяции 20, 100 поколений
• CFG scale: 7,5 (SD v1.5), 9,0 (SD v2), 7,0 (SDXL)
• Количество шагов выборки: 50 шагов PLMS
• Тестирование на одном GPU 4090

Результаты экспериментов

Основные результаты

Stable Diffusion v1.5 (ImageNet):

• Конфигурация e1: 9,01× ускорение, FID снизился с 27,64 до 27,31
• Превосходит методы сравнения во всех конфигурациях

Stable Diffusion v2 (ImageNet):

• Конфигурация e2: 7,25× ускорение, FID составляет 28,20
• По сравнению с исходной моделью FID снизился с 29,8 до 28,20

Набор данных COCO-30k:

• SD v1.5: 7× ускорение, FID снизился с 12,15 до 9,98 (снижение на 2,17)
• SD v2: 7,25× ускорение, FID снизился с 13,68 до 13,88

SDXL (MS COCO):

• 2,32× ускорение, FID снизился с 24,25 до 23,10
• Значительно превосходит DeepCache (1,75×) и DeepCache&ToMeSD (1,78×)

Абляционные исследования

Эффективность DiffScore: При различных параметрах кэширования и коэффициентах обрезки использование DiffScore постоянно улучшает оценку FID, доказывая эффективность оценки временной разности шума.

Влияние выравнивания CFG: С увеличением коэффициента обрезки выигрыш от конфигурации с выравниванием CFG постепенно увеличивается, при высоком коэффициенте обрезки (0,7) улучшение FID варьируется от 13 до 30 пунктов.

Анализ примеров

Результаты визуального сравнения показывают превосходную производительность DaTo в нескольких аспектах:

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

Экспериментальные находки

1. Восстановление динамичности признаков: DaTo успешно восстанавливает распределение различий в признаках до уровня, близкого к исходной Stable Diffusion
2. Эффект разреженного кодирования: умеренная обрезка токенов и кэширование признаков могут улучшить производительность модели, сосредоточиваясь на ключевых признаках
3. Обобщение стратегии: стратегия, найденная на SD v1.5, хорошо работает на SDXL и других наборах данных

Связанные работы

Эффективные модели диффузии

• Сокращение количества шагов выборки: DDIM, модели согласованности и т.д.
• Сжатие сети: квантизация, обрезка, дистилляция и т.д.
• Оптимизация архитектуры: улучшения U-Net, оптимизация Transformer и т.д.

Стратегии сокращения токенов

• Методы обучения: DynamicViT, A-ViT и другие, использующие вспомогательные модели для ранжирования и обрезки
• Эвристические методы: Token Pooling, Token Merging и другие методы, независимые от обучения
• Применение в моделях диффузии: ToMeSD, AT-EDM и другие адаптации для генеративных задач

Механизмы кэширования

• Кэширование U-Net: DeepCache использует временную избыточность для кэширования признаков
• Кэширование DiT: Δ-DiT — стратегия кэширования для Diffusion Transformer
• Проблемы оптимизации: баланс между повышением эффективности и сохранением качества генерации

Заключение и обсуждение

Основные выводы

1. DaTo успешно решает проблему потери динамичности признаков, вызванной кэшированием признаков
2. Адаптивная стратегия, полученная через эволюционный поиск, значительно превосходит фиксированные конфигурации
3. Метод достигает значительного ускорения и улучшения качества на нескольких моделях и наборах данных

Ограничения

1. Стоимость поиска: хотя ≤20 GPU часов приемлемо, всё ещё требуются дополнительные вычислительные ресурсы
2. Зависимость от оборудования: повышение производительности может варьироваться в зависимости от конфигурации оборудования
3. Ограничения в экстремальных условиях: чрезмерно высокий коэффициент обрезки или слишком низкая частота обновления кэша могут снизить производительность

Направления будущих исследований

1. Адаптивное обучение стратегии: разработка более интеллектуальных адаптивных стратегий кэширования и обрезки
2. Адаптация к другим архитектурам: расширение на больше архитектур моделей диффузии
3. Теоретический анализ: углубленное понимание роли принципов разреженного кодирования в моделях диффузии

Глубокая оценка

Преимущества

1. Высокая инновационность: первая систематическая работа по решению проблемы потери динамичности при кэшировании признаков
2. Практичность метода: независимость от обучения, простота развёртывания и интеграции
3. Полнота экспериментов: комплексная оценка на нескольких моделях и наборах данных
4. Теоретическая поддержка: предоставляет теоретическое объяснение на основе разреженного кодирования
5. Открытый исходный код: предоставляет полную реализацию кода

Недостатки

1. Недостаточная глубина теоретического анализа: теоретическое объяснение того, почему метод улучшает FID, относительно простое
2. Зависимость от алгоритма поиска: необходимость эволюционного поиска для нахождения оптимальной стратегии увеличивает сложность использования
3. Ограниченность метрик оценки: в основном полагается на FID для оценки, отсутствуют более разнообразные метрики качества
4. Отсутствие пользовательских исследований: не проведена оценка человеком для проверки качества генерации

Влияние

1. Академическая ценность: предоставляет новые идеи и методы для ускорения моделей диффузии
2. Практическая ценность: может быть напрямую применена к существующим моделям Stable Diffusion
3. Воспроизводимость: предоставляет подробные детали реализации и открытый исходный код
4. Вдохновляющий характер: служит примером применения оптимизации на уровне токенов в генеративных моделях

Применимые сценарии

1. Среды с ограниченными ресурсами: мобильные устройства, вычисления на периферии
2. Приложения реального времени: интерактивные приложения, требующие быстрой генерации
3. Массовая генерация: задачи крупномасштабной генерации изображений
4. Исследовательские прототипы: проекты, требующие быстрой итерации

Библиография

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

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