210M параметров и одна GPU: как обучить диффузионный text-to-image с нуля за 3,5 дня

Что именно обучили

В блоге Hugging Face вышел разбор Ивана Михненкова: он обучил модель генерации изображений по тексту с нуля, чтобы разобраться в современных подходах не по статьям, а на практике. Результат — tinydit, диффузионный трансформер на 210 млн параметров, обученный за 3,5 дня на одной GPU RTX PRO 6000 на 4,2 млн изображений разрешением 256².

Обучаемых частей ровно столько, сколько написано автором: сам трансформер, рецепт обучения, пайплайн данных и оценка оригинальны, а опирается модель на две замороженные предобученные компоненты — автоэнкодер FLUX.2 и текстовый энкодер flan-t5-base. Веса tinydit-256, код, каждая проектная развилка с источниками и браузерное демо опубликованы открыто.

Данные решили исход ещё до старта обучения

Первая попытка масштабирования несколькими месяцами ранее использовала веб-краул примерно на миллион изображений с перегенерированными подписями — и модель вышла хуже, чем обучение на одном COCO: качество картинок из интернета случайное, а подписи им не соответствовали. Поэтому во второй раз автор сначала достал по двадцать реальных строк из каждого набора-кандидата и посмотрел на них так, как их увидит модель 256².

Итоговая смесь: стоковые фото Pexels в полном кадре 640 px — 2,8 млн изображений и 60% батчей; части FLUX-Reason-6M, отфильтрованные по оценке качества, — 1,2 млн и 25%; COCO 2017 train — 118 тыс. и 15%.

Три детали дали реальный эффект. У каждой картинки есть длинная подпись (до 128 токенов T5) и короткая, и обучение берёт длинную в 50% случаев, короткую в 40%, а пустую — в 10% для classifier-free guidance, поскольку только длинные подписи портят короткие промпты, а только короткие — композицию. Синтетический набор ограничен четвертью выборки: модель, которой скармливают в основном сгенерированные изображения, перенимает «почерк» генератора. И ни одно изображение не обрезается в квадрат — наивный фильтр «оставить только близкие к квадрату» выбросил бы 94% COCO и 99% Pexels, ведь настоящая фотография это 4:3 и 3:2.

Вместо этого модель учится на пяти наборах соотношений сторон примерно по 256 токенов — 256×256, 288×224, 224×288, 320×208 и 208×320, каждый батч берётся из одного набора с самого первого шага, а средняя обрезка составляет 3–5% длинной стороны. Позиция задаётся двумерными rotary-эмбеддингами, поэтому трансформер не нужно менять под каждую форму, и он обобщается даже на кадр 448×256, которого не видел.

Архитектура: известный рецепт плюс registers

В самом трансформере ничего не изобреталось — и это осознанный выбор: взят блок, который используют все актуальные модели, но написанный с нуля, чтобы понимать каждую строку. Это DiT с cross-attention на текстовые токены (линия PixArt и SANA, тогда как FLUX и SD3 применяют joint attention), 2D RoPE, QK-нормализация, SwiGLU и adaLN-single: одна общая сеть модуляции плюс небольшая таблица на блок вместо отдельного линейного слоя на каждый блок. Одна эта замена освободила 27% прежнего числа параметров без затрат вычислений, и их автор вложил в глубину — ширина 896, 16 блоков, 14 голов.

Единственное заимствование из свежих исследований — registers. Шестнадцать обучаемых токенов присоединяются к последовательности изображения на третьем блоке и отбрасываются перед выходом, а каждый блок cross-attention получает две обучаемые «пустые» пары ключ/значение, к которым всегда можно обратиться вниманием. Идея пришла из работы Vision Transformers Need Registers и статей 2026 года, перенёсших её в диффузионные трансформеры.

Эффект измеримый: на среднем уровне шума в среднем блоке две пустые пары получают около 90% всей массы cross-attention, токен конца последовательности, который раньше служил стоком, — 4%, а содержательные слова — оставшиеся проценты, уже привязанные к своим объектам. Векторы registers к средним блокам вырастают до 4–13 норм токена изображения, что и есть подпись токенов, работающих как черновик, а каждый патч изображения отправляет на них 8–18% своего self-attention.

Рецепт: сдвиг таймстепов, две дополнительные функции потерь и compile

