Упрощенное разреженное внимание через токены сути

Аннотация

Разреженное внимание может снизить стоимость вывода в длинном контексте, но большинство вариантов вводят новые архитектурные компоненты. Мы представляем упрощенное разреженное внимание (SSA) — более простой подход к разреженному вниманию, не требующий архитектурных изменений. В частности, мы сначала проводим продолженное предварительное обучение на последовательностях, чередующихся с токенами сути. Мы оптимизируем стандартную функцию потерь следующего токена как обычно, но токены сути используют маску внимания, чтобы ограничить, к каким частям контекста может обращаться языковая модель; это учит модель упаковывать важную информацию каждого блока в токены сути. Во время вывода SSA оценивает блоки через внимание между текущим запросом и небольшим набором токенов сути, выборочно развертывая top-k блоков путем повторного введения соответствующих исходных токенов. Поскольку запрос оценивается только по отношению к токенам сути, мы избегаем затрат пропускной способности памяти, связанных с наивной оценкой по полному KV-кэшу, без необходимости в дополнительном KV-кэше, используемом методами разреженного внимания. На LongBench SSA стабильно превосходит базовые методы сжатия и разреженного внимания во время вывода при том же коэффициенте сжатия. Еще более примечательно, что в генерации с дополнением на основе поиска SSA может даже превосходить полное внимание после продолженного предварительного обучения более чем на 5,7 пункта. Мы объясняем это способностью выборочного развертывания SSA, которое концентрирует внимание на релевантных запросу блоках и эффективно отфильтровывает шум. SSA также расширяется до иерархического варианта «суть сути» (H-SSA), который достигает логарифмически-линейной сложности декодирования, сохраняя или повышая точность при высоких коэффициентах сжатия до 32x. Код доступен по адресу https://github.com/yuzhenmao/simplified-sparse-attention/.

English

Sparse attention can reduce the cost of long-context inference, but most variants introduce new architectural components. We introduce Simplified Sparse Attention (SSA), a simpler approach to sparse attention that requires no architectural changes. Concretely, we first perform continued pretraining on sequences interleaved with gist tokens. We optimize the standard next-token loss as usual, but the gist tokens use an attention mask to restrict what parts of the context the language model can attend to; this teaches the model to pack each chunk's important information into the gist tokens. At inference time, SSA scores chunks via attention between the current query and the small set of gist tokens, selectively unfolding the top-k chunks by reintroducing their corresponding raw tokens. Since the query is scored only against the gist tokens, we avoid the memory-bandwidth cost associated with naive scoring against the full KV cache, without requiring the auxiliary KV cache approach used by sparse attention methods. On LongBench, SSA consistently outperforms compression and inference-time sparse-attention baselines under the same compression ratio. More strikingly, in retrieval-augmented generation, SSA can even outperform full attention after continued pretraining by over 5.7 points. We attribute this to the ability of SSA's selective unfolding, which concentrates attention on the query-relevant chunks and effectively filters out noise. SSA further extends to a hierarchical gist-of-gist variant (H-SSA) that achieves log-linear decoding complexity while maintaining or improving accuracy at high compression ratios up to 32x. The code is available at https://github.com/yuzhenmao/simplified-sparse-attention/.