File size: 6,025 Bytes
b4f5f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | # Архитектура
Три части, каждая со своей задачей. Всё, что ниже, проверено запуском.
## 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 в лоссе, кропы одной картинки вместо набора.
|