Research map / Architectures and efficiency / Attention and transformers
Loki: Low-rank Keys for Efficient Sparse Attention
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.
Loki approximates attention by computing top-k key selection scores in a low-dimensional PCA-projected key space (calibrated offline), then computing exact attention over the selected keys, with Triton kernels giving up to 45% attention speedup on Llama2-13B.
Problem
Self-attention in LLM inference is costly in memory bandwidth and compute for long sequences because the KV-cache must be read every decoding step; existing sparse methods either permanently evict tokens or use fixed sparsity patterns and degrade quality.
Approach
Keys lie in a low-dimensional subspace consistently across models and datasets; PCA-project queries and keys, rank tokens using only the top-d principal dimensions, then compute exact attention on the top-k tokens without deleting any cache entries.
Key results
- Loki with k_f=0.25, d_f=0.25 keeps Llama2-7B perplexity at 5.2017 vs 5.1101 full attention, versus 5.2810 for H2O.
- Loki achieves nearly 45% attention speedup over vanilla attention for Llama2-13B at prompt length 3072 and generation length 512 (Triton kernels, excluding KV-cache append time).
- Top-k selection agreement (Jaccard similarity) between Loki and exact attention is around 0.9 for the main settings on Llama2-7B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Loki sparse attention (k_f=0.25, d_f=0.25) instead of full attention | Llama2-7B · WikiText-2 test | perplexity: 5.1101 -> 5.2017 | Hurt |
| Use Loki sparse attention (k_f=0.25, d_f=0.25) instead of full attention | Llama2-13B · WikiText-2 test | perplexity: 4.5680 -> 4.6102 | Hurt |
| Use Loki sparse attention (k_f=0.25, d_f=0.25) instead of full attention | Llama3-8B · WikiText-2 test | perplexity: 5.5696 -> 5.6648 | Hurt |
| Use Loki sparse attention (k_f=0.25, d_f=0.25) instead of full attention | Mistral-7B · WikiText-2 test | perplexity: 4.9140 -> 4.9233 | Hurt |
| Use pre-rotary rather than post-rotary PCA transformation (ablation) | Loki · Llama3-8B and Mistral-7B | perplexity: about 5 for the full attention to over 10 | Helped |
| Use k_f=0.25, d_f=0.25 rather than k_f=0.125, d_f=0.5 (ablation) | Loki · Short-context and LongBench tasks | accuracy | Helped |
| Increase key budget and reduced dimension from k_f=0.25, d_f=0.25 to k_f=0.5, d_f=0.5 (ablation) | Loki on Llama2-7B · Short-context tasks | average score: 43.50 -> 47.45 | Helped |
| Use k_f=0.125, d_f=0.5 rather than k_f=0.25, d_f=0.25 (ablation) | Loki on Llama2-7B · Short-context tasks | average score: 43.50 -> 39.76 | Hurt |
| Use PCA transforms calibrated on different datasets (ablation) | Loki · Short-context tasks | accuracy | No clear effect |
| Use optimized Triton kernels rather than the compared SparQ kernels | Loki attention implementation · Certain kernel scenarios | speedup: 2-3× | Helped |
| Use Loki rather than vanilla attention | Llama2-13B-like setup · Prompt length 3072, generation length 512 | attention compute speedup: nearly a 45% speedup | Helped |
| Vary the per-layer reduced-dimension policy (ablation) | Loki · Appendix experiment | model quality | No clear effect |
Surprising findings
- Post-rotary PCA transforms perform significantly worse than pre-rotary for Llama3-8B and Mistral-7B, despite attention being computed from post-rotary keys.
- More than 80% of Llama2-13B inference time in Hugging Face is spent appending KV-cache entries, which limits end-to-end gains.
Limitations
- Does not reduce KV-cache memory footprint.
- Top-k selection in PyTorch is a bottleneck; not yet integrated with FlashAttention.
- Llama3-8B and Mistral-7B degrade with post-rotary PCA, and Llama3-8B shows notably worse short-context performance.
Experimental setup
Datasets: WikiText-2 WikiText-103 C4 BookCorpus LongBench GSM8K MMLU
Benchmarks: WikiText-2 perplexity HellaSwag TruthfulQA Winogrande ARC GSM8K MMLU LongBench
Metrics: perplexity accuracy Jaccard similarity of top-k attention time speedup
Baselines: Full Attention Exact-TopK H2O SparQ Attention InfiniGen
Builds on: SparQ Attention H2O RoPE Triton
Closest papers by meaning
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.91)
- More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding (NeurIPS 2026 · similarity 0.91)
- SALS: Sparse Attention in Latent Space for KV Cache Compression (NeurIPS 2025 · similarity 0.90)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.90)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.90)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 2024 · 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.