В декабре 2023 по ML-тусовке прокатилась волна заголовков в духе «трансформерам конец». Поводом стала статья двух исследователей — Альберта Гу и Три Дао — со скучным названием: «Mamba: моделирование линейно-временных последовательностей с использованием селективных пространств состояний». Внутри была архитектура, в которой не было механизма внимания, того самого attention, на котором держится весь современный тир-лист нейронок. И при этом она работала на длинных текстах в несколько раз быстрее трансформера, при меньшем расходе памяти.
Прошло уже много времени, так что не будет спойлером сказать, что свой трон трансформеры не потеряли. Но история на этом не закончилась, и развязка интереснее, чем «очередной хайп-трейн не взлетел».
Что не так с трансформерами
Чтобы понять, против чего бунтовала Mamba, надо разобраться с одной особенностью трансформеров.
Языковая модель читает текст по кусочкам. Как можно было догадаться, речь идет о токенах.
Важно понимать, что токен — это не просто слово. Это может быть слово, часть слова или пунктуация. Это значит, что фактический текст, который может рассмотреть модель, может быть короче, чем вы изначально ожидали. Хотя, для базового понимания, можно думать о токенах как о словах.
Когда трансформер обрабатывает текст, он использует механизм внимания, и работает следующим образом: для каждого токена модель вычисляет, насколько с ним связан каждый другой токен в тексте. Слово «он» получает, например, высокую связь с «директором» из прошлого абзаца, низкую с предлогами вокруг. Из этих связей модель и понимает, что к чему относится.
Дальше считаем, если в тексте N токенов и каждый токен оценивает связь с каждым другим, то всего таких попарных вычислений порядка N умножить на N, то есть N².
Вот в этом и есть вся проблема. Для короткого вопроса она незаметна. Но как только даешь модели книгу, длинный диалог или большой файл с кодом, квадрат начинает душить: и по времени, и по памяти.
К этому добавляется второй расход. Чтобы каждый новый токен мог свериться со всеми предыдущими, модель хранит для каждого уже обработанного токена набор векторов (их называют ключами и значениями). Этот объем растет линейно с длиной текста и занимает память видеокарты. Чем длиннее контекст, тем больше памяти уходит только на то, чтобы держать под рукой все прошлое.
Из этого и вырос вопрос, на который отвечает Mamba: обязательно ли каждому токену сверяться с каждым другим токеном?
Новые GPU в облаке Selectel от 132,18 ₽/час
Видеокарты для ресурсоемких задач — NVIDIA® H200, RTX™ 6000 Pro.
Альтернатива: одно состояние вместо внимания
Другой способ обработать последовательность известен давно. Можно не хранить все токены и не сравнивать их попарно, а вести одно бегущее состояние: некую сводку прочитанного, которую модель обновляет с приходом каждого нового токена.
Так работали рекуррентные нейросети, RNN, главный инструмент для текста до 2017 года, пока их не потеснили трансформеры. У RNN было два минуса, из-за которых от них отказались. Первый: их тяжело обучать быстро, потому что каждый шаг зависит от предыдущего и идет строго по очереди, а параллельные ядра видеокарты при этом простаивают. Второй: на длинных текстах состояние размывалось, и начало забывалось.
Mamba выросла из этой же линии. Но чтобы подход наконец заработал на тексте, авторам пришлось разобраться с обеими старыми проблемами. Сначала посмотрим, как устроено это состояние, и уже потом как его сделали умным и быстрым.
Как работает модель пространства состояний
Конструкция, на которой стоит Mamba, называется моделью пространства состояний, по-английски State Space Model или SSM. Технически скрытое состояние — это вектор чисел фиксированной длины, в котором закодировано сжатое представление всего прочитанного.
Обработка одного токена идет по шагам:
-
Берется состояние, накопленное к этому моменту;
-
Оно преобразуется (параметры A): часть информации сохраняется, часть ослабляется;
-
В состояние подмешивается новый токен (параметры B задают, насколько сильно он повлияет);
-
Из обновленного состояния вычисляется выход (параметры C определяют, что достать);
-
Переход к следующему токену.
Отсюда два прямых следствия. Первое: на каждый токен тратится фиксированная работа, не зависящая от длины текста, поэтому общая стоимость растет линейно, без квадрата. Второе: память под состояние не растет вообще, в отличие от KV-кэша (ключ, значение), который пухнет с каждым словом.
Выглядит как готовая победа. Но если бы все было так просто, подобные модели давно вытеснили бы трансформеры. Засада есть, и именно из-за нее SSM годами проигрывали на тексте.
Засада старых SSM и главный прием Mamba
Засада в параметрах A, B и C. В классических SSM они были фиксированными: одинаковыми для всех токенов и неизменными после обучения. То есть модель обрабатывала важное слово, проходной предлог и ключевое имя одним и тем же движением, с одинаковой силой.
Для текста это плохо: одно слово нужно надолго удержать, другое можно сразу отбросить, а фиксированные параметры применяют ко всему один закон.
Вот что изменила Mamba. Параметры (а с ними и сила, с которой токен впускается в состояние) вычисляются заново для каждого токена, в зависимости от того, какой токен пришел. Разница видна на двух версиях одного шага:
# Классическая SSM: один и тот же закон для всех токенов
def step_fixed(h, x, A=0.9, B=0.5):
return A * h + B * x
# Mamba: насколько впустить токен, зависит от самого токена
def step_selective(h, x):
importance = abs(x) # упрощенно: важность берем из токена
gate = importance / (1 + importance) # 0..1, считается на лету
return (1 - gate) * h + gate * x
print(step_selective(0.0, 5.0)) # 4.17 — крупный токен впущен сильно
print(step_selective(0.0, 0.1)) # 0.01 — мелкий почти проигнорированВ фиксированной версии множитель при x всегда 0.5. В селективной он зависит от самого токена: важный впускается почти целиком, шумовой отсекается. На каждом шаге модель сама решает, что взять в состояние, а что пропустить. Этот прием называется механизмом отбора.
Из-за отбора важная информация держится в состоянии долго, а шум гаснет почти сразу.
По сути это способ делать то, ради чего обычно нужно внимание — выделять значимое и отсеивать шум, — но не сравнивая все со всем за квадратичную цену, а управляя одним состоянием за линейную.
Почему это работает быстро, если внутри все еще рекуррентность
Здесь должно возникнуть возражение. Выше сказано, что RNN отложили в том числе за медленное обучение: шаги идут по очереди, видеокарта простаивает. А Mamba ведет такое же пошаговое состояние, где каждый шаг зависит от предыдущего. В чем разница.
Разница в способе вычисления. Обновление состояния обладает свойством ассоциативности: длинную цепочку шагов можно посчитать не строго по очереди, а разбить на участки, обработать их независимо и параллельно, а затем соединить промежуточные результаты. Покажу на той же схеме h = a*h + b, где каждый шаг — это пара чисел (a, b):
import functools
def combine(L, R): # объединяем два последовательных шага в один
aL, bL = L
aR, bR = R
return (aR * aL, aR * bL + bR) # тот же результат, что два шага подряд
steps = [(0.9, 1.0), (0.8, 0.0), (0.5, 2.0), (0.2, 0.0)]
h = 0.0 # обычный последовательный проход
for a, b in steps:
h = a * h + b
left = functools.reduce(combine, steps[:2]) # первую половину...
right = functools.reduce(combine, steps[2:]) # ...и вторую можно посчитать
h_par = combine(left, right)[1] # параллельно, потом соединить
print(round(h, 4), round(h_par, 4)) # 0.48 0.48 — результат одинаковыйИтог тот же, что при последовательном проходе, но половины можно считать одновременно. На обучении, когда весь текст известен заранее, это позволяет загрузить все ядра видеокарты сразу. Авторы вдобавок написали операцию так, чтобы она бережно расходовала память и держала промежуточные данные в самой быстрой части видеокарты. В результате рекуррентность, из-за которой когда-то отказались от RNN, перестала быть тормозом.
Когда модель уже обучена и просто генерирует ответ, токены и так выходят по одному, и пошаговость не мешает. При этом на каждый новый токен Mamba тратит фиксированную работу и память, а трансформер заново сверяется со всей растущей историей. Чем длиннее текст, тем больше разрыв.
Вторая версия, Mamba-2, переработала эту механику еще ближе к тому, что эффективно считается на видеокартах, и ускорила обучение в разы относительно первой. Принцип остался прежним: отбирающее состояние плюс параллельный пересчет.
Где Mamba проигрывает
У фиксированного состояния есть оборотная сторона, и она прямо вытекает из устройства.
Состояние имеет постоянный размер, а значит, вся прочитанная информация в него сжимается, и детали теряются. Прогоним поток токенов и попробуем потом найти в состоянии самый первый:
import numpy as np
A = np.array([0.9, 0.8, 0.5, 0.2]); B = np.array([0.1, 0.3, 0.6, 1.0])
h = np.zeros(4)
tokens = [5, 1, 9, 2, 7, 3, 8] # хотим потом точно вспомнить первый — 5
for x in tokens:
h = A * h + B * x
print(h.round(2)) # [2.7 6.37 7.3 8.91]
# в этих четырех числах перемешан весь текст;
# вытащить из них ровно "первый токен был 5" уже нельзяУ трансформера такой проблемы нет: он хранит все токены явно и в любой момент обращается к конкретной позиции. Поэтому на задачах точного извлечения трансформеры стабильно сильнее: найти в большом документе нужный идентификатор, дословно скопировать фрагмент, ответить на дотошный вопрос про конкретное место. Та самая квадратичная дороговизна оборачивается их сильной стороной: умение буквально сопоставлять все со всем и есть умение точно находить.
Есть и второе слабое место. Модели на основе внимания лучше схватывают новые шаблоны прямо из примеров в запросе, без дообучения. Дашь в подсказке пару образцов формата — трансформер увереннее уловит закономерность, чем чистая SSM. Причина, по-видимому, та же.
Расклад следующий. Трансформер дороже, но точен в деталях и цепок к позициям. Mamba дешевле и быстрее, особенно на длинных текстах, но хранит сжатый смысл, а не дословную картину.
Так убила или нет
Теперь, когда обе стороны на руках, можно вернуться к заголовку.
Нет, не убила. Спустя время после той статьи трансформеры по-прежнему лежат в основе большинства флагманских моделей. Причина ровно та, что разобрана выше: для универсального помощника точное извлечение и цепкость к деталям слишком важны, чтобы отказываться от них целиком.
Но пока спорили, кто кого, индустрия сделала третий ход. Если у трансформера сила в точном поиске, а у Mamba в дешевой длинной памяти, то модель можно собрать из тех и других слоев сразу. Большую часть слоев отдать Mamba, чтобы было быстро и не съедало память, а немного слоев внимания вставить точечно, там, где нужен точный поиск.
Так и сделали. Уже сегодня сложилось практическое правило: примерно один слой внимания на семь слоев Mamba. При таком соотношении из модели исчезает основная часть раздувающегося KV-кэша, а способность точно находить нужное сохраняется за счет немногих слоев внимания.
Итог получился неожиданный. Mamba шла вытеснять трансформеры, а в результате встала с ними в один ряд внутри гибридных моделей: взяла на себя основную массу слоев и тяжелую рутину, оставив немногим слоям внимания тонкую работу. Так что, вполне возможно, модель, которой вы пользовались сегодня, (еще раз, возможно) отчасти построена на Mamba, просто об этом не пишут на главном экране.
Автор: inkedsymon
