Research map / Architectures and efficiency / Attention and transformers
MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection
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 MatryoshkaKV, which compresses the KV cache along the feature dimension by distilling trainable orthogonal (Cayley-parameterized, PCA-initialized) projections with a Matryoshka rank-sampling objective and a greedy per-head search for adaptive compression rates.
Problem
KV cache size becomes a storage and memory-transfer bottleneck for LLM inference; feature-dimension compression via training-free PCA projections degrades sharply at moderate-to-high compression rates.
Approach
Directly fine-tune orthogonal projection matrices (Cayley parameterization, initialized from PCA) with a KL distillation loss under a Matryoshka strategy that samples truncation ranks per layer/head so that the first r columns form a usable low-rank projection, then greedily allocate per-head key/value ranks under a budget.
Key results
- On LLaMA2-7B-base with CPT (~200M tokens), MatryoshkaKV retains 93.10% of baseline average accuracy at 37.5% cache budget, versus PCA collapsing to 33.14 at that budget.
- On Mistral-v0.3-7B-base, MatryoshkaKV retains 92.63% of baseline average accuracy at 37.5% cache budget.
- In SFT on LLaMA2-7B-base, the method retains 92.47% of baseline average accuracy at a 50% cache budget.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace PCA projections with trained MatryoshkaKV projections | LLaMA2-7B-base · CPT zero-shot, 50.0% cache | average accuracy: 43.07 -> 59.11 | Helped |
| Replace PCA projections with trained MatryoshkaKV projections | Mistral-v0.3-7B-base · CPT zero-shot, 50.0% cache | average accuracy: 34.98 -> 63.75 | Helped |
| Use MatryoshkaKV instead of ASVD | LLaMA2-7B-base · CPT, 90% cache | average accuracy: 60.06 -> 62.03 | Helped |
| Remove greedy search for adaptive compression levels (use uniform compression) (ablation) | MatryoshkaKV · CPT zero-shot, 37.5% cache | average accuracy: 56.94 -> 55.02 | Hurt |
| Remove Matryoshka training and train projections at a fixed 50% cache budget (ablation) | MatryoshkaKV · SFT on GSM8K | accuracy | Hurt |
| Remove the orthogonal constraint during projection training (ablation) | MatryoshkaKV · SFT on GSM8K, full KV cache | accuracy | Hurt |
| Replace PCA initialization with random orthogonal initialization (ablation) | MatryoshkaKV projections · SFT on GSM8K | evaluation loss | Hurt |
| Change the Matryoshka schedule from {i/8 d} to {i/4 d} (ablation) | MatryoshkaKV · CPT zero-shot, LLaMA2-7B-base | average accuracy | Mixed |
| Combine MatryoshkaKV with H₂O compression | MatryoshkaKV · LongBench, 10% overall cache budget | perplexity: +1.02 | Hurt |
| Combine MatryoshkaKV with 2-bit KIVI quantization | MatryoshkaKV · CPT zero-shot, 50.0% cache | average accuracy: 56.98 -> 56.35 | Hurt |
Surprising findings
- PCA projection without training drops to near-random accuracy at 50% budget on several tasks (e.g., HellaSwag 42.11 vs 74.00 baseline), while the distilled method keeps 66.78.
- Values can be compressed more aggressively than keys in the greedy search (key budget 32.28% vs value budget 42.72% at a 37.5% total budget).
Limitations
- Compute budget for training is stated as 30 GPU-hours with limited tokens; results rely on small RedPajama sample.
- Per-task adaptive compression requires a calibration search on task samples (5-10 samples) before inference.
- No system-level optimization; inference speed gains are not realized in the current implementation.
Experimental setup
Datasets: RedPajama-Data-1T-Sample GSM8K PIQA HellaSwag OpenbookQA ARC-C ARC-E WinoGrande
Benchmarks: HellaSwag ARC-C ARC-E PIQA WinoGrande CommonSenseQA GSM8K OpenbookQA LongBench
Metrics: accuracy average accuracy perplexity tokens per second relative accuracy vs baseline
Baselines: PCA projection (Eigen-Attention) ASVD LoRA SFT baseline H2O KIVI
Builds on: Eigen-Attention Matryoshka representation learning LoRA H2O KIVI
Closest papers by meaning
- Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (ICLR 2024 · similarity 0.91)
- EchoKV: Efficient KV Cache Compression via Similarity-Based Reconstruction (NeurIPS 2026 · similarity 0.90)
- SALS: Sparse Attention in Latent Space for KV Cache Compression (NeurIPS 2025 · similarity 0.90)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.90)
- SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget (ICLR 2025 · similarity 0.90)
- MiniCache: KV Cache Compression in Depth Dimension for Large Language Models (NeurIPS 2024 · similarity 0.90)
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.