PaperlayerView exact revision on arXiv

Source-linked AI summary

PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling

Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Yucheng Li, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Junjie Hu, Wen Xiao

arXiv:2406.02069v4cs.CLcs.AI

TL;DR

Long-context LLMs face attention-computation and KV-cache memory constraints, while the suitability of uniform retention across layers remains unresolved. The paper identifies Pyramidal Information Funneling and introduces PyramidKV, which varies cache allocation by layer; it preserves benchmark performance with 12% cache and reaches 100.0 Acc. with 128 entries on Needle In A Haystack.

  • Problem

    Long-context processing creates quadratic attention costs and memory constraints, while whether uniform KV-cache strategies suit all layers remains unresolved.

  • Method

    PyramidKV analyzes layerwise attention aggregation and dynamically allocates more KV cache to lower layers and less to higher layers.

  • Results

    12% KV cache preserves LongBench performance, and 128 entries enable LLaMa-3-70B-Instruct to achieve 100.0 Acc. matching full-cache performance on Needle In A Haystack.

  • Takeaways & Limitations

    PyramidKV reduces KV-cache memory while preserving long-context understanding, with advantages especially pronounced under memory constraints.

  • Takeaways & Limitations

    Attention-recall improvements diminish as the KV-cache budget increases, including at a 2k cache size.

Abstract

from arXiv · show

In this study, we investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusing on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these insights, we developed PyramidKV, a novel and effective KV cache compression method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reducing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques, achieving up to a 20.5 absolute accuracy improvement on TREC dataset. In the Needle-in-a-Haystack experiment, PyramidKV outperforms competing methods in maintaining long-context comprehension in LLMs; notably, retaining just 128 KV cache entries enables the LLAMA-3-70B model to achieve 100.0 Acc. performance.

1 Introduction

Long-context inference creates quadratic attention delays and KV-cache memory constraints, while prior compression methods generally use uniform retention across layers. PyramidKV instead follows layerwise attention aggregation, preserving performance with substantially less cache.

  • Long-context attention incurs quadratic computation costs, motivating KV-state caching but creating substantial memory constraints.
  • Prior KV-cache strategies leave open whether uniform cache sizes are appropriate across all LLM layers.
  • PyramidKV allocates more cache to lower layers and less to higher layers, matching dispersed-to-concentrated information flow.
  • 12.0% KV cache preserves LongBench performance, while extreme compression to 0.7% significantly outperforms other methods.
  • 128 KV cache entries let LLaMa-3-70B-Instruct achieve 100.0 Acc. in Needle In A Haystack, matching full-cache performance.

2 Related Work

Related work reduces KV-cache memory through adaptive retention, attention-based selection, clustering, or eviction policies.

  • FastGen adapts KV-cache retention tactics to the specific nature of attention heads.
  • SnapKV selects or clusters significant KV positions using their attention scores.
  • H2O dynamically evicts cache entries while balancing recent and historically significant information.

3 Pyramidal Information Funneling

A layerwise study of multi-document question answering identifies a pyramidal attention pattern: broad global aggregation in lower layers becomes localized and then concentrates on key tokens in upper layers.

  • The analysis examines how LLMs aggregate dispersed information from multiple interrelated documents during long-context question answering.
  • Attention behavior is measured by averaging attention across heads within each layer and visualizing selected layers.
  • Lower layers distribute attention approximately uniformly across available content, indicating broad-spectrum global aggregation.
  • Middle layers transition toward localized attention within individual documents as information encoding progresses.
  • Upper layers concentrate attention overwhelmingly on a few key tokens, producing the observed massive-attention phenomenon.

4 PyramidKV

