Разделённые Диффузионные Трансформеры: ускорение генерации высококачественных изображений через разделение семантики и деталей

Диффузионные трансформеры в генерации изображений

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

Стратегии повышения эффективности

Для решения этих проблем исследователи применяют различные подходы для повышения эффективности диффузионных трансформеров. Оптимизированные механизмы внимания, такие как линейное и разреженное внимание, снижают вычислительные затраты. Улучшенные методы выборки, включая логнормальное повторное сэмплирование и перенормировку потерь, стабилизируют обучение. Кроме того, использование индуктивных смещений, специфичных для домена (REPA, RCG, DoD), и маскированное моделирование усиливают структурированное обучение признаков и способности к рассуждению. Модели DiT, SiT, SD3, Lumina и PixArt расширяют фреймворк диффузионных трансформеров на задачи генерации текста в изображение и видео.

Представление Разделённого Диффузионного Трансформера (DDT)

Исследователи из Университета Нанкина и ByteDance Seed Vision предложили Разделённый Диффузионный Трансформер (DDT), который разделяет архитектуру на два компонента: энкодер условий для извлечения семантики и декодер скорости для генерации деталей. Такое разделение устраняет конфликт оптимизации, обрабатывая низкочастотную семантику и высокочастотные детали отдельно. Модель DDT-XL/2 достигает лучших в своем классе значений FID – 1.31 и 1.28 для ImageNet разрешений 256×256 и 512×512 соответственно, при этом обучаясь в 4 раза быстрее.

Механизмы работы DDT

Энкодер условий извлекает семантические признаки (zt) из зашумленных входов, временных шагов и классовых меток. Декодер скорости использует эти признаки для оценки векторного поля скорости, необходимого для генерации изображения. Для обеспечения согласованности представлений zt на разных шагах применяются методы выравнивания представлений и супервизия декодера. Во время вывода используется механизм совместного самокондиционирования, который повторно использует выходы энкодера zt на выбранных шагах для сокращения вычислений. Алгоритм динамического программирования оптимально выбирает эти шаги, обеспечивая баланс между скоростью и качеством изображения.

Обучение и оценка

Модели DDT обучались на ImageNet с разрешением 256×256 с размером батча 256 без обрезки градиентов и разогрева. Для оценки использовались метрики FID, sFID, IS, Precision и Recall с применением VAE-ft-EMA и выборки Эйлера. Исследователи улучшили базовые модели с помощью SwiGLU, RoPE, RMSNorm и логнормальной выборки. DDT стабильно превосходил предыдущие модели, особенно в больших конфигурациях, и сходился значительно быстрее, чем REPA. Дополнительное улучшение было достигнуто за счёт стратегий совместного использования энкодера и настройки соотношения размеров энкодера и декодера.

Влияние и перспективы

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