Саяк Пол из Hugging Face: Как трансформеры захватили мир диффузионных моделей

В рамках курса Stanford CS25 Саяк Пол, ведущий инженер из Hugging Face, представил глубокий технический обзор эволюции диффузионных моделей. Главная тема выступления — масштабный переход индустрии от традиционных сверточных архитектур (UNet) к архитектурам на базе трансформеров, таким как DiT и MMDiT, которые сегодня лежат в основе Stable Diffusion 3, Flux и Sora.

🎨 От шума к изображению: Основы диффузии 2:44

Диффузионные модели представляют собой итеративный процесс превращения случайного гауссова шума в реалистичное изображение . В отличие от GAN (генеративно-состязательных сетей), которые работают «в один проход», диффузия последовательно очищает данные от шума.

Саяк Пол выделяет ключевые компоненты современной системы генерации текста в изображение:

• Текстовые энкодеры: Преобразуют промпт в эмбеддинги. Современные модели, такие как Stable Diffusion 3, используют сразу три энкодера для повышения точности .
• Зашумленные латенты (Noisy Latents): Исходная точка в виде сжатого представления шума.
• Временной шаг (Time step): Сообщает модели, на каком этапе траектории она находится.
• Ядро диффузионной сети: Основной вычислительный блок, который предсказывает шум.
• Декодер: Превращает очищенный латентный код в итоговые пиксели .

По мнению Пола, сегодня доминируют модели в латентном пространстве (Latent Space Diffusion), так как работа напрямую с пикселями слишком затратна с точки зрения вычислений . Новым трендом становится Flow-matching (используется в Flux и SD3), где путь от шума к данным моделируется как прямая линия, что упрощает обучение .

🏗️ Почему UNet уходит в прошлое? 13:47

Долгое время архитектура UNet была стандартом в генерации изображений (начиная с DDPM и заканчивая SDXL в 2023 году). Однако Саяк Пол называет структуру гигантских UNet «плохими новостями» для разработчиков из-за их невероятной сложности .

Проблемы традиционного UNet:

• Чрезмерная громоздкость: сочетание сверточных слоев, ResNet-блоков и кастомных трансформер-блоков .
• Сложность масштабирования: сверточные слои не обладают такой же «божественной благосклонностью» к масштабированию, как чистые трансформеры .
• Трудности интеграции: UNet сложнее объединять с другими модальностями (например, с LLM-бэкендами) .

Саяк Пол утверждает, что переход к чистым трансформерам был неизбежен, чтобы использовать все наработки из области NLP и Vision Transformers (ViT) .

🚀 DiT: Диффузионный трансформер 21:29

Архитектура DiT (Diffusion Transformer) сохраняет стандартный проход Vision Transformer, но адаптирует его под задачи генерации.

Технические особенности DiT:

• Патчификация (Patchification): Входное изображение разбивается на мелкие квадраты (патчи), которые превращаются в токены .
• Адаптивная нормализация слоев (adaLN): По словам Пола, это лучший способ внедрения условий (текста или времени) в модель. Вместо кросс-внимания используется модуляция параметров масштаба и сдвига в слоях нормализации .
• Обучение: Каждый блок инициализируется как идентичный (Identity block), что критически важно для стабильности обучения на ранних этапах .

Эксперименты показывают, что adaLN значительно превосходит кросс-внимание при работе с простыми условиями, такими как метки классов, и при этом требует меньше ресурсов .

📝 PixArt-alpha: Текст вместо классов 28:45

Модель PixArt-alpha стала одним из первых шагов к полноценному использованию DiT для генерации по тексту. Саяк Пол выделяет несколько ключевых решений этого проекта:

• Использование T5-xxl: Переход от CLIP к более мощному текстовому энкодеру T5 позволил модели понимать длинные и сложные промпты .
• Оптимизация adaLN: Команда PixArt заменила сложные вычисления на таблицы эмбеддингов, что сократило объем вычислений на 27% .
• Инициализация: Обучение ускорялось за счет использования весов от уже обученных классификационных DiT-моделей .

При размере всего в 0.6 млрд параметров PixArt-alpha показала впечатляющую точность следования инструкциям (prompt following) .

🔍 Решение проблемы 4K и квадратичной сложности 42:28

Стандартное внимание (Attention) имеет квадратичную сложность. При генерации изображений в высоком разрешении (например, 4K) объем оперативной памяти может достигать 190 ГБ даже для одного изображения .

Для решения этой проблемы была предложена архитектура SANA:

1. Линейное внимание: Вместо стандартного механизма $N \times N$ используется линейный вариант, где сложность ограничена $O(N)$ .
2. Mix-FFN блоки: Поскольку при линейном внимании теряется локальность патчей, в модель добавляются небольшие сверточные слои для восстановления связей между соседними пикселями .
3. Отказ от позиционных эмбеддингов: Использование сверток в Mix-FFN позволило полностью отказаться от явного указания позиций (NoPe — No Positional Embeddings) .

💎 Stable Diffusion 3 и MMDiT 48:11

В Stable Diffusion 3 была представлена архитектура MMDiT (Multi-modal Diffusion Transformer). Главная идея здесь — разделение модальностей. По мнению Пола, объединение текста и изображений в одном пространстве на ранних этапах приводит к наследованию взаимных смещений (biases) .

Особенности MMDiT:

• Раздельные проекции: Текст и изображения имеют собственные наборы матриц QKV и слоев нормализации .
• Ко-эволюция: Эмбеддинги обеих модальностей развиваются параллельно, но взаимодействуют через общие блоки внимания .
• QK-нормализация: Необходима для предотвращения нестабильности обучения при больших масштабах .

Саяк Пол отмечает, что MMDiT — это «вычислительно ненасытная» архитектура . Чтобы оптимизировать её, современные модели (например, Flux от Black Forest Labs) используют комбинацию из MMDiT-блоков и стандартных DiT-блоков, где модальности объединяются после первичной обработки .

🛠️ Будущее: Контроль, Видео и In-Context Learning 1:08:43

В завершение Саяк Пол обозначил горизонты развития отрасли:

• Контроль (Structural Control): Развитие ControlNet и его аналогов для DiT (например, FluxControl) позволяет задавать позы, карты глубины и сегментации .
• Видео: Главный вызов здесь — 3D-внимание. По мнению спикера, факторизованное внимание (разделение на время и пространство) работает плохо . Самые многообещающие модели сейчас — Wan и LTX-Video.
• In-Context Learning: Исследователи пытаются научить диффузионные модели обучаться «в контексте», как это делают LLM. Проекты вроде Transfusion и FuseDiT (разработка группы Саяка) объединяют авторегрессию на дискретных токенах и диффузию на непрерывных данных .

Саяк Пол также пригласил всех изучать библиотеку diffusers от Hugging Face, где реализованы практически все упомянутые модели в открытом доступе .