Research map / Architectures and efficiency / Attention and transformers
ThinK: Thinner Key Cache by Query-Driven Pruning
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 ThinK, a query-dependent channel pruning method that removes the least important key-cache channels per head (greedy top-T by query-key interaction norm), complementary to KV eviction and KIVI quantization, cutting key cache memory by 20-40% with near-lossless LongBench accuracy.
Problem
KV cache memory grows linearly with batch size and sequence length in long-context LLM inference, and prior work reduces it along sequence length, layers, or precision but overlooks the channel (head dimension) axis.
Approach
Observe that key cache channel magnitudes are highly imbalanced and attention is low-rank, then prune key channels per head using a query-driven score ||Q[:,j] K[:,j]^T||_F (computed over an observation window), instead of magnitude-only norms, so that pruning minimizes attention-weight loss.
Key results
- On LongBench with LLaMA-3-8B-Instruct and KV-size 2048, SnapKV+ThinK(0.4) reaches avg 41.91 vs 41.58 for SnapKV alone, surpassing the full-KV baseline (41.86 avg on the 8B model's full cache at the table's reported numbers is 41.86).
- ThinK integrated with KIVI achieves 2.8x peak memory reduction on LLaMA-2-7B-chat per the abstract, enabling batch size increase from 4x to 5x on a single GPU.
- Key cache memory on LLaMA-2-7B: KIVI(4/4) 61.7 GB vs KIVI(4/4)+ThinK(0.4) 53.3 GB and ThinK(0.5) 51.2 GB.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add query-driven pruning of 40% of key-cache channels | SnapKV · LongBench, LLaMA-3-8B-Instruct, KV-size 512 | average score: 40.10 -> 40.55 | Helped |
| Add query-driven pruning of 40% of key-cache channels | SnapKV · LongBench, Mistral-7B-Instruct-v0.2, KV-size 512 | average score: 40.46 -> 40.40 | Hurt |
| Add query-driven pruning of 40% of key-cache channels | SnapKV · Needle-in-a-Haystack, LLaMA3-8B-Instruct, KV-size 128 | accuracy: 79.6 -> 79.6 | No clear effect |
| Increase key-channel pruning from 40% to 50% (ablation) | SnapKV · Needle-in-a-Haystack, LLaMA3-8B-Instruct, KV-size 128 | accuracy: 79.6 -> 77.4 | Hurt |
| Add key-cache channel pruning at 40% | KIVI · LongBench | average score: 31.92 -> 31.77 | Hurt |
| Prune 40% of key-cache channels on the vanilla model | LLaMA-3-8B-Instruct · LongBench, full KV cache | average score: 41.86 -> 42.27 | Helped |
| Add key-cache channel pruning at 40% | KIVI (4/4) · LLaMA-2-7B, batch size 300 | TPOT (ms/token): 0.27 -> 0.25 | Helped |
| Add key-cache channel pruning at 40% | KIVI (4/4) · LLaMA-2-7B, batch size 300 | TTFT (ms): 7.0 -> 10.4 | Hurt |
| Add key-cache channel pruning at 40% | KIVI (4/4) · LLaMA-2-7B | throughput (tokens/s): 5168 -> 5518 | Helped |
| Increase recent-size from 0 to 32 (ablation) | H2O + ThinK (0.4) · LongBench, Mistral-7B-Instruct-v0.2 | average score: 38.91 -> 40.47 | Helped |
| Increase recent-size from 32 to 128 (ablation) | H2O + ThinK (0.4) · LongBench, Mistral-7B-Instruct-v0.2 | average score: 40.47 -> 40.51 | No clear effect |
| Also prune 30% of value-cache channels (ablation) | H2O + ThinK · LongBench, LLaMA-3-8B-Instruct, KV-size 512 | average score: 37.23 -> 37.23 | No clear effect |
Surprising findings
- Pruning 40% of key channels often matches or slightly improves LongBench accuracy over unpruned SnapKV/H2O.
- Magnitude-only l1/l2 channel pruning degrades sharply at 40% pruning, while the query-driven criterion does not.
Limitations
- Slightly increases time-to-first-token due to channel importance computation.
- Value cache pruning is less effective; value channels have more uniform magnitudes.
- Performance degrades at high pruning ratios (>=0.5) with small KV sizes, e.g. on Needle-in-a-Haystack.
Experimental setup
Datasets: LongBench Needle-in-a-Haystack
Benchmarks: LongBench Needle-in-a-Haystack
Metrics: accuracy (LongBench score) needle retrieval accuracy peak memory throughput TPOT TTFT
Baselines: H2O SnapKV KIVI l1-norm channel pruning l2-norm channel pruning Palu ASVD
Builds on: SnapKV H2O KIVI
Closest papers by meaning
- Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference (NeurIPS 2025 · similarity 0.92)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.92)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 2024 · similarity 0.92)
- MiniCache: KV Cache Compression in Depth Dimension for Large Language Models (NeurIPS 2024 · similarity 0.92)
- SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget (ICLR 2025 · similarity 0.91)
- SCBench: A KV Cache-Centric Analysis of Long-Context Methods (ICLR 2025 · similarity 0.91)
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.