От каскадных моделей до картинок в 4к: как эволюционировали диффузионки

На дворе 2025 год. Генерацией картинок и видео в интернете больше никого не удивишь. Генеративный контент повсюду, а его качество настолько высоко, что бывает трудно отличить синтетическую картинку от реальной.

Прогресс не стоит на месте, но какой ценой? Мир сильно усложнился со времён выхода первых моделей Stable Diffusion, подробные рассказы о которых (раз, два) описывали главную и до недавнего времени единственную доступную широкому кругу пользователей качественную открытую модель.

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

Цель этого обзора — дать общее представление о прогрессе и трендах. Некоторые детали в нём могут показаться сложными и непонятными для неспециалиста в области генеративного моделирования. Чтобы помочь разобраться в теории и основных инженерных практиках, стоящих за современными генеративными моделями, мы выпустили бесплатный курс лекций и семинаров в рамках CV week — открытого интенсива ШАДа, который подготовили практикующие разработчики и исследователи генеративных диффузионных моделей, авторы технологии YandexART.

Итак, чем же жил мир генеративных нейросетей последние два года и что сильнее всего повлияло на качественный прорыв в генерации изображений? Будем искать ответы вместе.

Simple diffusion

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

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

И тут группа авторов из Google Brain представляет свою работу simple diffusion. В ней авторы не пытаются спорить с пунктом про более высокую вычислительную эффективность каскадной и латентной парадигм. В то же время они показывают, что концептуально можно обучить диффузию сразу генерировать изображения в высоком разрешении без использования дополнительных моделей Super Resolution и вариационных автокодировщиков.

Добится отличного качества генерации сразу в высоком разрешении помогли несколько трюков.

Адаптация расписания прямого процесса под картинки высокого разрешения 

Дело в том, что чем выше разрешение картинок, тем, как правило, у них выше SNR (signal‑to‑noise ratio, соотношение «сигнал — шум»). А наиболее распространённое на момент выхода статьи alpha‑cosine‑расписание и без того зашумляет картинки медленнее, чем изначально предложенное линейное. Для иллюстрации давайте обратимся к картинке из статьи.

Верхняя строка показывает, как меняется уровень шума на изображении при использовании косинусного расписания. Авторы замечают, что при таком расписании во время обратного процесса у модели будет слишком мало времени, чтобы сформировать основную структуру изображения. При этом чем выше разрешение, тем выраженнее данная проблема.

Введение «сдвинутого» (shifted) расписания (нижняя строка) решает эту проблему и существенно улучшает качество генераций.

Масштабирование архитектуры

Вторая проблема генерации в высоком разрешении — большие карты активаций, потому что они требуют много видеопамяти на хранение данных во время проведения вычислений. Чтобы максимально избежать работы с большими активациями, авторы предлагают такие решения:

1. Сразу даунскейлить входное изображение с помощью вейвлет‑преобразования с довольно хитрым образом подобранными параметрами. Эту операцию можно упростить до свёрточного блока с большим страйдом и ядром, немного увеличив вычислительную сложность, но упростив схему модели.

2. Использовать архитектуру U‑Vision Transformer (U‑ViT), которую можно воспринимать как небольшой U‑Net, который уменьшает размерность активаций вплоть до 16×16. Вариант с обычным U‑Net авторы тоже сделали, но он немного проигрывает по Inception Score.

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

SDXL

Обычно лето — время затишья. Но только не для инженеров Stability.ai! В июле 2023 года выходит статья про модель SDXL, которая произвела настоящий фурор. На сегодняшний день она собрала более 1300 цитирований по версии Google Scholar, что довольно впечатляющий результат.

В чём же заключаются нововведения модели SDXL?

Первая по-настоящему большая opensource-модель

Предыдущие поколения (SDv1.x и SDv2.x) использовали похожую архитектуру с 860 и 865 миллионами параметров в денойзере соответственно, а в качестве текстовых энкодеров — не самые большие модели семейства CLIP. В этой итерации авторы сделали следущее:

1. Подняли размер основной диффузионки до 2,8 миллиардов параметров. Чем больше модель, тем сложнее её обучать, а диффузия в 2,8 миллиарда и сегодня считается не маленькой.

