Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention

2026-04-22

AI Summary

This work addresses the quadratic computational overhead of full attention mechanisms in large language models when processing long-context inputs. The authors propose Gist-based Sparse Attention (GSA), an end-to-end learnable sparse attention mechanism that first compresses the context into coarse-grained, learnable gist tokens and then selectively expands the most relevant original segments for fine-grained computation, thereby implementing a coarse-to-fine multi-level access strategy. Innovatively integrating KV cache compression with sparse attention, GSA enables recursive construction of multi-resolution “gist-of-gist” structures. Experimental results demonstrate that under compression ratios of 8× to 32×, GSA significantly outperforms existing context compression and inference-time sparse attention methods on both LongBench and RAG benchmarks.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Scaling large language models to long contexts is challenging due to the quadratic computational cost of full attention. Mitigation approaches include KV-cache selection or compression techniques. We instead provide an effective and end-to-end learnable bridge between the two without requiring architecture modification. In particular, our key insight is that interleaved gist compression tokens -- which provide a learnable summary of sets of raw tokens -- can serve as routing signals for sparse attention. Building on this, we introduce selective unfolding via GSA, which first compresses the context into gist tokens, then selects the most relevant gists, and subsequently restores the corresponding raw chunks for detailed attention. This yields a simple coarse-to-fine mechanism that combines compact global representations with targeted access to fine-grained evidence. We further incorporate this process directly into training in an end-to-end fashion, avoiding the need for external retrieval modules. In addition, we extend the framework hierarchically via recursive gist-of-gist construction, enabling multi-resolution context access with logarithmic per-step decoding complexity. Empirical results on LongBench and RAG benchmarks demonstrate that our method consistently outperforms other compression baselines as well as inference-time sparse attention methods across compression ratios from $8\times$ to $32\times$. The code is available at: https://github.com/yuzhenmao/gist-sparse-attention/

Problem

• long-context
• attention efficiency
• KV-cache compression
• sparse attention
• learnable compression

Innovation

• Gist Sparse Attention
• learnable compression
• selective unfolding
• sparse attention
• long-context modeling