PyramidKV compresses KV caches by allocating layer-specific budgets that follow pyramidal information flow, then retaining the most important KV states within each layer. It preserves recent instruction tokens across layers and selects additional tokens using attention from those tokens.

  • Motivation: KV cache compression stores key and value matrices to reduce redundant autoregressive computations, but long contexts make the cache memory-intensive.The method seeks smaller sub-matrices while preserving evaluation performance relative to the full cache.
  • PyramidKV: PyramidKV uses two steps: dynamically allocating different KV cache budgets across layers and selecting important KV vectors within each attention head.The allocation follows aggregated information flow across layers.
  • Motivation: Fixed cache sizes can be suboptimal because lower layers have denser attention while higher layers have sparser attention.Uniform allocation may retain unimportant higher-layer tokens while overlooking crucial lower-layer tokens.
  • KV Cache Size/Budget Allocation: PyramidKV retains the last α input tokens across all layers as instruction tokens because they contain immediate task-related information.These tokens are also called a local window in prior literature.
  • KV Cache Size/Budget Allocation: The remaining budget forms a pyramid: the top and bottom layer sizes are determined first, and intermediate layer sizes follow an arithmetic sequence.The top-layer budget is controlled by β, while the bottom-layer budget is set to satisfy the total cache budget.
  • KV Cache Selection: Within each layer and attention head, PyramidKV scores non-instruction tokens by attention from instruction tokens and retains the top k_l tokens.A pooling layer is used on attention scores to reduce the risk of massive activations misleading selection.

5 Experiment

Experiments evaluate PyramidKV across LongBench tasks, models, cache sizes, and memory settings, with comparisons against fixed-size baselines and FullKV. PyramidKV generally preserves or improves performance, especially under severe cache constraints, while maintaining long-context retrieval.

  • Experimental Setup: PyramidKV uses varying cache sizes across layers, whereas the baselines keep a fixed size; experiments match average cache size and total memory for fairness.All methods use the same prompts, and baseline comparisons include StreamingLLM, H2O, SnapKV, and FullKV.
  • LongBench Results: 12.0% of the KV cache preserves LongBench performance, while PyramidKV consistently surpasses competing methods across cache sizes and backbone models.Its advantage becomes particularly pronounced in memory-constrained settings retaining about 0.8% of the prompt KV cache.
  • LongBench Results: At KV cache size 2048, PyramidKV improves over baseline methods and even outperforms FullKV in the performance-preserving scenario.The evaluation reports results for KV cache sizes 64 and 2048, representing memory-efficient and performance-preserving scenarios.
  • LongBench Results: PyramidKV’s gains are largest on tasks with room for improvement, including TREC, while its small deficits occur mainly on saturated tasks and remain marginal.The reported pattern yields higher overall average performance and especially strong gains on in-context learning tasks.
  • Needle In A Haystack: In Needle In A Haystack, PyramidKV outperforms competing compression methods; 128 KV cache entries let LLaMa-3-70B-Instruct reach 100.0 Acc., matching FullKV.The test uses context lengths up to 8k and compares PyramidKV, SnapKV, and H2O at cache size 128 and full cache.

6 Conclusion

PyramidKV identifies Pyramidal Information Funneling in LLM attention and uses it to compress KV caches with layer-specific budgets. The method preserves long-context understanding while reducing memory usage, but its evaluation remains limited in model and language coverage.

  • Conclusion: PyramidKV uses Pyramidal Information Funneling to compress KV caches according to layer-specific attention patterns.The approach allocates cache based on how information is distributed across layers.
  • Conclusion: PyramidKV preserves long-context understanding while significantly reducing memory usage in memory-constrained settings.
  • Limitations: The experiments cover only three base models and English-language inputs, limiting the demonstrated scope of the findings.The authors identify broader model-family and multilingual evaluation as future work.
  • Future Work: PyramidKV may support more few-shot examples under constrained memory, although this application is presented as potential future use.

D Pyramidal Information Funneling

Across transformer layers, attention progressively narrows from broad contextual coverage toward localized and highly concentrated token selection. PyramidKV mirrors this pattern by allocating larger lower-layer caches and smaller higher-layer caches while selecting important vectors per head.

  • Pyramidal Information Funneling: Pyramidal Information Funneling emerges across all transformer layers as attention progressively narrows its focus.The analysis examines layers 0 through 30 rather than only isolated lower and upper layers.
  • Attention Patterns: Lower layers attend broadly across global contexts, whereas higher layers concentrate on local tokens and a small set of critical tokens.
  • Massive Attention Mechanism: The critical upper-layer tokens can occur at regular sequence intervals, not only at the leading positions.This differs from the initial-token-only massive attention reported in the cited prior setting.
  • Token Selection: These observations motivate selecting upper-layer tokens by their highest attention scores rather than relying only on earlier positions.
  • PyramidKV: PyramidKV dynamically assigns different cache budgets across layers and selects important KV vectors within each attention head.
  • Allocation Strategy: The method adopts an arithmetic allocation sequence because lower-layer attention is dispersed while higher-layer attention concentrates on fewer tokens.