2. Добавили дополнительный текстовый энкодер. В этой версии начали использовать большой CLIP (OpenCLIP ViT‑bigG), а суммарный размер текстовых энкодеров превысил 800 миллионов параметров. Текстовые эмбеддинги из нескольких энкодеров просто конкатенируются. Этот механизм обрёл популярность и активно используется и сейчас.

3. Добавили refiner — дополнительный модуль, который использует механизм SDEdit для увеличения чёткости генераций. Это самый неоднозначный блок, поскольку уже на релизе модели возникли вопросы о том, действительно ли оптимально «доуточнять» генерацию именно в низком разрешении (обычно это делают на финальной стадии уже в высоком разрешении). Его внедрение оказалось не ключевым, поскольку в более поздних итерациях SD его убрали, а другие крупные игроки это решение так и не адаптировали.

Контроль и качество генерации

Micro‑Conditioning — так авторы назвали введение дополнительного вектора‑условия. Через него в модель подаётся дополнительная информация о разрешении и соотношении сторон изображения, которое нужно сгенерировать.

Также в качестве условия добавили параметры кроппинга. Дело в том, что более ранние модели были склонны генерировать обрезанные изображения, как в примере ниже.

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

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

Кстати, SDXL и сейчас остаётся сильным бейзлайном по качеству генерации, а его файнтюны используются во многих открытых проектах.

EMU

Следующим пиком громких анонсов стал сентябрь 2023, когда вышло сразу несколько технических репортов от крупных технологических компаний. Модель EMU — первая text‑to‑image генеративная модель, представленная компанией Meta*.

Перед обсуждением этого и других техрепортов не будет лишним напомнить, чем они концептуально отличаются от научных статей. В первую очередь техрепорт — это описание готовой системы, ориентированной на практическое использование. Изменения, сделанные при создании такой системы, не всегда всесторонне протестированы изолированно друг от друга, а прогресс в качестве, как правило, достигается за счёт большой комбинации различных факторов. Тем не менее техрепорты — крайне полезные источники знаний с точки зрения практика, поэтому мы в Яндексе активно за ними следим.

Итак, на мой взгляд, наиболее интересны в статье EMU два нововведения.

Разделение обучения на pre-training и supervised fine-tuning (SFT)

Читателям, знакомым с другими генеративными моделями, например с LLM, такой подход не покажется удивительным. Действительно, разделение стадий обучения по целям и задачам выглядит вполне естественным.

Суть в том, что сначала, на стадии pre‑training, мы учим модель на максимально широком распределении данных различного качества. В случае EMU это 1,1 миллиард пар «картинка — текст», про которые мы мало что знаем.

После этого модель файнтюнится на небольшом наборе семплов исключительного качества с целью сдвинуть выученное распределение в сторону максимально эстетичных изображений. Главный инсайт этой работы: SFT‑датасет не должен быть большим, но должен быть крайне качественным.

Воронка формирования файнтюн‑датасета выглядит следующим образом: начиная с 200 тысяч достаточно качественных семплов, с помощью автоматических фильтраций и ручного отбора авторы получают всего 2000 семплов, которые впоследствии используют для SFT. Высокое качество генераций показывает важность этой стадии и степень прогресса, которую можно достичь с её помощью.

Использование многоканального автоэнкодера

Латентная диффузия требует использования автокодировщика. Его качество в принципе сильно влияет на качество диффузии. Латентное пространство, формируемое энкодером, определяет пространство, в котором учится генеративная модель. Декодер, в свою очередь, задаёт предел реконструкционного качества генераций.

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

Все предыдущие работы, включая SDXL с его огромным 2,8B денойзером, использовали четырёхканальные вариационные автокодировщики (VAE). В EMU впервые было предложено использовать модели с бо́льшим числом каналов. По утверждениям авторов, это сильно увеличило качество и детализацию генераций.

В будущем всё больше систем будет переходить на многоканальные автокодировщики.

DALL-E 3

Одним из наиболее громких анонсов 2023 года стал выход обновлённой версии модели DALL‑E от OpenAI. Она произвела фурор не только отличным качеством и последующей интеграцией в ChatGPT, но и феерическими графиками жёлтым по белому с полностью дефолтными настройками matplotlib.

Чтобы понять причину высокого качества DALL‑E 3 относительно современников, необходимо понять, как устроено формирование данных, на которых учат диффузионки.

