text-to-image · 177.3M параметров · одна RTX 4080

UDiT — пирамидальный U-трансформер

Ширина растёт к низу: 192 → 512 → 1024 → 1536. Свёртки на мелком уровне, внимание на грубом, текст на всех уровнях. Декодер готового автоэнкодера работает хвостом-рендерером, поэтому тело считает 32 числа на клетку 8×8, а не пиксели.

177.3Mпараметров всего

97 мсшаг обучения · bs8 · 256²

20.4GFLOP forward у тела

0.647детали (hf) против AE 0.859

8.1 ГБпик памяти из 16

оригинал, реконструкция AE и две генерации

слева направо: оригинал · реконструкция AE (потолок 0.859 hf) · две генерации модели на 30k шагов с разными сидами

тело · 143.3M

пирамида 192/512/1024/1536, поток в латентном пространстве

текст · 10.0M

Ettin-150M заморожен, обучается только выбор слоёв

хвост · 23.9M

декодер AE, дообучается как flow по пикселям

UDiT · архитектура1 / 12

три части

Модель не считает пиксели — только 32 числа на клетку 8×8

Латент вместо пикселей

Один токен = 32 числа на клетку 8×8. При 256² это 1024 токена против 196 608 пикселей.

Декодер уже обучен

Хвост наследует декодер AE: телу не нужно учиться рисовать текстуру с нуля.

Текст не экономим

Кросс-внимание в каждом блоке до самого мелкого уровня: +2.5M, «switch» 0.43 → 0.78.

UDiT · три части2 / 12

почему пирамида

Плоский стек ширины 1024 платит за каждый токен по полной

метрика	плоский 1024	пирамида	что это значит
GFLOP на forward (тело)	232	20.4	в 11 раз меньше арифметики при том же числе параметров
шаг обучения, bs8 · 256²	196 мс	97 мс	вдвое больше шагов за то же время
детали: hf(gen)/hf(AE)	21%	51–75%	и это при 3–6 кратном меньшем числе шагов обучения

Стоимость слоя ≈ токены × ширина². Мы держим ширину там, где токенов мало, и снимаем с мелкого уровня 28-кратный перерасход — эти же деньги уходят в глубину и разрешение.

UDiT · сравнение3 / 12

лестница

Ширина растёт там, где токенов мало — стоимость уровней сбалансирована

уровень	токенов 256²	токенов 512²	ширина	блоки	механика
L0 ÷1	1024	4096	192	6 enc / 4 dec + self-attn	свёртки + текст
L1 ÷2	256	1024	512	4 / 3	трансформер
L2 ÷4	64	256	1024	3 / 2	трансформер
L3 ÷8	16	64	1536	2 (боттлнек)	трансформер

энкодер поднимается по лестнице, декодер спускается; на каждом уровне прибавляется скип

UDiT · бюджет уровней4 / 12

роли уровней

Каждый уровень отвечает за свой масштаб — как у художника

L3 · «прищур»

Вся картинка — 16 токенов. Полное внимание: все видят всех. Здесь решается композиция и симметрия.

L2 · крупные формы

Тело, крылья, общая раскраска. Ширина 1024 — запас «смысла» на грубое описание.

L1 · средние детали

Форма крыла, крупный рисунок, где тёмное и где светлое. Ширина 512.

L0 · текстура

Жилки, кромки, крапинки. Здесь свёртка 3×3, а не внимание: текстура локальна, а свёртка в 20 раз дешевле.

Внимание на грубом уровне стоит копейки: 16 токенов × 16 = 256 пар. Внимание на мелком — 1024 × 1024 = миллион пар. Поэтому «прищур» мы делаем вниманием, а текстуру — свёрткой.

UDiT · роли уровней5 / 12

путь и скипы

Энкодер поднимается к смыслу, декодер возвращается к деталям — помня эскиз

Почему аддитивные

Ширина входа и выхода уровня совпадает, проекции не нужны — так делают U-Net’ы SD.

Почему pool, а не stride

Среднее 2×2 плюс линейный слой: сжатие учится и не ломает фазу позиционного кодирования.

Что это даёт

Тонкая структура строится на своём уровне, а не вытягивается вниманием с одного масштаба.

UDiT · путь и скипы6 / 12

текст

Подпись не «берётся из последнего слоя» — модель сама выбирает, какие слои читать

SliceMixer · 10.04M обучаемых · идея из moss/adapter (0.8B + 158M голова вместо 9B-энкодера)

Два входа текста в каждый блок

кросс-внимание: каждый токен-клетка спрашивает «что в подписи про моё место?»
модуляция: общий вектор «t + усреднённый текст» настраивает shift/scale/gate фильтров

что сделали	цена	эффект
кросс-внимание и на L0	+2.5M	+9 мс
«switch» (один шум, два капшена)	—	0.43 → 0.78
детали (hf) не пострадали	—	0.44 → 0.65

«switch» — насколько латент меняется от смены подписи при том же шуме

