MiniCache: KV Cache Compression in Depth Dimension for Large Language Models

Akide Liu; Jing Liu; Zizheng Pan; Yefei He; Gholamreza Haffari; Bohan Zhuang

Source-linked AI summary

MiniCache: KV Cache Compression in Depth Dimension for Large Language Models

Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Gholamreza Haffari, Bohan Zhuang

arXiv:2405.14366v2May 23, 2024cs.CLcs.AIcs.LG

TL;DR

Longer contexts make KV caches a major inference-memory bottleneck. MiniCache compresses redundant KV states across adjacent middle-to-deep layers while retaining distinct token pairs, achieving strong compression and throughput with near-lossless performance.

Problem
KV-cache size grows with sequence length, creating substantial memory demands for long-context LLM inference.
Method
MiniCache merges adjacent-layer KV states by interpolating their directions while preserving magnitudes, and retains distinct token pairs separately.
Results
4-bit MiniCache achieves up to 5.02× compression, approximately 5× higher inference throughput, and 41% lower memory than the FP16 full-cache baseline with near-lossless performance.
Takeaways & Limitations
MiniCache provides a training-free depth-wise KV-cache compression approach that complements quantization and sparsity for more efficient LLM inference.
Takeaways & Limitations
Truthfulness, security, adversarial attacks, and data leakage remain unresolved challenges alongside improvements in LLM efficiency.

Abstract

from arXiv · show

A critical approach for efficiently deploying computationally demanding large language models (LLMs) is Key-Value (KV) caching. The KV cache stores key-value states of previously generated tokens, significantly reducing the need for repetitive computations and thereby lowering latency in autoregressive generation. However, the size of the KV cache grows linearly with sequence length, posing challenges for applications requiring long context input and extensive sequence generation. In this paper, we present a simple yet effective approach, called MiniCache, to compress the KV cache across layers from a novel depth perspective, significantly reducing the memory footprint for LLM inference. Our approach is based on the observation that KV cache states exhibit high similarity between the adjacent layers in the middle-to-deep portion of LLMs. To facilitate merging, we propose disentangling the states into the magnitude and direction components, interpolating the directions of the state vectors while preserving their lengths unchanged. Furthermore, we introduce a token retention strategy to keep highly distinct state pairs unmerged, thus preserving the information with minimal additional storage overhead. Our MiniCache is training-free and general, complementing existing KV cache compression strategies, such as quantization and sparsity. We conduct a comprehensive evaluation of MiniCache utilizing various models including LLaMA-2, LLaMA-3, Phi-3, Mistral, and Mixtral across multiple benchmarks, demonstrating its exceptional performance in achieving superior compression ratios and high throughput. On the ShareGPT dataset, LLaMA-2-7B with 4-bit MiniCache achieves a remarkable compression ratio of up to 5.02x, enhances inference throughput by approximately 5x, and reduces the memory footprint by 41% compared to the FP16 full cache baseline, all while maintaining near-lossless performance.

1 Introduction

MiniCache addresses the growing memory burden of KV caches by compressing redundant states across adjacent layers in the middle-to-deep portions of LLMs. It combines cross-layer merging with token retention and achieves strong efficiency with near-lossless performance.

Motivation: KV caches can dominate inference memory for long sequences, making compression important for efficient LLM deployment.A 175B GPT-3 configuration with batch size 64 and sequence length 4,096 requires approximately 1,208GB of GPU memory.
Contribution: MiniCache compresses KV caches across the depth dimension, exploiting high similarity between adjacent middle-to-deep layers.This complements existing quantization and sparsity approaches.
Method: The method reparameterizes cache states into magnitude and direction, interpolates directions, and preserves state norms during merging.A token retention mechanism separately stores low-similarity state pairs that are unsuitable for merging.
Evaluation: MiniCache is training-free and evaluated across diverse LLMs and question-answering, generation, and long-context benchmarks.Experiments include Mixtral-8x7B, Phi-3-Mini, LLaMA-3 8B and 70B, and LongBench.
Results: 4-bit MiniCache reaches up to 5.02× compression, approximately 5× higher inference throughput, and 41% lower memory than the FP16 full-cache baseline with near-lossless performance.These results are reported as comparisons against the FP16 full cache baseline.

2 Related Work

Prior efficient-inference research reduces LLM cost through dynamic computation, quantization, sparsity, and model merging. These approaches operate on model structure, numerical precision, retained elements, or aggregated parameters and activations.

Efficient inference: Dynamic inference methods select computational substructures or skip unimportant layers according to the input.Examples include mixture-of-experts, mixture-of-depths, and layer-skipping methods.
Efficient inference: Quantization converts weights and activations to low-bit formats, reducing memory footprint and computational intensity.Sparsification removes unnecessary elements from model weights and token representations.
Model merging: Model merging aggregates model parameters and activations at different granularities to improve inference efficiency and exploit redundancy.Weight averaging is presented as a common merge-compression technique.

3 Motivation