Обучение идёт на rectified flow со сдвигом таймстепов 2,8: модель предсказывает скорость перехода от шума к изображению, а таймстепы сэмплируются лог-нормально, как в SD3. Сдвиг нужен из-за 32-канального латента FLUX.2 — изображение остаётся узнаваемым при таком шуме, который уничтожил бы 4-канальный латент SD, поэтому для одинаково сложных задач нужно больше шума: со сдвигом половина обучающих примеров имеет больше 74% шума, и тот же сдвиг изгибает 20-шаговое расписание сэмплинга так, что 15 шагов попадают именно в эту область, где решается композиция.

Две вспомогательные функции потерь: косинусный член на направление скорости из LightningDiT — при высоком шуме MSE определяется величиной, тогда как изображение задаёт направление, — и dispersive loss на признаках пятого блока (Wang и He, 2025), который расталкивает внутренние представления разных примеров без внешнего энкодера и не даёт представлению схлопнуться.

Отдельно компилировался шаг обучения: Dynamo трассирует forward, AOTAutograd выводит backward, Inductor сливает мелкие поэлементные операции вокруг умножений матриц — на реальном шаге это дало ускорение в 2,4 раза и вдвое меньше памяти, а ограничение в виде статических форм закрывается тем, что каждая из пяти форм получает свой граф.

Остальные детали: AdamW с 2e-4 и betas (0,9; 0,95), fp32-мастер-веса под bf16 autocast, EMA весов с decay 0,9999 (горизонт 10 000 шагов — обычные 0,999 для прогона в 400 тыс. шагов слишком коротки), батч 256 и линейное снижение learning rate до нуля в последней четверти прогона; 0,76 секунды на шаг, 400 тыс. шагов за 3,5 дня.

Лосс — не сигнал качества

Функция потерь flow-matching за весь прогон упала с 0,805 до 0,754, хотя изображения прошли путь от пятен до фотографий. Причина структурная: одному и тому же зашумлённому входу соответствует множество пар «изображение и шум», лучший возможный прогноз — их среднее, а разброс вокруг среднего никто не может убрать; при высоком шуме, где проходит половина обучения, эта неустранимая часть и составляет почти весь лосс.

Поэтому лосс использовался как сигнал здоровья — обучение и отложенная выборка совпадали до третьего знака на протяжении 24 эпох, то есть запоминания не было, — а качество отслеживалось каждые 10 тыс. шагов на замороженном наборе отложенных подписей: FID и FD-DINOv2 против реальных изображений, точность по объектам (детектор, обученный на COCO, ищет запрошенный класс в подписи «a photo of a {class}» для всех 80 классов), а также модели человеческих предпочтений PickScore и HPSv2.1.

FID упал с 33,7 на 10 тыс. шагов до 28,1 на 100 тыс. и 27,0 на 400 тыс.; FD-DINOv2 — с 570 до 274 и 218; точность по объектам — с 65% до 88% и 90%; PickScore — с 19,5 до 20,6 и 20,9; HPSv2.1 — с 0,199 до 0,238 и 0,254. FID и точность по объектам вышли на плато около 300 тыс. шагов, а FD-DINOv2 и оценки предпочтений улучшались до самого конца, и снижение learning rate на последних 100 тыс. шагов добавило ещё 5% по FD-DINOv2.

Что получается, а что пока нет

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

На 2 456 отложенных промптах переход с 20 шагов на 50 стоит 0,4 FID и не даёт ничего по оценкам предпочтений, а отказ от сдвига таймстепов обходится дороже, чем удвоение числа шагов; композиция фиксируется уже после третьего шага сэмплинга, и всё дальнейшее — только повышение резкости.

Пять выводов и что дальше

Автор формулирует пять советов для тех, кто запускает такой прогон: сначала смотреть на данные в том разрешении, в котором обучаешь; согласовывать распределение таймстепов с латентом, ведь 32-канальный латент хочет сдвиг около 2,8; использовать наборы соотношений сторон с первого шага вместо разогрева на квадратной обрезке; не читать прогресс по лоссу, а выбрать две-три метрики, считаемые каждый час, и заморозить набор промптов; компилировать шаг обучения.

Базовая модель намеренно не дистиллирована и использует обычный classifier-free guidance, потому что следующая фаза — обучение с подкреплением через Flow-GRPO с теми же моделями предпочтений в роли награды и списком провалов в роли целей; результат обещан вторым релизом в том же репозитории.