UDiT · текст7 / 12

хвост и сэмплинг

Тело ведёт латент, готовый декодер рисует пиксели

на каком латенте учить хвост	что видит рендерер	что происходит
только z0 (чистый латент)	идеальный латент	на сэмплинге встречает ошибку тела, к которой не привык
только предсказание тела	мусор в начале обучения	учится «исправлять мусор» и разучивается рисовать: 29.6 → 21.3 дБ
z0 + (ошибка тела).detach()	реальную ошибку	видит промах, но градиент через него не идёт; 24.7–25.9 дБ стабильно

Сэмпл: 8 шагов тела + 4 шага декодера ≈ 35 мс на 4080 при 256². CFG 1.5 — выше 2 при недоученной безусловной ветви ухудшает латент (замерено 0.19 → 0.28). Метрика tail(z0) в каждом логе ловит деградацию рендерера.

UDiT · хвост8 / 12

результат на узком домене

Бабочки: детали растут с шагами и не насыщаются

шаги	hf / AE	mse_lat	high	switch
4 000	0.446 · 52%	0.21	0.42	0.43
10 000	0.436 · 51%	0.164	0.32	0.68
20 000	0.540 · 63%	0.107	0.19	0.78
30 000	0.647 · 75%	0.072	0.114	0.78

hf — энергия высоких частот (лапласиан). AE-реконструкция того же вида: 0.859, оригинал 0.913. «high» — ошибка на высоком шуме, то есть «какую особь рисовать»; она упала в 3.7 раза.

Три сида дают три разные особи: 0.642 / 0.646 / 0.654 — модель не сколлапсировала в один мод.

Четыре вида бабочек, сгенерированные моделью

генерации на 30k шагов: Delias, Cepora, Troides, Catasticta — вид задаётся подписью из 16 токенов

UDiT · результат9 / 12

цена

177.3M параметров, 97 мс на шаг, 8.1 ГБ — обучение на одной карте

Параметры по частям

Арифметика на forward

Скорость и память

Почему хватает 16 ГБ

Память ест не модель (2.5 ГБ весов и состояний), а активации: они ~ токены × ширина × блоки. Пирамида держит широкие слои на 16–64 токенах, поэтому активации маленькие.

Где запас

Шаг упирается в память, а не в счёт. Запас есть на батч 16 (14.8 ГБ с хвостом) и на 512² с батчем 1–2 — то есть на разрешение, где деталь физически возможна.

UDiT · цена10 / 12

следующий прогон

FLUX-Reason: 2 эпохи по 1.5M картинок, поток без кэша

параметр	значение
датасет	LucasFang/FLUX-Reason-6M, часть Aesthetics-Part01 (415 шардов, ~250 ГБ)
разрешение / батч	256² · bs16
подпись	caption_entity → caption_composition (40–400 симв.), иначе caption_detail
бюджет	2 эпохи = 187 500 шагов ≈ 12.5 ч
оптимизатор	AdamW8bit, прогрев 500, cosine до 10%
дропауты	подпись 10% (CFG) · слова 5%
старт	веса бабочек (30k шагов)

Что уже видно (шаг 800)

«high» здесь в 6 раз выше, чем на бабочках: 1.5M разных сцен, «какую именно рисовать» по подписи не угадать — это нормальный старт, а не поломка.

Смотреть

tmp/flux/log.csv · tmp/flux/samples/latest.png (лиса, стеклянный город, неон, импрессионист)

Остановить / продолжить

kill процесса · тот же запуск с --resume --out tmp/flux

Честно про «2 эпохи»

эпоха по Part01 = 1.5M картинок; все 4 части ≈ 1 ТБ, а на диске 352 ГБ

UDiT · FLUX-Reason11 / 12

правила

Всё, что ниже — вывод из замеров, а не из вкуса

1 Ширина растёт к низу. Плоский 1024 на всех уровнях: 196 мс и hf 0.18. Пирамида: 97 мс и 0.54.

2 Текст не экономим. Кросс-внимание на самом мелком уровне: +2.5M и +9 мс, «switch» 0.43 → 0.78.

3 Хвост учится на z0 + ошибка тела. Только z0 — не видит реальных ошибок; только предсказания тела — деградация 29.6 → 21.3 дБ.

4 Бакеты кратны 64. Три даунсэмпла требуют латента, кратного 8. Аспект — центр-кроп, не растяжение.

5 Метрики вместо впечатлений. mse_lat (общий, low, high), mse_pix, switch, tail(z0), hf(gen)/hf(AE), skip.

6 Никаких костылей. Фиксированный шум, клип латента, «mix» латентов — выброшены: они прячут то, что мы измеряем.

Дальше

дожать FLUX-Reason до 2 эпох по Part01

Потом

люди: 5749 фото, 256² — лица и кисти

Затем

512² на людях: деталь физически разрешима

Только потом

ёмкость: ширина боттлнека и число блоков L0

UDiT · правила12 / 12