MiniMax представила метод разреженного внимания MSA для длинных контекстов

Компания MiniMax представила метод разреженного внимания MSA (MiniMax Sparse Attention), построенный на основе Grouped Query Attention (GQA). Подход нацелен на снижение вычислительных затрат, возникающих из‑за квадратичной сложности softmax‑внимания при длинных контекстах. Технологию протестировали внутри модели Mixture-of-Experts с 109B параметров, обученной на нативных мультимодальных данных. Также опубликован open-source inference‑kernel и представлен продакшн‑модель MiniMax-M3.

MSA разделяет механизм внимания на две стадии: Index Branch и Main Branch. Первая ветка определяет, какие блоки ключей‑значений должна читать каждая query, после чего вторая выполняет точное softmax‑внимание только по выбранным блокам. Отбор происходит на уровне блоков, а не отдельных токенов. Базовый размер блока составляет B k = 128 токенов, при этом каждая query и группа GQA сохраняют k = 16 блоков, всего kB k = 2,048 key-value токенов на запрос. В отличие от плотного GQA, где вычисления масштабируются как O(N), MSA имеет сложность O(kB k), остающуюся постоянной при росте длины контекста. По данным MiniMax, метод показывает сопоставимые результаты с GQA на downstream‑бенчмарках и снижает вычисления внимания на токен в 28.4× при контексте 1M.

Для обучения используется KL alignment loss, поскольку операция Top-k недифференцируема. Распределение Index Branch согласуется с паттерном внимания Main Branch, выступающим «учителем». В схеме применяются три стабилизирующих механизма: Gradient Detach, Indexer Warmup и принудительное включение локального блока. MSA поддерживает два режима обучения: MSA-PT, обучение с нуля после 40B-token indexer warmup, и MSA-CPT, преобразование dense GQA checkpoint, обученного на 2.6T токенов, с последующим дообучением на 400B токенов, включая 40B токенов warmup.