Недавно исследователи из Калифорнийского университета в Ирвине и других учреждений использовали стратегии отложенного маскирования, MoE и послойного расширения, чтобы снизить стоимость обучения диффузионной модели до 1890 долларов. Это на порядок ниже, чем предыдущий самый дешевый метод (Wuerstchen) за 28400 долларов, в то время как стоимость таких моделей, как Stable Diffusion, еще выше.
Исследователи опробовали различные подходы для снижения этих огромных затрат. Например, оригинальная диффузионная модель требовала около 1000 шагов от шума до изображения, теперь это сократилось до примерно 20 шагов или даже меньше. Когда базовые модули в диффузионных моделях постепенно заменялись с Unet (CNN) на DiT (Transformer), появились некоторые оптимизации, основанные на особенностях Transformer, такие как квантование, пропуск некоторых избыточных вычислений в Attention, конвейеризация и т.д.
Это исследование позволяет обучить диффузионную модель с 11,6 миллиардами параметров с нуля всего за 1890 долларов. Это на порядок лучше по сравнению с SOTA, давая обычным людям надежду на возможность предварительного обучения. Что еще важнее, технологии снижения затрат не повлияли на производительность модели, 11,6 миллиардов параметров дали очень хорошие результаты. Помимо визуального восприятия, числовые показатели модели также отличные, например, оценка FID очень близка к Stable Diffusion 1.5 и DALL·E 2.
Секреты экономии включают:
-
Стратегия отложенного маскирования: использование patch-mixer для предварительной обработки перед маскированием, встраивание информации отброшенных патчей в выжившие, значительно снижая падение производительности из-за высокого маскирования.
-
Тонкая настройка: небольшая тонкая настройка (без маски) после предварительного обучения (с маской) для уменьшения нежелательных артефактов генерации из-за использования маски.
-
MoE и послойное расширение: использование упрощенных слоев MoE на основе выбора экспертов для увеличения параметров и выразительной способности модели без значительного увеличения затрат на обучение. Также рассматривается метод послойного масштабирования, линейно увеличивающий ширину блоков Transformer.
В экспериментальной настройке авторы использовали два варианта DiT: DiT-Tiny/2 и DiT-Xl/2 с размером патча 2. Все модели обучались с помощью оптимизатора AdamW с косинусным затуханием скорости обучения и высоким затуханием весов. На входе модели использовался четырехканальный вариационный автоэнкодер (VAE) из модели Stable-Diffusion-XL для извлечения признаков изображений, также тестировалась производительность новейшего 16-канального VAE в крупномасштабном обучении (экономичная версия).
Авторы использовали фреймворк EDM в качестве единой настройки обучения для всех диффузионных моделей, используя FID и оценку CLIP для измерения производительности моделей генерации изображений. В качестве текстового энкодера была выбрана наиболее часто используемая модель CLIP.
Для обучения использовались три набора данных реальных изображений (Conceptual Captions, Segment Anything, TextCaps), содержащих 22 миллиона пар изображение-текст. Поскольку SA1B не предоставляет реальных подписей, здесь использовались синтетические подписи, сгенерированные моделью LLaVA.