透過主旨標記的簡化稀疏注意力

摘要

稀疏注意力可降低長上下文推論的成本，但大多數變體會引入新的架構元件。我們提出簡化稀疏注意力（SSA），這是一種無需架構調整的簡潔稀疏注意力方法。具體而言，我們首先對交錯插入要點詞元的序列進行持續預訓練。如同常規方式最佳化標準的下一詞元損失函數，但這些要點詞元透過注意力遮罩限制語言模型能關注的上下文範圍，從而使模型學會將每個區塊的重要資訊壓縮至要點詞元中。在推論階段，SSA 透過當前查詢與少量要點詞元間的注意力計算對區塊進行評分，並透過重新引入對應的原始詞元來選擇性展開前 k 個區塊。由於查詢僅與要點詞元進行比對，我們避免了直接對完整 KV 快取進行評分所導致的記憶體頻寬消耗，亦無需採用稀疏注意力方法所需的輔助 KV 快取機制。在 LongBench 基準上，SSA 在相同壓縮比率下一致優於壓縮式與推論階段稀疏注意力基準方法。更引人注目的是，在檢索增強生成場景中，SSA 經持續預訓練後甚至能比完整注意力方法提升超過 5.7 個百分點。我們將此歸因於 SSA 的選擇性展開能力，其能將注意力集中於與查詢相關的區塊，並有效濾除雜訊。SSA 更進一步擴展為階層式要點詞元變體（H-SSA），在最高達 32 倍壓縮比率下，實現對數線性解碼複雜度，同時維持或提升準確率。程式碼已公開於 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/.