Кэшированное адаптивное объединение токенов: динамическое сокращение токенов и исключение избыточных вычислений в модели диффузии

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

• ID статьи: 2501.00946
• Название: Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model
• Авторы: Omid Saghatchian, Atiyeh Gh. Moghadam, Ahmad Nickabadi (Amirkabir University of Technology)
• Категория: cs.CV (Компьютерное зрение)
• Дата публикации: 1 января 2025 г. (препринт arXiv)
• Ссылка на статью: https://arxiv.org/abs/2501.00946
• Ссылка на код: https://github.com/omidiu/ca_tome

Аннотация

Модели диффузии стали перспективным методом для генерации высококачественных многомерных изображений. Однако эти модели ограничены высокими вычислительными затратами и медленной скоростью вывода, отчасти из-за квадратичной сложности вычислений механизма самовнимания относительно размера входных данных. В данной статье предлагается метод кэшированного адаптивного объединения токенов (CA-ToMe), который решает эту проблему путём вычисления сходства между токенами и объединения токенов с коэффициентом сходства, превышающим пороговый параметр t. Благодаря повторяющимся закономерностям, наблюдаемым в соседних шагах, и изменениям в частоте сходства, метод улучшен путём внедрения адаптивного порога и добавления механизма кэширования. Экспериментальные результаты показывают, что метод, как метод ускорения без обучения, достигает 1,24-кратного ускорения в процессе удаления шума, сохраняя при этом тот же показатель FID, что и существующие методы.

Исследовательский контекст и мотивация

Определение проблемы

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

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

Важность проблемы

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

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

1. Методы сокращения количества шагов выборки обычно требуют переобучения или использования сложных численных решателей
2. Методы обрезки токенов приводят к потере информации и снижению производительности
3. Традиционное объединение токенов (ToMe) использует фиксированный коэффициент объединения, не способный адаптироваться к изменениям распределения сходства токенов на разных временных шагах и слоях

Исследовательская мотивация

На основе двух ключевых наблюдаемых явлений:

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

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

1. Предложение механизма адаптивного порога: динамическая регулировка стратегии объединения в соответствии с распределением сходства токенов, замена фиксированного коэффициента объединения
2. Разработка механизма кэширования: использование сходства между соседними шагами, кэширование пар токенов для сокращения повторных вычислений
3. Реализация ускорения без обучения: метод может быть применён непосредственно к предварительно обученным моделям без переобучения
4. Достижение лучшего компромисса качество-скорость: по сравнению с базовым методом ToMe, достижение более быстрого вывода при сохранении качества изображения

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

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

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

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

1. Механизм адаптивного объединения токенов

Традиционный метод ToMe использует фиксированное соотношение r для объединения токенов, в то время как CA-ToMe вводит пороговое значение сходства t:

Основная идея:

• Разделение изображения на области шага размером sx × sy
• Выбор левого верхнего токена каждой области шага в качестве целевого токена
• Вычисление косинусного сходства между исходными и целевыми токенами
• Объединение только пар токенов с коэффициентом сходства, превышающим пороговое значение t

Анализ преимуществ:

• Сценарий A: когда большинство токенов имеют низкое сходство, фиксированный коэффициент объединения вынуждает объединять несходные токены, приводя к потере информации. Адаптивный порог гарантирует объединение только высокосходных токенов
• Сценарий B: когда большинство токенов имеют высокое сходство (например, на ранних этапах удаления шума), фиксированный коэффициент объединения ограничивает количество объединяемых токенов. Адаптивный порог позволяет объединять больше токенов, повышая эффективность

2. Конструкция механизма кэширования

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

$JaccardDistance(A_n, A_{n+1}) = 1 - \frac{|A_n \cap A_{n+1}|}{|A_n \cup A_{n+1}|}$

где An представляет множество всех пар исходный-целевой токен на шаге n.

Стратегия реализации:

• Установка контрольных точек (checkpoints), вычисление матриц сходства только на определённых временных шагах
• Повторное использование ранее вычисленных пар токенов на шагах без контрольных точек
• Значительное сокращение затрат на повторные вычисления матриц сходства

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

