# Архитектура

Три части, каждая со своей задачей. Всё, что ниже, проверено запуском.

## 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 в лоссе, кропы одной картинки вместо набора.