Архитектура

Прошлые работы показали, что увеличение размерности латентных кодов повышает качество. Следуя данному наблюдению, авторы попробовали автоэнкодер с 4 (как и в прошлых SD), 8, 16 каналами. Увеличение ширины стабильно накидывает по метрикам.

Самое примечательное - архитектура трансформера.
В типичной диффузионной UNet архитектуре условие подается в виде модуляции сверток и через cross attention с патчами изображения/латента.

Здесь же токены изображения и текстовые конкатенируются и обрабатываются совместно единым мультимодальным трансформером. В данной работе авторы используют 3 текстовых энкодера разной размерности - CLIP-L/14, CLIP-G/14 и T5XXL. Эмбеддинги первых двух энкодеров конкатятся и паддятся до размерности T5XXL (которая больше, чем у L/14 и G/14) вместе взятых. То есть имеем 3️⃣ потока токенов - текстовые токены от CLIP-L/14 + CLIP-G/14, текстовые токены от T5XXL, токены изображения. Причем FeedForward слои, нормализации и модуляционные слои имеют свои веса для каждого из трех перечисленных выше потоков токенов. Attention проекции общие для всех токенов.

Иногда пре-attention выдает слишком большие значения, и для повышения стабильности обучения авторы нормализуют ключи и значения перед attention c помощью RMSNorm.

Предложенная архитектура работает лучше, чем UViT (UNet с token merging и splitting), CrossAttention и DiT с общими MLP, модуляциями для всех типов токенов. Иметь отдельные слои для разных текстовых энкодеров полезно, хоть выигрыш сравнительно небольшой. Модели сравниваются на основе динамики валидационного лосса / метрик на СС12M.

Следуя прошлым работам, авторы добавляют синтетические описания, сгенерированные через CogVLM в пропорции 50/50%. Их добавление в обучение заметно улучшает согласованность текста и изображения по мнению разметчиков.