Atención Dispersa Simplificada mediante Tokens de Esencia

Resumen

La atención dispersa puede reducir el costo de la inferencia en contextos largos, pero la mayoría de las variantes introducen nuevos componentes arquitectónicos. Presentamos Atención Dispersa Simplificada (SSA), un enfoque más sencillo para la atención dispersa que no requiere cambios arquitectónicos. Concretamente, primero realizamos un preentrenamiento continuo en secuencias intercaladas con tokens de resumen. Optimizamos la pérdida estándar de siguiente token como es habitual, pero los tokens de resumen utilizan una máscara de atención para restringir qué partes del contexto puede atender el modelo de lenguaje; esto enseña al modelo a compactar la información importante de cada fragmento en los tokens de resumen. En el momento de la inferencia, SSA puntúa los fragmentos mediante la atención entre la consulta actual y el pequeño conjunto de tokens de resumen, desplegando de forma selectiva los mejores k fragmentos al reintroducir sus correspondientes tokens originales. Dado que la consulta se puntúa solo frente a los tokens de resumen, evitamos el coste de ancho de banda de memoria asociado con la puntuación ingenua frente a toda la caché KV, sin necesidad del enfoque auxiliar de caché KV utilizado por los métodos de atención dispersa. En LongBench, SSA supera consistentemente a los métodos de compresión y de atención dispersa en tiempo de inferencia bajo la misma tasa de compresión. De manera más llamativa, en la generación aumentada por recuperación, SSA puede incluso superar a la atención completa después del preentrenamiento continuo en más de 5.7 puntos. Atribuimos esto a la capacidad de despliegue selectivo de SSA, que concentra la atención en los fragmentos relevantes para la consulta y filtra eficazmente el ruido. SSA se extiende además a una variante jerárquica de resumen de resúmenes (H-SSA) que logra una complejidad de decodificación log-lineal mientras mantiene o mejora la precisión en altas tasas de compresión de hasta 32x. El código está disponible en 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/.