F Details of Evaluation

The evaluation uses LongBench and ablations to study PyramidKV’s allocation and inference behavior, alongside comparisons with MInference and PyramidInfer. Results support linear allocation, limited overhead, and compatibility with complementary acceleration methods.

  • Evaluation Setup: LongBench evaluates PyramidKV on long-context tasks using a multi-task benchmark designed for extended documents and complex information sequences.Dataset statistics and metric definitions are provided separately in Table 3.
  • Allocation Ablation: A relatively stable linear arithmetic decrease aligns with the observed attention pattern and outperforms alternative pyramidal allocation strategies.The alternatives include geometric and exponential decay strategies.
  • Hyperparameter Ablation: Small β values yield better outcomes, while PyramidKV remains generally robust to β selection.
  • Comparison with MInference: PyramidKV and MInference address different inference stages and can be integrated into a hybrid approach.PyramidKV targets decoding, whereas MInference accelerates KV-cache generation during pre-filling.
  • Comparison with PyramidInfer: PyramidKV consistently outperforms PyramidInfer, whose geometric decay differs from PyramidKV’s arithmetic decay.
  • Inference Efficiency: PyramidKV adds minimal inference overhead, with budget allocation requiring only a one-time computation and comparable speed to baseline methods.

N PyramidKV Excels in all KV Cache Size Limitation

Across LongBench configurations and backbone models, PyramidKV generally performs best at smaller KV-cache sizes, where its allocation strategy is most advantageous. Its attention-recall gains are also largest under the smallest budgets.

  • LongBench Results: PyramidKV consistently surpasses other methods across KV-cache sizes and backbone models, especially in memory-constrained settings.The comparisons include LLaMA-3-8B, Mistral-7B, and LLaMA-3-70B.
  • LongBench Results: PyramidKV’s performance strengths become more evident at small KV-cache sizes across the reported LongBench comparisons.
  • Attention Recall Rate: 1.87% is PyramidKV’s average attention-recall improvement at KV Cache size 64, decreasing to 0.36% at size 2048.
  • Attention Recall Rate: With a small budget, PyramidKV more effectively preserves useful attention information through improved cache allocation.

O LongBench results for 128 context length

PyramidKV was additionally evaluated at 128k context length, with results summarized as a comparison against baselines.

  • 128k-context experiments used Llama-3-8B-Instruct-Gradient-1048k with a sequence length of 128k.
  • The passage reports that these experiments provide further validation of PyramidKV’s scalability and robustness.
  • Table 14 compares PyramidKV with baseline methods at 128k context length.

P PyramidKV Preserves the Long-Context Understanding Ability

Needle-in-a-Haystack evaluations across models, context lengths, cache sizes, and methods show that PyramidKV preserves long-context retrieval more closely to FullKV than competing compression methods.

  • PyramidKV achieves Acc. performance closest to FullKV across the evaluated models and KV cache sizes.The comparisons include Mistral-7B-Instruct-32k, LLaMA-3-8B-Instruct-8k, and LLaMA-3-70B-Instruct-8k with cache sizes 64, 96, and 128.
  • 100.0 Acc. performance is achieved by PyramidKV with 128 KV cache size for LLaMA-3-70B-Instruct at 8k context, matching FullKV.
  • Retrieval heads are predominantly located in higher layers, while no retrieval heads are observed in bottom layers.Additional bottom-layer analysis found no massive attention in any individual head.

R PyramidKV Implementation at vLLM

The vLLM implementation compares compressed and vanilla-attention throughput, but uneven per-layer cache budgets create memory-management inefficiencies in naive implementations.

  • Relative throughput under compression decreases as the new input context length approaches the limit.The resulting wait causes new sequences to be added to the decoding batch later.
  • Naive PyramidKV allocation, release, movement, and access of very small memory chunks may cause inefficiency and fragmentation in vLLM.
  • Per-layer paged attention can address the implementation issue by retrieving each layer’s KV cache through layer-specific block tables.This avoids relying on fixed memory offsets.
Loading 2406.02069v4…