Разреженное внимание на уровне токенов: Эффективный вывод в длинных контекстах с чередующимся выбором токенов
Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection
February 3, 2026
Авторы: Dongwon Jo, Beomseok Kang, Jiwon Song, Jae-Joon Kim
cs.AI
Аннотация
Квадратичная сложность механизма внимания остается ключевым ограничением при выводе моделей в контекстах большой длины. Существующие методы ускорения либо разрежают матрицу внимания с помощью структурированных паттернов, либо безвозвратно исключают токены на определенных слоях, что может сохранять нерелевантные токены или опираться на необратимые ранние решения, игнорируя динамику важности токенов между слоями и головами внимания. В данной работе мы предлагаем Token Sparse Attention — легковесный и динамичный механизм разрежения на уровне токенов, который сжимает Q, K, V для каждого заголовка в уменьшенное множество токенов при вычислении внимания, а затем декомпрессирует выход обратно в исходную последовательность, позволяя переоценивать информацию о токенах на последующих слоях. Кроме того, наш метод открывает новые возможности проектирования на стыке селекции токенов и разреженного внимания. Подход полностью совместим с реализациями плотного внимания, включая Flash Attention, и может быть легко интегрирован с существующими ядрами разреженного внимания. Эксперименты показывают, что Token Sparse Attention стабильно улучшает компромисс между точностью и задержкой, обеспечивая до 3.23-кратного ускорения внимания для контекста длиной 128K токенов при деградации точности менее 1%. Эти результаты демонстрируют, что динамическое и чередующееся разрежение на уровне токенов является дополнительной и эффективной стратегией для масштабируемого вывода в длинных контекстах.
English
The quadratic complexity of attention remains the central bottleneck in long-context inference for large language models. Prior acceleration methods either sparsify the attention map with structured patterns or permanently evict tokens at specific layers, which can retain irrelevant tokens or rely on irreversible early decisions despite the layer-/head-wise dynamics of token importance. In this paper, we propose Token Sparse Attention, a lightweight and dynamic token-level sparsification mechanism that compresses per-head Q, K, V to a reduced token set during attention and then decompresses the output back to the original sequence, enabling token information to be reconsidered in subsequent layers. Furthermore, Token Sparse Attention exposes a new design point at the intersection of token selection and sparse attention. Our approach is fully compatible with dense attention implementations, including Flash Attention, and can be seamlessly composed with existing sparse attention kernels. Experimental results show that Token Sparse Attention consistently improves accuracy-latency trade-off, achieving up to times3.23 attention speedup at 128K context with less than 1% accuracy degradation. These results demonstrate that dynamic and interleaved token-level sparsification is a complementary and effective strategy for scalable long-context inference.