MSA снижает стоимость инференса для длинного контекста в 28 раз!

MiniMax представили свой механизм MSA (MiniMax Sparse Attention), который снижает вычислительные затраты на токен для внимания в 28,4 раза при контексте в один миллион токенов, сохраняя при этом паритет с обычным вниманием на стандартных бенчмарках. Они внедрили его в свою модель MiniMax-M3 и выпустили инференс-ядро под лицензией MIT.

Вот какую проблему они решают: компонент «внимания» во всех современных языковых моделях требует сравнения каждого токена с каждым другим токеном, что даёт квадратичную сложность относительно длины входных данных. Если мы хотим обработать целую кодовую базу или книгу (то есть длинный контекст), то эти квадратичные вычисления становятся слишком медленными и слишком дорогими. Разреженное внимание призвано избегать ненужных сравнений, но нам хочется добиться этого без потери интеллектуальности модели.

1Как работает приём с двумя ветками

MSA разбивает задачу на два этапа. Сначала «индексная ветка» вычисляет оценки для частей входных данных (блоков по 128 токенов) и выбирает 16 наиболее релевантных токенов для каждого запроса, при этом блок самого запроса всегда остаётся локальным. Затем «основная ветка» выполняет точные вычисления внимания, но только между выбранными блоками, а не по всему входу целиком. Теперь каждый токен «смотрит» примерно на 2048 других токенов независимо от длины документа — этим и объясняется, почему выигрыш настолько велик при длинном контексте: в 28,4 раза меньше вычислений на токен при обработке контекста в 1 млн токенов, а также в 14,2 раза более быстрый prefill и в 7,6 раза более быстрый decoding на GPU Nvidia H800.

На самом деле важно то, сохраняется ли качество работы модели. По данным MiniMax, их модель на MSA, обученная с нуля, показывает результаты почти идентичные полному вниманию по нескольким метрикам (MMLU 67,2 против 67,0, GSM8K 77,7 против 76,2, HumanEval 64,0 против 61,0) и даже превосходит его в задаче извлечения из длинного контекста (RULER-8K, 84,2 против 79,8). А значит, большую часть выигрыша вы получаете бесплатно! Их базовая модель — это Mixture-of-Experts на 109 млрд параметров, обученная на 3 триллионах токенов, а выпущенное ими ядро (fmha_sm100) поддерживает современные типы точности вроде FP8 и FP4.

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

Тем не менее есть оговорки. На данный момент открытое ядро работает только на новейших GPU Nvidia серии SM100; для других устройств нужны отдельные реализации. По отдельным подзадачам по сравнению с полным вниманием могут оставаться небольшие отличия, а все измерения были получены с помощью внутреннего тестового фреймворка MiniMax без внешней проверки. И всё же внедрение этого механизма в реальный продукт, а не просто в качестве доказательства концепции, показывает, что борьба за эффективную обработку длинного контекста становится всё жарче.