Skip to content

PrefixKV: Adaptive Prefix KV Cache is What Vision Instruction-Following Models Need for Efficient Generation

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

PrefixKV is a KV cache compression method for large vision-language models (LVLMs) that addresses the overlooked heterogeneity of importance distributions across transformer layers. Instead of applying a uniform cache size per layer, it uses binary search over cumulative priority sequences to find an optimal global prefix configuration—layer-wise retention ratios that collectively preserve maximal contextual information under a fixed compression budget. At 20% budget, it achieves 1.8× throughput speedup on LLaVA-1.5-7B with minimal quality loss.

Problem

Existing KV cache compression methods (H2O, StreamingLLM, Elastic Cache, PyramidKV) apply a uniform per-layer cache size, ignoring that different transformer layers have markedly different importance distributions (characterized via Lorenz curves and Gini coefficients). Layers with dispersed importance distributions lose disproportionate contextual information under uniform allocation, degrading generation quality. Even manually designed non-uniform schemes (PyramidKV) use fixed hand-crafted patterns rather than data-driven optimal allocation.

Method

Importance metric: Attention scores are averaged across heads per token per layer to yield per-KV-vector importance \(I_l^n\). Within each layer, KV vectors are sorted by descending importance to form a priority sequence; the prefix of this sequence (top-\(R_l\) fraction) constitutes the retained cache.

Global prefix configuration: The set \(\{R_1, \ldots, R_L\}\) of per-layer retention ratios, subject to \(\sum_l R_l = rL\) (compression budget \(r\)), defines the search space. PrefixKV finds a scalar information-retention threshold \(p\) such that for each layer \(l\), \(R_l = \min(\{o \mid P_l^o \geq p\})\), where \(P_l^o\) is the cumulative priority (fraction of total importance) when retaining the top-\(o\) fraction.

Binary search: \(p \in [0,1]\) is found via binary search (converges in 7–12 steps) until \(\sum_l R_l \approx rL\).

Offline estimation: Because per-layer importance distributions are empirically stable across samples (low std ≈0.01–0.024), the configuration \(\{R_l\}\) is estimated offline from as few as 1–10 random training samples and applied at inference without per-sample overhead. During decoding, newly generated tokens' KV vectors are pruned at a fixed distance from the latest token to maintain the configured ratios (following Elastic Cache).

Key Contributions

  • Empirical demonstration that KV importance distributions differ substantially across layers (Lorenz curve / Gini coefficient analysis), motivating adaptive per-layer budgeting.
  • Global prefix configuration formulation: reduces layer-wise cache sizing to finding a single scalar threshold \(p\) via binary search.
  • Offline calibration procedure requiring only ~10 samples, adding no online inference overhead.
  • State-of-the-art performance on LLaVA-Description and MM-Vet across compression ratios 10%–90%, and generalization to Qwen-VL and LLaMA-3-8B on LongBench/Needle-in-a-Haystack.

Results

  • LLaVA-Description, LLaVA-1.5-7B, 20% budget: PrefixKV PPL 3.69 vs. Pyramid 12.4, H2O 48.3, Elastic 14.0 (uncompressed: 3.20).
  • MM-Vet, LLaVA-1.5-13B, 40% budget: PrefixKV PPL 4.78 vs. Elastic 6.31, H2O 10.6.
  • Inference throughput, LLaVA-1.5-7B, batch=16, 20% budget: 486.7 tok/s vs. 266.6 tok/s full cache (≈1.8× speedup); full cache OOMs at batch=48 while PrefixKV achieves 934.4 tok/s.
  • LLaVA-1.5-13B, 13% baseline PPL gap closed: at ≥30% budget, PrefixKV matches uncompressed PPL (2.73) on LLaVA-Description.
  • LongBench (LLaMA-3-8B, budget=128 tokens): PrefixKV 37.5 vs. PyramidKV 37.1, H2O 35.0, Local 32.0, baseline 41.6.
  • Needle-in-a-Haystack: PrefixKV 97.6 vs. PyramidKV 97.3, H2O 49.3, Local 30.4.
  • Combined with KIVI quantization at 3.125% budget: achieves 40.0 on LongBench.
  • Feature MAE at 50% budget: PrefixKV consistently lower than Elastic, H2O, Local across all layers.

Limitations

  • Evaluation limited to LLaVA-1.5-7B/13B and Qwen-VL; coverage of more diverse LVLM architectures (e.g., those with visual token compression or cross-attention designs) is not demonstrated.
  • Offline configuration assumes importance distribution stability across domains and instruction lengths; while robustness is shown empirically, distribution shift under highly novel or out-of-domain inputs is not fully characterized.
  • The method inherits the general KV eviction framework and still drops tokens; feature-similarity-based merging provides only marginal improvement and positional merging can hurt, leaving residual information loss at very low budgets (10–20%).
  • Benchmarks are perplexity and ROUGE against the model's own uncompressed output rather than human evaluation or task-accuracy metrics (e.g., VQA accuracy, MMMU), which may not capture all quality dimensions.

Relevance to Vision-Language Models

KV cache overhead is a primary bottleneck for deploying LVLMs in production, where image tokens substantially lengthen sequences relative to text-only LLMs. PrefixKV directly targets this bottleneck by revealing and exploiting a previously underexplored layer-wise heterogeneity property specific to LVLM inference. The offline calibration approach is particularly practical: it integrates with existing LVLM checkpoints (LLaVA, Qwen-VL) without retraining and is complementary to other compression axes (quantization, prefill acceleration via FastV). For researchers tracking VLMs, this work establishes layer-adaptive KV budgeting as a key design axis alongside uniform-budget eviction and merging methods.