Joint Subspace Sparse Attention: новый метод для эффективного декодирования длинных контекстов в LLM
В сфере больших языковых моделей (LLM) проблема обработки контекста длиной в сотни тысяч токен сталкивается с узким местом пропускной способности памяти. Новая методология, опубликованная на arXiv под названием Joint Subspace Sparse Attention (JSSA), предлагает радикально увеличить скорость генерации текста, сохранив при этом точность, сравнимую с полным вниманием к ключам и значениям (full-KV). Алгоритм решает две критические задачи: повышает качество выбора токенов и оптимизирует их сборку в памяти, демонстрируя ускорение до 6,43 раза на GPU NVIDIA B200 при контексте 128K.
# Joint Subspace Sparse Attention: новый метод для эффективного декодирования длинных контекстов в LLM
Проблема пропускной способности памяти в эпоху длинного контекста
Современные большие языковые модели сталкиваются с растущим давлением на вычислительные ресурсы по мере увеличения длины контекста. Основной потребитель памяти и времени обработки при генерации текста — это механизм внимания, работающий с кэшем ключей и значений (KV-cache). В традиционной архитектуре, известной как "full-KV attention", модель должна обрабатывать каждое новое сгенерированное слово относительно всех предыдущих токенов в истории, что приводит к линейному росту стоимости вычислений.
Команды исследователей из arXiv:2610.19026 указывают в своем абстрактном изложении, что существующие методы разреженного внимания (sparse attention) часто работают на грубой зернистости, что ограничивает их эффективность. Попытки реализовать выбор на уровне отдельных токенов (token-level selection) обещают более тонкий контроль, но вводят новые сложности. Во-первых, использование низкоразмерного скоринга может привести к потере структуры, необходимой для точного ранжирования запросов (query) и ключей (key). Во-вторых, сбор разрозненно выбранных токенов может нивелировать выигранное время, если не оптимизировать сам процесс их агрегации.
Техническое пояснение: В простейших терминах, "кэш ключей и значений" (KV-cache) — это аналог заметок на полях страницы книги, которые модель читает при каждой попытке написать новую фразу. Если заметки занимают слишком много места и времени на поиск, процесс чтения (генерации) замедляется. Метод JSSA пытается прочитать только самые важные заметки, не теряя смысл всего абзаца.
Joint Subspace Sparse Attention: теория и алгоритм
Предлагаемое исследование представляет Joint Subspace Sparse Attention (JSSA) как решение, которое одновременно атакует проблемы точности и исполнения. Авторы разработали подход, при котором модель учит разделяемое низкоразмерное подпространство запросов и ключей. Это делается путем минимизации ошибки сохранения токен-скоринга, что позволяет уловить необходимую структуру данных, не перегружая систему.
С технической точки зрения, авторы аппроксимируют полученную негладкую математическую задачу (non-convex objective) более простыми маргинальными целями (tractable marginal objective), которые имеют замкнутое спектральное решение. Это означает, что вычисления могут выполняться эффективно, используя методы спектральной теории, без необходимости сложных итеративных уточнений, которые не дали измеримого улучшения в tested settings (тестовых настройках).
Для обеспечения реальной эффективности на "железе" была создана специальная версия ядра внимания (indexed-attention kernel), свободная от операций сбора данных (gather-free). Этот механизм поддерживает различные выборы на уровне голов внимания (per-KV-head selections) в рамках стандартной компоновки кэша группового запроса (Grouped-Query Attention, GQA). Это позволяет избежать архитектурных зависимостей, которые часто приводят к сбоям или неравномерной производительности в других разреженных методах.
Результаты: от 128K контекста до масштабирования моделей
Экспериментальная часть исследования была проведена на различных типах моделей архитектуры внимания, включая стандартные GQA, гибридные скользящие окна с глобальным вниманием и модели с множественными головными латентными вниманием (Multi-head Latent Attention, MLA). Результаты показали, что JSSA стабильно сохраняет точность, близкую к полному вниманию, при этом избегая значительного "набухания" длины генерации, характерного для конкурирующих методов.
Ключевой метрикой, демонстрирующей эффективность, стало ускорение декодирования. На GPU NVIDIA B200 метод JSSA достиг ускорения пропускной способности до 6,43 раза по сравнению с полным KV-вниманием при обработке контекста длиной 128 000 токенов.
Исследователи отметили также отсутствие необходимости в дополнительной тренировке модели (model-native DSA) для достижения сопоставимой точности логического рассуждения. Масштабируемость подхода также была проверена на моделях объемом 700 миллиардов параметров. Даже в таких экстремальных условиях JSSA достиг 96–99% своей максимальной пропускной способности, что свидетельствует о robustness (устойчивости) алгоритма при работе с гигантскими языковыми моделями.
Этот результат важен, поскольку ранее попытки оптимизировать внимание на таких масштабах часто приводили к катастрофическому падению качества или полной неработоспособности на специфических архитектурах. JSSA, напротив, предлагает универсальное решение, работающее на стандартных компонентах индустрии.