Research map / Architectures and efficiency / Attention and transformers
UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes UniGist, a sequence-level long-context compression method that replaces raw tokens with interleaved gist tokens under a unified chunk-free sparse attention layout, with a hardware-aligned Triton kernel (gist shift trick) for faster training and real KV removal at inference.
Problem
Gist-token-based sequence-level KV compression loses detail and cross-chunk information because chunk-wise training lets tokens rely on nearby uncompressed tokens, and its fragmented sparse attention is inefficient on GPUs.
Approach
Remove chunk-wise training in favor of one-pass training with a unified sparse gist attention pattern (sink tokens, gist tokens, local window), align position ids so gists share the next raw token's position, and apply a gist shift trick that moves gist tokens to the right so GPU kernels can skip non-visible blocks.
Key results
- On HELMET with Llama-3.1-8B-Instruct, UniGist averages 60.4 versus 63.9 for full attention and 51.8 for Activation Beacon, the best compression baseline.
- On Llama-3.2-3B-Instruct, UniGist averages 47.8 versus 37.9 for Activation Beacon and 51.9 for full attention.
- Attention speed-up approaches the compression ratio (about 4x at ratio 4) as context length grows, per the per-layer latency measurements.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use UniGist instead of Activation Beacon | Llama-3.1-8B-Instruct · HELMET Synthetic | score: 91.3 vs 61.8 | Helped |
| Use UniGist instead of Activation Beacon | Llama-3.2-3B-Instruct · HELMET ICL | score: 73.8 vs 57.4 | Helped |
| Use UniGist instead of Activation Beacon | Llama-3.1-8B-Instruct · HELMET Summ. | score: 22.9 vs 21.0 | Helped |
| Use UniGist instead of Activation Beacon | Llama-3.2-3B-Instruct · HELMET Average | score: 47.8 vs 37.9 | Helped |
| Train UniGist with 128 sink tokens rather than without sink tokens (ablation) | UniGist · Training steps 50–100 | loss | Helped |
| Train UniGist with 128 sink tokens rather than without sink tokens (ablation) | UniGist · Training steps 50–100 | gradient norm stability | Helped |
| Vary the local attention window size (ablation) | UniGist · RAG and synthetic recall tasks | score | Mixed |
| Set the local attention window size to 128 (ablation) | UniGist · RAG and synthetic recall tasks | score | Helped |
| Apply compression-specific training with UniGist rather than full attention | UniGist · MMLU-Pro, GSM8K, and HellaSwag | score: 47.7 vs 47.6; 83.6 vs 83.9; 60.2 vs 60.1 | No clear effect |
Surprising findings
- Full-attention fine-tuning on the same data slightly lowers HELMET average for Llama-3.1-8B (61.4 vs 63.9), suggesting the SFT data itself hurts long-context performance.
Limitations
- Experiments limited to models up to 8B parameters and 32K context length.
- Adapts existing models via continued pretraining rather than training from scratch, so the architecture's inductive biases are not tested from scratch.
Experimental setup
Datasets: ProLong 64K dataset (restructured to 32K) Magpie-Llama-3.1-Pro-MT-300K-Filtered LongAlpaca BookSum
Benchmarks: HELMET MMLU-Pro GSM8K HellaSwag
Metrics: SubEM SumEM NDCG@5 ROUGE Accuracy peak memory forward/backward latency
Baselines: StreamingLLM SnapKV PyramidKV AdaKV AutoCompressors Activation Beacon Full Attention Full Attention (FT)
Builds on: gist tokens Activation Beacon AutoCompressors FlashAttention StreamingLLM
Closest papers by meaning
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.91)
- ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference (NeurIPS 2025 · similarity 0.90)
- ParallelComp: Parallel Long-Context Compressor for Length Extrapolation (ICML 2025 · similarity 0.89)
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.89)
- Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers (NeurIPS 2025 · similarity 0.89)
- One Token per Highly Selective Frame: Towards Extreme Compression for Long Video Understanding (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.