Дело в том, что большинство text‑to‑image‑моделей обучают на крупных наборах данных из интернета вроде LAION 5B, состоящих из произвольных картинок и прикартиночных текстов. Эти тексты могут быть абсолютно любыми, в том числе совершенно никак не относящимися к картинке. Но поскольку текст и картинка находятся рядом (например, на одной веб‑странице), то предполагается, что такой текст лучше связан с картинкой, чем любой другой произвольный текст — в целом это выполняется.

Работа по системе DALL‑E 3 называется Improving Image Generation with Better Captions. И это неспроста. Авторы решают очевидную проблему — устраняют плохие тексты, из‑за которых модели не могут выучить нужные концепты и плохо следуют указаниям пользователей. Вместо этого авторы генерируют синтетические тексты с помощью собственной, специально обученной для этой задачи модели. В ходе экспериментов авторы пробуют разные распределения текстов (более короткие или более подробные), а также смотрят на то, какую долю исходных текстов стоит заменить на синтетику. На практике лучше всего работает замена на подробную синтетику в 95% пар «картинка — текст».

К сожалению, как и в других техрепортах, в этой работе есть много белых пятен. Например, нет графиков, как ведут себя метрики в случае замены 100% текстов, а само разделение на короткие и более подробные тексты весьма условное и сильно зависит от контекста.

В любом случае работа полезна тем, что показывает возможность замены текстов на синтетические, а рост качества подчёркивает важность этого шага.

DiT и PixArt-α

Сентябрь 2023 уже должен был закончиться, и тут появляется работа от коллектива из Huawei Noah»s Ark Lab, Dalian University of Technology, HKU и HKUST. Их подход отличается очень быстрым обучением на малом числе данных.

Несмотря на наличие большого количества нововведений, сообществу эта работа в основном запомнилась как первая серьёзная диффузионная модель, обученная с трансформерной архитектурой денойзера. Про это стоит рассказать отдельно, поскольку данная работа положила начало повсеместному переходу на трансформеры.

На самом деле трансформерная архитектура для диффузионных моделей DiT была предложена ещё в конце 2022 года и мало чем отличается от ViT — кроме разве что дополнительного обусловливания на timestep t и метку класса. И то поскольку в то время ещё было актуально заниматься class‑condition‑генерацией на датасетах вроде ImageNet.

Авторы PixArt‑α модифицировали DiT, добавив возможность подавать текст в качестве условия генерации. В качестве текстового энкодера использовали большую языковую модель T5.

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

ADD и SDXL-Turbo

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

Для ускорения этого процесса были предложены методы дистилляции. Первый из них — Adversarial Diffusion Distillation (ADD) от авторов из Stability.ai. Они сразу же применили метод для получения SDXL‑Turbo — ускоренной версии модели SDXL, способной генерировать изображения, сопоставимые по качеству с генерациями исходной модели всего за четыре шага (то есть примерно на порядок быстрее).

Метод устроен достаточно просто. Модель‑студент инициализируется моделью‑учителем и обучается с использованием двухкомпонентного лосса:

• с одной стороны, модель‑ученик пытается сгенерировать картинку, которая будет максимально похожа на исходную реальную (adversarial loss);

• с другой — пытается воспроизвести генерацию учителя (distillation loss).

Стоит обратить внимание, что использование adversarial loss делает модель в некотором смысле не чисто диффузионной. Адверсариальная компонента, по сути, делает из модели GAN, инициализированный диффузией, — и у этого есть свои практические последствия.

Главная проблема GAN как генеративных моделей — mode collapse, то есть состояние, в котором модель учится генерировать реалистичные изображения, но только из очень узкого распределения. Конечно, инициализация диффузией и использование distillation loss позволяет уменьшить этот эффект, однако факт остаётся фактом: сильная адверсариальная дистилляция приводит к падению вариативности генераций в угоду скорости.

С самой же работой связан довольно забавный факт: только спустя год она была опубликована и представлена на конференции. Дело в том, что при первой подаче авторы получили плохие оценки рецензентов и, несмотря на огромную популярность и цитируемость, работа была отклонена. На ECCV 2024 справедливость была восстановлена, а работа получила свой заслуженный oral на одной из главных конференций по компьютерному зрению.

Несмотря на минусы, ADD и сегодня остаётся популярным методом дистилляции диффузионных моделей. Правильный подбор гиперпараметров и длительности обучения позволяет получить все плюсы, почти не почувствовав минусов метода.

SD 3

