| # Архитектура |
|
|
| Три части, каждая со своей задачей. Всё, что ниже, проверено запуском. |
|
|
| ## 1. Тело: пирамидальный U-трансформер |
|
|
| Принцип: **ширина растёт там, где токенов мало.** Стоимость слоя ~ `токены × C²`, поэтому на |
| высоком разрешении держим узкие свёрточные блоки, а ширину и полное внимание — внизу. |
|
|
| | уровень | токенов при 512² | ширина | блоки | механика | |
| |---|---|---|---|---| |
| | L0 ÷1 | 4096 (512²) / 1024 (256²) | **192** | 6 enc / 4 dec + 1 self-attn с каждой стороны | depthwise 3×3 + expand 4× (ConvNeXt-стиль), кросс-внимание к тексту в каждом блоке | |
| | L1 ÷2 | 1024 / 256 | **512** | 4 enc / 3 dec | трансформер: self-attn + cross-attn к тексту + FFN | |
| | L2 ÷4 | 256 / 64 | **1024** | 3 enc / 2 dec | то же, полное внимание дёшево | |
| | L3 ÷8 | 64 / 16 | **1536** | 2 (боттлнек) | то же; широкий низ = глобальные ограничения (симметрия, композиция) | |
|
|
| Переходы и связи: |
|
|
| - **вниз**: mean-pool 2×2 + `Linear`, вверх: `Linear` + nearest 2×; |
| - **скипы аддитивные** (SD-стиль): выходы `e0/e1/e2` прибавляются к декодеру того же уровня; |
| ширина на входе и выходе уровня совпадает, поэтому сложение без проекций; |
| - **RoPE на каждом уровне**, фаза в **латентных пикселях** (L1 сдвинут вдвое, L2 вчетверо), поэтому |
| одна фаза означает одно расстояние независимо от уровня; |
| - **модуляция**: общий вектор `t + pooled-текст` → `(shift, scale, gate)` на уровень; в conv-блоках |
| L0 пять чанков (три на conv-путь, два на вход кросс-внимания); |
| - **zero-init**: выход conv-блоков и FFN обнулён, поэтому новые ветки не ломают старт. |
|
|
| Размеры: тело **143.3M**. Шаг обучения 97 мс при bs8 на 256² (8.1 ГБ). Для сравнения, прежний |
| плоский стек ширины 1024 на всех уровнях: 196 мс и hf 0.18 против 0.54 у пирамиды (архитектура |
| удалена в `tmp/archive`). |
|
|
| ## 2. Текст |
|
|
| - **Ettin-150M** (`text_encoder/`): ModernBERT, 768, 22 слоя, byte-level BPE, MIT, заморожен. |
| - **SliceMixer** (`src/text.py`, 10.04M обучаемых): 7 срезов (выходы после глобальных слоёв |
| 2, 5, 8, 11, 14, 17, 20) → пер-срезный LayerNorm (нормы растут 52 → 548, без нормировки ранние |
| слои не влияют) → 2 блока attention **по оси срезов** → `Linear(7→1)` → zero-init добавка к |
| среднему по срезам. Идея из `moss/adapter`, где 0.8B + 158M голова заменили 9B-энкодер. |
| - Вход в тело: кросс-внимание на **всех** уровнях (общие ключи/значения на уровень) плюс |
| pooled-текст в модуляции. |
|
|
| ## 3. Хвост |
|
|
| Декодер AE без изменений архитектуры: патчи 8×8, 6 блоков @512, adaLN с низкоранговыми дельтами, |
| long skip из блока n−2, финальный `PixelShuffle + conv3×3` ×3. 23.87M. |
|
|
| Обучается как обычный flow по пикселям: случайные `t`, случайный шум, цель — картинка, **условие — |
| `z0 + (ошибка тела).detach()`**. Условие без градиента: рендерер видит реальную ошибку тела, но не |
| может «научиться исправлять мусор» (именно это раньше увело его с 29.6 на 21.3 дБ). |
|
|
| ## 4. Сэмплинг |
|
|
| ``` |
| z ~ N(0,1) |
| 4–8 шагов Эйлера по телу: x0 = тело(z, t, текст); z ← z − (z − x0)/(t·N) |
| CFG 1.5 (выше 2 при недообученной безусловной ветви ухудшает латент) |
| 1 + 3 шага декодера AE: изображение |
| ``` |
|
|
| При 256² сэмпл стоит ~35 мс на 4800-й карте; при 512² — примерно вчетверо дороже. |
|
|
| ## 5. Бюджет |
|
|
| | часть | параметры | |
| |---|---| |
| | тело | 143.3M | |
| | выбор слоёв текста | 10.04M | |
| | хвост | 23.87M | |
| | **итого** | **177.3M** | |
|
|
| ## 6. Что зафиксировано и почему |
|
|
| - **x0, а не velocity.** Смена параметризации тела не затрагивает декодер (это отдельный flow со |
| своей целью), но требует переучивания тела с нуля — поэтому оставлено как есть. |
| - **Текст не экономим.** Кросс-внимание на L0 добавило 2.5M параметров и 9 мс, `switch` вырос |
| 0.43 → 0.78 при том же уровне деталей. |
| - **Никаких фиксированных шумов и клипов в обучении** — только стандартный flow-matching. |
| - **Удалено как тупиковое:** плоский 1024-стек, резидуальная U-ветка, Mix-FFN поверх плоского |
| стека, velocity-переключение посреди обучения, LPIPS в лоссе, кропы одной картинки вместо набора. |
|
|