The motivation is to exploit redundancy in KV caches across adjacent layers rather than only within layers. Experiments show strong middle-to-deep similarity, promising average merging, and a need to retain distinct token pairs.

Explorations and observations: Figure 2 compares average merging, MiniCache, and the full-cache baseline across five datasets.Its panels cover the average-merging baseline, adjacent-layer pairwise similarity, and cross-dataset performance.
Cross-Layer Redundancy: KV-cache similarity is low in shallow layers but high between adjacent middle-to-deep layers.This pattern is measured using angular distance on LLaMA-3-70B across COQA, GSM8K, and TruthfulQA.
Cross-Layer Redundancy: Starting from the middle layer, adjacent-layer KV-cache merging shows favorable performance across four LLMs on GSM8K.The models are LLaMA-2-7B, LLaMA-2-13B, LLaMA-3-8B, and Mixtral-8x7B.
Cross-Layer Redundancy: Most adjacent-layer token pairs are highly similar, but a few semantically distinct outlier pairs are non-mergeable and require retention.Merging these distinct tokens degrades performance, as reflected by the γ = 0 row in Table 2.

4 Method

MiniCache compresses adjacent-layer KV caches from the model midpoint onward by merging shared directional states while preserving magnitudes and retaining difficult tokens. The resulting shared representation supports approximate restoration with lightweight extra storage and reduces cache memory usage.

Cross-Layer Compression: MiniCache begins at S = L/2 and consolidates sufficiently similar KV pairs from adjacent middle-to-deep layers into a shared cache.The shared representation eliminates independent storage and processing of the original key and value states for each paired layer.
Reparameterization-Based Cache Merging: Direct averaging can lose information, so MiniCache reparameterizes each state into magnitude and direction, then interpolates directions while preserving original ℓ2 norms.SLERP is primarily used for directional interpolation; it follows the shortest path on the unit sphere, while Fmax offers lower overhead but lower accuracy.
Unmergeable Token Retention: MiniCache retains highly distinct token pairs instead of merging them, protecting layer-specific information associated with semantically important tokens.The retained index is selected using angular-distance thresholds controlled by γ, and retained tokens are restored at their original positions.
Cache Restoration: The shared cache stores merged states, retained tokens, per-layer magnitudes, and token indices, after which restoration rescales directions and recovers retained tokens.These additional components are described as lightweight relative to full-layer caches.
Efficiency Discussion: 4brh(s + n) becomes 3brh(s + n) when every two middle-to-deep layers share one KV state, reducing FP16 decoding-cache usage.Magnitude vectors add only a single-channel norm vector, while a retention threshold of 0.05 contributes an additional brh(0.05(s + n)) term.

5 Experiments

MiniCache is evaluated across models, benchmarks, compression baselines, and long-context workloads. It preserves performance while reducing memory use and increasing throughput, including a 5.02× compression rate with KIVI on LongBench.

Experimental setup: MiniCache is evaluated on diverse models and tasks against fully cached, quantized, and sparsity-oriented baselines.Experiments include Phi-3-Mini, Mixtral-8x7B, LLaMA-3 8B and 70B, LongBench, and lm-eval-harness tasks.
Main results: MiniCache maintains robust performance while merging middle-to-deep KV-cache layers, outperforming averaging and achieving its strongest compression on larger models.On LLaMA-3-70B, merging 87.5% of layers causes nearly zero performance drop on COQA.
LongBench: 5.02× compression is achieved when MiniCache is combined with 4-bit KIVI on LongBench, with the strongest compression rate and best performance among compared methods.MiniCache is orthogonal to existing quantization and sparsity methods at model and token levels.
Main results: Merging more layers increases memory reduction, while Figure 4 compares MiniCache with averaging and unmerged full-cache baselines across multiple models and datasets.The figure varies the number of merged layers on the x-axis.
Efficiency analysis: 41% memory saving and approximately 5× higher throughput are achieved by 4-bit MiniCache versus the FP16 baseline at batch size 128.MiniCache reduces memory usage by 25GB and reaches 1.29× the throughput of 2-bit KIVI.

6 Ablation Study

The ablations examine the interpolation parameter t and token-retention threshold γ. They identify t = 0.6 and γ = 0.05 as settings that balance merged-cache performance, retained information, and memory demand.

Interpolation parameter t: t = 0.6 produces the most robust merged representation, whereas t = 0.5 behaves like less-effective average merging.The result indicates that more information is derived from the second SLERP term at the optimal setting.
Interpolation parameter t: High relative-magnitude-ratio frequencies cluster around 0.4 and 0.6, correlating with the optimal interpolation parameter and motivating dynamic t selection.Dynamic t could provide layer-wise control of SLERP weights.
Token retention threshold γ: γ = 0.05 provides the best balance between performance and efficiency across three datasets.Higher retention thresholds generally preserve more tokens but increase memory demand.

7 Conclusion and Future Work

MiniCache explores depth-wise KV cache compression and reports substantial memory and throughput gains, while identifying multi-layer merging as future work.