Наконец, мы врываемся в 2024 год. На дворе весна — и новый релиз от Stability.ai! Как и их прошлые работы, эта включает в себя множество нововведений и изменений по сравнению с предыдущими версиями. Большинство из них мы уже обсудили выше:

• использование трансформерной архитектуры, вроде предложенной в PixArt‑α, но с несколькими изменениями;

• увеличение количества текстовых энкодеров. Теперь их аж три: две модели CLIP и LLM в лице T5 XXL;

• улучшенный 16-канальный автоэнкодер, как в EMU, и использование синтетических текстов, как в DALL‑E 3.

Отличительная фишка этого поколения — использование альтернативной формулировки диффузии под названием Rectified Flow, что и дало название работе: Scaling Rectified Flow Transformers for High‑Resolution Image Synthesis.

По своей сути Rectified Flow — это один из вариантов диффузии, обучение которой приводит к более «спрямлённым» траекториям генерации. Это, в свою очередь, означает, что обученная таким образом модель должна быть способна исходно генерировать изображения за меньшее число шагов при сохранении качества, а ещё лучше поддаваться дистилляции.

Отдельного внимания в работе заслуживает систематическое сравнение архитектур и постановок диффузии. Авторы используют Rectified Flow со своим вариантом архитектуры. И не просто из общих соображений и веры в результат, а на основании экстенсивных экспериментов, в которых показывают, что данное сочетание является наиболее эффективным.

К сожалению, в дело вмешался человеческий фактор и релиз совпал с кадровыми перестановками в компании. Виноваты ли они или какие‑то другие факторы, но модель по итогу получила много критики и не обрела такой широкой популярности, как предыдущие релизы. Тем не менее статья написана очень добротно и содержит массу полезных инсайтов.

Kolors

Лето 2024 ознаменовалось выходом довольно интересного техрепорта от авторов из Китая по модели Kolors. Лично для меня он был интересен тем, что, несмотря на большое количество публикаций по генеративным моделям, масштабных техрепортов по готовым системам из Поднебесной почти не поступало, поэтому этот редкий зверь сразу привлёк моё внимание.

Основные характеристики системы:

• U‑Net на 2,6 миллиардов параметров в качестве денойзера. Выбор свёрточной модели кажется интересным в контексте того, что большинство к этому моменту уже успели перейти на трансформеры.

• Много сил потрачено на чистку данных, в том числе с помощью VLM и OCR. Тут прослеживается некое ноу‑хау, поскольку ранее такие комбинации не встречались — звучит разумно.

• GLM — мультиязычная генеративная модель в качестве текстового энкодера. Естественно, под мультиязычностью в первую очередь понимается нативное понимание китайского.

На последнем пункте про GLM хотелось бы остановиться подробнее. Дело в том, что использование генеративных моделей в качестве текстовых энкодеров мы встречаем довольно редко. Причина у этого очень простая: генеративные модели учатся правдоподобно генерировать новый текст, для чего не обязательно требуется детально понимать все элементы входной последовательности. В то же время диффузионной модели требуется энкодер, который досконально поймёт все детали пользовательского ввода, поэтому на практике мы часто видим использование CLIP и BERT‑подобных моделей. Несмотря на это, авторы считают GLM хорошим выбором и подтверждают свою интуицию делом — получают очень достойные генерации.

В целом техрепорт выглядит очень солидно, будем ждать новых работ от азиатских исследователей.

Nvidia SANA

Несмотря на традиционно сильное исследовательское направление, до конца 2024 года компания Nvidia не зарелизила ни одного техрепорта. Это вызывало вопросы: либо они не верят в продуктизацию text‑to‑image и в целом не занимаются этим направлением, либо готовят что‑то крупное, ведь с каждым новым релизом выделиться становилось всё сложнее.

Представленная в ноябре модель SANA действительно оказалась оригинальной. Фокус работы — в возможности нативно, без использования дополнительных Super‑Resolution‑моделей, генерировать изображения в разрешении 4К, причём за разумное время. Более того, генерации в более привычных разрешениях 512×512 и 1024×1024 пикселей работают очень быстро без использования дистилляции. Как же им это удалось?

Во время разбора модели EMU мы уже обращались к теме важности выбора автокодировщика, в латентном пространстве которого будет учиться диффузия. Но до этого мы обсуждали только количество каналов в латентах. Теперь давайте обсудим другой параметр — глубину, или scaling factor (SF).

