Проблема длинного контекста
При работе с контекстом около 1 миллиона токенов плотное внимание (dense attention) становится вычислительно неподъемным как на этапе префилла, так и при декодировании. Существующие методы разреженного внимания (sparse attention) обычно оценивают все предыдущие блоки или токены, что дает сложность выбора $O(L)$ для каждого запроса. Иерархические методы снижают эту сложность до $O(\log L)$, но их обучаемые версии ранее демонстрировали точность только в пределах длины, на которой они были обучены, не позволяя эффективно экстраполировать результаты.
Решение: Recursive Sparse Attention (RecSA)
Авторы (Ruitian Wang, Xinyu Yang, Minshen Zhang, Hao Zhang) предлагают RecSA, который строит дерево над KV-кэшем. Каждый узел дерева суммирует свое поддерево с помощью одного ключа (key) и одного усредненного значения (mean-pooled value). Ключи извлекаются с помощью пуллера, общего для всех уровней и дистиллированного из плотного внимания. Ключевая особенность: каждый уровень выше токенов выбирает не более $k$ кандидатов, что обеспечивает сложность выбора $O(\log L)$ и чтение $O(\log L)$ записей. Нераспределенные кандидаты используют общий softmax с точными токенами, а позиционное кодирование RoPE адаптируется так, чтобы относительные расстояния оставались в диапазоне, видимом во время обучения.
Экспериментальные результаты
Модель была инициализирована на базе Qwen3-4B и дообучена на 0.62 миллиарда токенов при исходной длине обучения 32K. Тестирование проводилось на бенчмарке RULER в диапазоне от 131K до 2M токенов. RecSA превзошел все базовые методы на всех длинах. Особенно показателен результат при 1 миллионе токенов: RecSA набрал 60.5 баллов, тогда как лучший базовый метод (Dense + DCA) показал лишь 40.0.
| Метод | Контекст (токены) | Результат RULER | Ускорение префилла (H200) | Ускорение декодирования (H200) |
|---|---|---|---|---|
| RecSA (предложенный) | 1,000,000 | 60.5 | 3.2x | 3.8x |
| Dense + DCA (лучший baseline) | 1,000,000 | 40.0 | 1.0x (база) | 1.0x (база) |
| RecSA (оптимизированные ядра) | 524,288 | — | 3.2x | 3.8x |
Производительность и практическая значимость
С использованием оптимизированных ядер RecSA работает быстрее плотного внимания в end-to-end префилле и декодировании в диапазоне от 64K до 2M токенов. На одном GPU NVIDIA H200 при контексте 524K токенов метод обеспечивает ускорение префилла в 3.2 раза и декодирования в 3.8 раза. Это открывает путь к эффективной обработке сверхдлинных документов, таких как юридические архивы, полные истории кода или научные монографии, без катастрофического падения скорости инференса.
Источник: arXiv cs.CL ↗