Mamba: архитектура, которая шла убивать трансформеры

В сфере искусственного интеллекта произошло заметное событие. inkedsymon только что Mamba: архитектура, которая шла убивать трансформеры Средний 7 мин 8 Блог компании Selectel Машинное обучение * Искусственный интеллект Natural Language Processing * Научно-популярное Обзор В декабре 2023 по ML-тусовке прокатилась волна заголовков в духе «трансформерам конец». Поводом стала статья двух исследователей — Альберта Гу и Три Дао — со скучным названием: «Mamba: моделирование линейно-временных последовательностей с использованием селективных пространств состояний». Внутри была архитектура, в которой не было механизма внимания, того самого attention, на котором держится весь современный тир-лист нейронок.

И при этом она работала на длинных текстах в несколько раз быстрее трансформера, при меньшем расходе памяти. Прошло уже много времени, так что не будет спойлером сказать, что свой трон трансформеры не потеряли. Но история на этом не закончилась, и развязка интереснее, чем «очередной хайп-трейн не взлетел».

Технические детали

Что не так с трансформерамиЧтобы понять, против чего бунтовала Mamba, надо разобраться с одной особенностью трансформеров. Языковая модель читает текст по кусочкам. Как можно было догадаться, речь идет о токенах.

Как модель раскладывает текст на токены. Важно понимать, что токен — это не просто слово. Это может быть слово, часть слова или пунктуация.

Это значит, что фактический текст, который может рассмотреть модель, может быть короче, чем вы изначально ожидали. Хотя, для базового понимания, можно думать о токенах как о словах. Когда трансформер обрабатывает текст, он использует механизм внимания, и работает следующим образом: для каждого токена модель вычисляет, насколько с ним связан каждый другой токен в тексте.

Отраслевые последствия

Слово «он» получает, например, высокую связь с «директором» из прошлого абзаца, низкую с предлогами вокруг. Из этих связей модель и понимает, что к чему относится. Дальше считаем, если в тексте N токенов и каждый токен оценивает связь с каждым другим, то всего таких попарных вычислений порядка N умножить на N, то есть N².

Вот в этом и есть вся проблема. Для короткого вопроса она незаметна. Но как только даешь модели книгу, длинный диалог или большой файл с кодом, квадрат начинает душить: и по времени, и по памяти.

К этому добавляется второй расход. Чтобы каждый новый токен мог свериться со всеми предыдущими, модель хранит для каждого уже обработанного токена набор векторов (их называют ключами и значениями). Этот объем растет линейно с длиной текста и занимает память видеокарты.

Событие, по словам экспертов, усилит конкуренцию в сфере ИИ.