Дело в том, что у подавляющего большинства современных автокодировщиков SF равен 8. Это значит, если взять картинку 512×512×3 пикселей и пропустить через энкодер, получится латент с разрешением 64×64×С, где C — число каналов, которое мы обсуждали выше.

Автокодировщики используют SF 8 неспроста — он обеспечивает хороший баланс между степенью компрессии и качеством реконструкции изображений. Нам бы хотелось ещё увеличить SF, ведь это бы уменьшило активации нашей диффузионки, позволив учить более быстрые модели или модели генерации в большем разрешении. Однако наивные попытки увеличить SF до 16 или 32 приводят к полностью непригодным результатам, особенно на изображениях низкого разрешения. Но фундаментальное ли это ограничение?

Параллельно с SANA авторы из Nvidia выпускают другую статью под названием Deep Compression Autoencoder for Efficient High‑Resolution Diffusion Models. В ней авторы решают основные проблемы, мешающие сделать глубокие автокодировщики, после чего обучают целое семейство моделей и даже выкладывают их в опенсорс!

Для модели SANA выбрали основной автокодировщик с SF 32 и 32 каналами, который отлично показывает себя на разрешениях больше 1024×1024 пикселей. И он же позволяет сократить вычислительный бюджет на обучение диффузии более чем в 100 раз!

Интересно, что авторы SANA также используют decoder‑only генеративную LLM в качестве текстового энкодера и траснформерную архитекруту денойзера. Последнее примечательно ещё и тем, что благодаря низким пространственным размерам латентов, они могут позволить себе использовать патчи размера 1. Это значит, трансформерный денойзер учится исключительно генеративной задаче, а не даунсемплингу и генерации, как это обычно происходит.

Модель получилась классная, хотя шума вокруг неё случилось не много. Выкладывание Deep Compression Autoencoder в opensource — большое дело, поскольку это открывает пространство для исследований и обучения моделей (особенно для тех, кто не обладает большими вычислительными бюджетами).

Закрытые модели: Ideogram, Recraft, MidJourney, FLUX

Два года назад неоспоримым лидером среди платных проприетарных моделей была MidJourney. По сути, это был главный инструмент для профессионалов и техноэнтузиастов, готовых заплатить за лучшее качество генерации.

Сегодня MidJourney всего лишь одна из многих топовых моделей, предлагаемых по подписке, про внутреннюю реализацию которых мы ничего не знаем. Например, недавно появилось приложение Ideogram, позиционирующее себя как Instagram с генеративным контентом. Или Recraft, захватывающий рынок дизайнеров за счёт внедрения канвасов и качественной генерации текста.

А ещё выходцы из Stability.ai поняли ошибку предыдущей компании — чистый opensource монетизировать трудно. Поэтому свою новую модель FLUX они сделали открытой только частично.

Что же мы всё‑таки знаем про эти модели? На самом деле не очень много. Например, у Recraft есть блог, в котором они аккуратно и без лишних технических деталей повествуют о том, как им удалось добиться такого высокого качества генерации текстов. FLUX выложили базовую версию своих денойзера и автокодировщика — они оказались весьма годными, однако мы не знаем деталей того, как они этого добились.

Про остальные модели мы знаем ещё меньше: в основном это высокоуровневые слухи о том, что все учат большие латентные модели, потихоньку переходят на трансформеры и постановки диффузии Rectified Flow и Flow Matching.

За последние два года в мире генерации text‑to‑image произошло множество интересных событий. Сегодня мы остановились только на самых ярких — а сколько ещё осталось непокрытым! За рамками нашего обсуждения остались сотни научных работ, предлагающих хитроумные трюки для улучшения тех или иных аспектов моделей. Сами text‑to‑image‑модели всё больше используют для смежных задач: редактирования изображений, inpaining и outpainting, в качестве инициализации для text‑to‑video и text‑to-3d‑моделей, а также много чего другого.

Если не всё написанное в этом тексте было понятным, никогда не поздно наверстать упущенное. Для этого мы, разработчики и исследователи Яндекса, подготовили и выложили в открытый доступ бесплатный интенсив ШАДа по теории и практике диффузионных моделей — CV Week. Если у вас появятся вопросы, то смело задавайте их в комментариях.

_{*Компания Meta признана экстремистской организацией, а её продукты, Facebook и Instagram, запрещены на территории РФ}