7 Conclusion and Future Work: 41% lower memory and approximately 5× higher throughput than the FP16 baseline are reported for MiniCache.The paper describes this as a state-of-the-art balance between efficiency and performance.
7 Conclusion and Future Work: Future work targets cross-multiple-layer merging, advanced interpolation, and further memory optimization for large-scale deployments.

A Additional Experiment Results

Additional experiments compare MiniCache with token sparsity methods and position it as complementary because the approaches target different redundancy dimensions.

A Additional Experiment Results: MiniCache outperforms H2O in most LongBench tasks.The comparison uses Mistral-7B-Instruct on the LongBench dataset.
A Additional Experiment Results: MiniCache reduces inter-layer redundancy, whereas H2O reduces intra-layer redundancy, making the approaches orthogonal.
A Additional Experiment Results: KV-cache research includes quantization, sparsity, attention optimization, query grouping, token shrinking, and long-context generation strategies.

C Discussions and Limitations

The discussion frames MiniCache as a training-free, depth-wise compression approach with practical efficiency benefits, while noting algorithmic and broader LLM limitations.

C Discussions and Limitations: MiniCache merges similar neighboring-layer KV states and is adaptable to existing intra-layer pruning and quantization technologies.
C Discussions and Limitations: The method is presented as a preliminary exploration of depth-wise KV-cache compression for low-resource deployment and long-context generation.The discussion mentions mobile-device deployment, textbook understanding, batch inference, and long-context generation.
C Discussions and Limitations: Truthfulness, security vulnerabilities, adversarial attacks, and data leakage remain unresolved challenges for LLMs alongside efficiency improvements.
C Discussions and Limitations: SLERP merges only two vectors, restricting simultaneous multi-layer merging and limiting further compression-ratio improvements.The paper identifies more sophisticated merging algorithms as a needed direction.

D Additional Implementation Details

MiniCache begins cross-layer merging at a chosen layer, shares caches across adjacent layers, deletes redundant storage, and restores approximate scales and retained tokens during decoding.

D Additional Implementation Details: From the merging layer S onward, cross-layer merging is applied at odd-numbered layers while previous-layer caches are removed.Magnitudes and retention-sensitive tokens are stored for recovery.
D Additional Implementation Details: During decoding, even layers fetch shared caches and perform scale restoration and retention-token recovery before new KV states are reused.
D Additional Implementation Details: The merging function normalizes adjacent-layer KV pairs, interpolates along the shortest spherical path, and preserves geometric information through restoration.
D Additional Implementation Details: The inference flow fetches, merges, caches, deletes, retrieves, and error-suppresses KV pairs across prefilling and decoding.
D Additional Implementation Details: Retention indices preserve unmerged tokens selected by a distance threshold, while new tokens are concatenated along the token dimension.

E Detailed Efficiency Derivation

MiniCache derives memory usage by leaving the first half of layers unmerged, merging later layers pairwise, and accounting for restoration vectors and retained tokens.

E Detailed Efficiency Derivation: S = 1/2r marks the midpoint where MiniCache begins consolidating every two layers into one shared state space.The layer count is r, with b denoting batch size, h hidden size, and s and n input and output sequence lengths.
E Detailed Efficiency Derivation: 3brh(s + n) is the compressed cache usage, combining 2brh(s + n) for unmerged layers with brh(s + n) for merged layers.The unmerged portion spans layers 1 through S, while the merged portion spans layers S + 1 through r.
E Detailed Efficiency Derivation: 2br(s + n) accounts for additional normalized vectors saved during restoration, because each has one channel dimension.The restoration cost is then combined with a 0.1brh(s + n) allowance for retained tokens.
E Detailed Efficiency Derivation: 3brh(s + n) + 2br(s + n) + 0.1brh(s + n) gives MiniCache’s overall memory requirement after compression and restoration costs are combined.The expression groups compressed cache usage with additional restoration memory and 5% token retention overhead.
E Detailed Efficiency Derivation: br(s + n)(3.1h + 2) is the simplified total KV-cache memory cost reported for the MiniCache framework.This form follows by grouping common factors and simplifying the combined expression.

F Detailed Experiment Results

The detailed experiment section organizes model comparisons across GSM8K, COQA, and TruthfulQA for LLaMA-3, Mixtral-8x7B, and Phi-3-Mini.

F Detailed Experiment Results: LLaMA-3-70B has detailed tables for GSM8K, COQA, and TruthfulQA.The passages identify the GSM8K comparison and separate COQA and TruthfulQA tables.
F Detailed Experiment Results: LLaMA-3-8B has detailed tables for GSM8K, COQA, and TruthfulQA.Each table compares the model with MiniCache and reports a model mean in its header.
F Detailed Experiment Results: Mixtral-8x7B has detailed tables for GSM8K, COQA, and TruthfulQA.The supplied headers identify model, MiniCache, and model-mean comparison columns.
F Detailed Experiment Results: Phi-3-Mini has detailed tables for GSM8K, COQA, and TruthfulQA.The supplied passages identify corresponding detailed comparisons for all three datasets.