Vereinfachte Sparse Attention mittels Gist-Tokens

Zusammenfassung

Sparse Attention kann die Kosten von Inferenzen mit langen Kontexten reduzieren, aber die meisten Varianten führen neue architektonische Komponenten ein. Wir stellen Simplified Sparse Attention (SSA) vor, einen einfacheren Ansatz für sparse Attention, der keine architektonischen Änderungen erfordert. Konkret führen wir zunächst fortgesetztes Vortraining auf Sequenzen durch, die mit Gist-Tokens durchsetzt sind. Wir optimieren wie üblich den standardmäßigen Next-Token-Verlust, aber die Gist-Tokens verwenden eine Aufmerksamkeitsmaske, um einzuschränken, auf welche Teile des Kontexts das Sprachmodell achten kann; dies lehrt das Modell, die wichtigen Informationen jedes Chunks in die Gist-Tokens zu packen. Zum Inferenzzeitpunkt bewertet SSA Chunks über die Aufmerksamkeit zwischen der aktuellen Abfrage und der kleinen Menge von Gist-Tokens, wobei es die Top-k-Chunks selektiv entfaltet, indem es ihre entsprechenden Roh-Tokens wieder einführt. Da die Abfrage nur gegen die Gist-Tokens bewertet wird, vermeiden wir die mit naiver Bewertung gegen den gesamten KV-Cache verbundenen Speicherbandbreitenkosten, ohne den von sparse-Attention-Methoden verwendeten Ansatz des Hilfs-KV-Caches zu benötigen. Auf LongBench übertrifft SSA konsequent Kompressions- und Inferenzzeit-Sparse-Attention-Baselines unter demselben Kompressionsverhältnis. Noch auffälliger ist, dass SSA in der retrieval-gestützten Generierung nach fortgesetztem Vortraining die volle Aufmerksamkeit sogar um über 5,7 Punkte übertreffen kann. Dies führen wir auf die Fähigkeit der selektiven Entfaltung von SSA zurück, die die Aufmerksamkeit auf die abfragerelevanten Chunks konzentriert und Rauschen effektiv herausfiltert. SSA erweitert sich weiter zu einer hierarchischen Gist-of-Gist-Variante (H-SSA), die eine log-lineare Dekodierungskomplexität erreicht, während sie die Genauigkeit bei hohen Kompressionsverhältnissen von bis zu 32x beibehält oder verbessert. Der Code ist verfügbar unter 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/.