разборllmархитектураssmeva

Mamba и Jamba: что будет после трансформеров (или вместе с ними)

· 3 мин · BadRabbit

Последний разбор перед завтрашним большим постом про Еву. В её реестре недавно появилась Jamba-1.7-Mini, и это повод разобраться, чем она отличается от обычных моделей. Основа — визуальный гайд Мартена Гроотендорста по Mamba и анонс Jamba от AI21.

Две крайности

  • Трансформер при генерации каждого токена смотрит на все предыдущие. Это мощно, но дорого: вычисления растут с длиной, а KV-кеш растёт вместе с контекстом.
  • RNN хранит фиксированное «состояние» и обновляет его на каждом токене. Это дёшево и линейно, но RNN забывает далёкий контекст и плохо параллелится при обучении.

Модели пространства состояний

SSM пришли из теории управления. Скрытое состояние h меняется по правилу с четырьмя матрицами:

h' = A·h + B·x      # как меняется состояние под действием входа
y  = C·h + D·x      # что из состояния выдать наружу
  • A — как состояние живёт само по себе;
  • B — как вход влияет на состояние;
  • C — как состояние превращается в выход;
  • D — прямой проброс входа.

Фокус в том, что такую модель можно считать двумя способами: как свёртку (параллельно, быстро при обучении) и как рекуррентность (по шагу, быстро при генерации).

Что добавила Mamba: избирательность

У классических SSM матрицы одинаковы для всех токенов: модель одинаково «впитывает» и важное, и мусор. Mamba сделала B, C и шаг дискретизации Δ зависящими от входа. Теперь модель может решить: этот токен важен, запомню, а этот пропущу. Это и есть селективность.

Цена — свёрточный трюк больше не работает. Его заменили параллельным сканированием (используя ассоциативность операции) и аппаратно-оптимизированным ядром, которое, как и FlashAttention, держит промежуточные данные в быстрой памяти и пересчитывает их вместо хранения.

Jamba: гибрид

AI21 решила не выбирать, а смешать:

  • слои Mamba и трансформера чередуются: один слой внимания на восемь слоёв;
  • поверх MoE: 16 экспертов, на токен работают два;
  • 52 миллиарда параметров всего, 12 активных;
  • контекст 256 тысяч токенов, до 140 тысяч помещается на одну видеокарту на 80 ГБ;
  • на длинном контексте пропускная способность в 3 раза выше, чем у Mixtral 8x7B.

Внимание даёт точность «найти иголку», Mamba даёт дешёвый длинный контекст, MoE даёт ёмкость.

Мой комментарий

Для Евы гибридные модели интересны тем, что их состояние почти не растёт с контекстом: у слоёв Mamba нет KV-кеша. На 8 ГБ это значит, что длинный документ можно скормить модели, не выгоняя соседей из VRAM. Но с MoE-частью та же проблема: всех экспертов надо где-то держать. В реестре Евы Jamba записана с бюджетом 7000 МБ VRAM, то есть почти вся моя карта. Рядом с ней не поместится никто, поэтому она кандидат на гибридный режим CPU+GPU, и как раз сейчас бенчмарк Евы меряет, сколько стоит такой компромисс.