Как решить проблему памяти с MiniMax Sparse Attention

Itinai.com it company office background blured chaos 50 v 7b8006c7 4530 46ce 8e2f 40bbc769a42e 2

Что такое MiniMax Sparse Attention (MSA)

MiniMax Sparse Attention (MSA) — это метод разреженного внимания, построенный поверх Grouped Query Attention (GQA). Он решает одну главную проблему современных трансформеров: квадратичную зависимость вычислений от длины контекста. Вместо того чтобы просчитывать softmax‑весы по всему входу, MSA сначала выбирает лишь небольшие блоки токенов, а затем применяет точный softmax только к ним.

Двухэтапный процесс: Index‑Branch и Main‑Branch

В MSA внимание делится на два шага.

Index‑Branch — определяет, какие блоки key‑value следует прочитать. Выбор происходит на уровне блоков (по умолчанию 128 токенов), а каждый запрос хранит k = 16 блоков. Таким образом, бюджет токенов для каждого запроса фиксирован: k·Bₖ = 2 048.
Main‑Branch — выполняет обычный скалярный dot‑product softmax, но лишь над выбранными токенами.

Блоки выбираются независимо внутри каждой группы GQA, но совместно для всех запросов внутри группы. Это позволяет разным группам «смотреть» в разные дальние части последовательности, сохраняя при этом чтение KV‑пар подряд, что удобно для памяти GPU.

Как работает индексатор

Index‑Branch добавляет два простых проекционных слоя к стандартному GQA‑слою: один для index query и один общий index key. Сначала получаются оценки видимых токенов, затем они агрегируются по блокам (max‑pool). Оператор Top‑k выбирает k самых релевантных блоков, причём блок, содержащий текущий запрос, обязательно включён — так мы гарантируем, что локальный контекст не исчезнет.

Обучение разреженному вниманию

Выбор Top‑k не дифференцируем, поэтому обычный кросс‑энтропийный loss не может напрямую обучать индексатор. MiniMax решает это при помощи KL‑alignment loss: распределение соседних весов в Index‑Branch выравнивается с распределением Main‑Branch (учитель — усреднённые Softmax‑весы). Чтобы обучение было стабильным, используется:

Gradient Detach — запрещает обратный поток от KL‑потери в основной backbone.
Indexer Warmup — первые несколько итераций работают с полным вниманием, позволяя индексатору «подучиться» без риска плохих роутов.

Два пути к разреженной модели

MSA‑PT — обучение с нуля после 40 B токенов Warmup.
MSA‑CPT — преобразование уже обученного плотного GQA‑чекпоинта (2.6 T токенов), затем продолжение обучения ещё 400 B токенов с Warmup.

Кернел‑оптимизации

Алгоритм без сопутствующего кода не даст ускорения, поэтому MiniMax выпустил два GPU‑кернела:

Exp‑free Top‑k — использует сохраняемый порядок softmax‑оценок, обходя дорогостоящие exp‑и‑sum и достигая до 5.1× ускорения над torch.topk при контексте 128 K токенов.
KV‑outer sparse attention — перебирает KV‑блоки вместо запросов, повышая арифметическую интенсивность. Кернел упаковывает ⌈128/G⌉ позиций запросов в одну 128×128 матричную операцию.

Код доступен в репозитории GitHub и поддерживает BF16, FP8, а также более «лёгкие» форматы.

Сравнение с другими разреженными подходами

Главные отличия MSA:

Выбор блоков уровня GQA‑группы (не по токену, а по блоку — 128 токенов).
Top‑k применяется отдельно к каждой группе, но блоки остаются непрерывными, что упрощает кеширование.
Обучаемый индексатор (KL‑align), а не параметр‑free «скользящее окно».

По качеству результаты находятся в совокупности с полным вниманием. На тесте MMLU разреженные модели дают 67.0 % (MSA‑PT) и 67.2 % (MSA‑CPT) против 68.6 % у плотного baseline. На длинном контексте (128 K) MSA‑CPT почти не отстаёт от полного внимания.

Практические выводы

Если ваш проект работает с последовательностями > 8 K токенов и вы чувствуете, что O(N²)‑затраты влекут за собой дорогостоящие вычисления, MSA позволяет зафиксировать бюджет в ~2 K токенов независимо от длины.
Для GPU SM100 (H100) доступен готовый кернел, который легко подключить к PyTorch через torch.utils.cpp_extension.
Не забывайте про Warmup — без него индексатор может «запутаться» и ухудшить точность.
Если у вас уже есть обученный dense‑GQA‑модель, путь MSA‑CPT позволит конвертировать её без переобучения с нуля.