SmoothCache: Ускорение Инференса для Диффузионных Трансформеров

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

Проблема и Подход

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

Основные Принципы SmoothCache

• Анализ ошибок представления: SmoothCache анализирует ошибки представления слоев на основе небольшого калибровочного набора данных, чтобы адаптивно определять, какие слои можно кэшировать.
• Адаптивное кэширование: Используя эти ошибки, техника определяет оптимальный уровень кэширования на разных этапах процесса деноизинга, позволяя повторно использовать ключевые особенности без значительного ухудшения качества генерации.
• Универсальность: SmoothCache не требует специфических предположений о модели или дополнительного обучения, что делает её универсальной для различных архитектур DiT.

Результаты и Применение

Эксперименты показали, что SmoothCache может ускорить инференцию на 8% до 71% в зависимости от модальности и конфигурации модели, при этом качество генерации либо остаётся на прежнем уровне, либо улучшается. Техника была успешно применена к различным моделям:

• DiT-XL для генерации изображений: Улучшение скорости инференса с сохранением высокого качества изображений.
• Open-Sora для текст-в-видео: Значительное ускорение процесса генерации видео.
• StableAudioOpen для текст-в-аудио: Оптимизация инференса для создания аудиоконтента.

Технические Детали

SmoothCache работает на основе следующих принципов:

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

Ограничения и Будущее Развитие

Несмотря на свои преимущества, SmoothCache имеет ограничения:

• Зависимость от архитектуры: Эффективность зависит от наличия повторяющихся блоков DiT и остаточных соединений.
• Возможные ошибки: Кэширование может вводить ошибки, особенно если ошибки представления слоев накапливаются.

Будущие исследования могут быть направлены на:

• Улучшение управления зависимостями между слоями для минимизации ошибок.
• Исследование влияния вариаций ошибок калибровки на качество и скорость инференса.

Заключение

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