Правильно реализованное иерархическое разреженное внимание: к моделированию бесконечного контекста

Аннотация

Масштабирование современных больших языковых моделей (LLM) до длинных контекстов ограничено квадратичной вычислительной стоимостью и плохой экстраполяцией длины плотного внимания. Блочно-разреженное внимание предлагает многообещающую альтернативу, но все существующие методы уступают полному вниманию из-за неточного отбора блоков. Мы предлагаем иерархическое разреженное внимание с ориентирами (HiLS-Attention) — механизм блочно-разреженного внимания, который обучает отбор блоков сквозным образом с помощью функции потерь языковой модели (LM). HiLS иерархически факторизует внимание: каждый запрос независимо выполняет внимание с каждым извлечённым блоком для получения специфичной для блока информации, а результирующие выходы объединяются в соответствии с оценками извлечения блоков. Включая оценки извлечения в прямой проход вычислений внимания, HiLS оптимизирует их непосредственно с помощью потерь LM, что позволяет реализовать сквозное обучение извлечению и естественное разреженное обучение. Экспериментальные результаты показывают, что HiLS-Attention достигает производительности, сопоставимой, а в некоторых случаях и превосходящей полное внимание, на длинах контекстов внутри домена. В то же время HiLS-Attention экстраполирует длину контекста более чем в 64 раза относительно длины обучения с 90% точностью извлечения, далеко превосходя полное внимание. Кроме того, существующие модели полного внимания могут быть преобразованы в HiLS-Attention путём лёгкого продолженного предобучения, сохраняя производительность внутри домена и приобретая способность к сверхдлинной контекстной экстраполяции. В сочетании с разреженным доступом к KV и вычислениями HiLS-Attention нарушает обычный компромисс между эффективностью и производительностью, обеспечивая LLM с длинным контекстом, которые одновременно более эффективны и более результативны в общих задачах с длинным контекстом по сравнению с аналогами, использующими полное внимание.

English

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than 64times the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.