1. Динамическая адаптивность: автоматическая регулировка стратегии объединения в соответствии с распределением сходства, избежание ограничений фиксированных параметров
2. Оптимизация временного измерения: использование избыточности между временными шагами, сокращение объёма вычислений через кэширование
3. Избирательное применение на уровне слоёв: специальное применение оптимизации к верхним слоям U-Net, требующим интенсивных вычислений (D1 и U1)
4. Отсутствие необходимости переобучения: как метод ускорения plug-and-play, может быть применён непосредственно к существующим моделям

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

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

• Набор данных ImageNet-1k: генерация 2000 изображений разрешением 512×512 (по 2 изображения на класс, всего 1000 классов)
• Набор для проверки: использование 5000 изображений для проверки ImageNet-1k для вычисления показателя FID
• Шаблон подсказки: "A high-quality photograph of a classname."

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

1. FID (Fréchet Inception Distance): основной показатель для оценки качества генерируемых изображений
2. Время вывода: среднее время генерации 2000 изображений
3. PSNR: пиковое отношение сигнал-шум, оценка качества реконструкции на уровне пикселей
4. SSIM: индекс структурного сходства, оценка пространственной и структурной согласованности

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

• Базовый метод: исходная Stable Diffusion v1.5
• ToMe: традиционный метод объединения токенов (r=50%)

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

• Оборудование: GPU Tesla V100S
• Шаги диффузии: 50 шагов выборки PLMS
• Масштаб CFG: 7.5
• Размер шага: фиксированный 2×2
• Применяемые слои: применение только к слоям D1 и U1 U-Net

Экспериментальные результаты

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

Ключевые выводы:

• CA-ToMe достигает самой быстрой скорости вывода (6.09s)
• Показатель FID (34.05) превосходит ToMe (34.16) и близок к базовому методу (33.66)
• Достигнут оптимальный баланс между скоростью и качеством

Анализ параметра порога

Наблюдаемые результаты:

• Изменения в диапазоне порога 0.4-0.6 относительно небольшие, так как большинство токенов имеют сходство ≥0.6
• Пороговое значение 0.7 обеспечивает оптимальный компромисс между качеством и скоростью
• Более высокие пороговые значения улучшают качество, но снижают скорость

Сравнение конфигураций кэширования

CONFIG 3 показывает лучшие результаты, согласующиеся с анализом расстояния Жаккара, с большим количеством контрольных точек на шагах 8, 11, 13 и на заключительных шагах.

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

Путём сравнения вклада различных компонентов:

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

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

Методы ускорения моделей диффузии

1. Сокращение количества шагов выборки:
   • Методы дистилляции знаний 26,51,28
   • Неявная выборка 32
   • Продвинутые решатели дифференциальных уравнений 52,33
   • Большинство требуют переобучения
2. Сокращение вычислений на каждом шаге:
   • Методы квантизации 31,36
   • Сокращение токенов 21,40,41,43,44
   • Техники кэширования 24,37,38,39
   • Plug-and-play, без необходимости переобучения

Технологии сокращения токенов

1. Обрезка токенов: прямое удаление неважных токенов, может привести к потере информации
2. Объединение токенов: объединение похожих токенов, сохранение полноты информации
   • ToMe 21: использование фиксированного коэффициента объединения
   • CA-ToMe данной работы: адаптивный порог + механизм кэширования

Техники кэширования

Существующие методы кэширования ориентированы на различные компоненты:

• Кэширование перекрёстного внимания 38
• Кэширование кодировщика U-Net 39
• Кэширование высокоуровневых признаков 24

Данная работа впервые применяет кэширование к вычислению сходства при объединении токенов.

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

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

1. Адаптивный порог эффективно решает ограничения фиксированного коэффициента объединения, динамически регулируя стратегию объединения в соответствии с распределением сходства
2. Механизм кэширования использует избыточность между временными шагами, значительно сокращая повторные вычисления
3. Метод CA-ToMe достигает 1,24-кратного ускорения, сохраняя и даже слегка улучшая качество изображения
4. Характеристика отсутствия необходимости обучения обеспечивает хорошую практичность и масштабируемость метода

Ограничения

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

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

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

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

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

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

Недостатки

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

Влияние

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

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

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

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

Данная работа цитирует 54 связанные публикации, включая в основном:

• Фундаментальную теорию моделей диффузии 1,2,3
• Приложения генерации изображений 4,5,18,19,20
• Методы ускорения 24,25,26,27,28
• Методы обработки токенов 21,40,41,43,44
• Техники кэширования 24,37,38,39

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