Ra / docs /ARCH.md
recoilme's picture
Ra fp16 (step 212000)
b4f5f26 verified
|
Raw History Blame Contribute Delete
6.03 kB

Архитектура

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

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