PaperlayerView exact revision on arXiv

Source-linked AI summary

Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time

Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, Anshumali Shrivastava

arXiv:2305.17118v2cs.LGcs.CL

TL;DR

Large language model serving is limited by KV-cache memory, which can exceed model-weight memory and constrain inference batch size. Scissorhands uses persistent token importance to maintain a fixed-budget cache without finetuning, combining empirical, theoretical, and system-level analysis. It reports up to 5× KV-cache memory reduction without compromising quality and demonstrates compatibility with quantization for up to 20× compression.

  • Problem

    KV-cache memory can exceed model-weight memory and constrain batch size, creating a central memory bottleneck for high-throughput LLM inference.

  • Method

    Scissorhands maintains a fixed-budget KV cache without finetuning by preferentially storing tokens identified as persistently important.

  • Results

    Up to 5× KV-cache memory reduction is achieved without compromising model quality, with accuracy maintained at substantial compression levels.

  • Takeaways & Limitations

    The method supports memory reduction across both representation and sequence-length dimensions through compatibility with quantization.

  • Takeaways & Limitations

    The largest experimentally fitted model was OPT-66B, and the authors could not access its training process to determine how the observed behavior arose.

Abstract

from arXiv · show

Large language models(LLMs) have sparked a new wave of exciting AI applications. Hosting these models at scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights are memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints on the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had a substantial influence at one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis, we propose Scissorhands, a system that maintains the memory usage of the KV cache at a fixed budget without finetuning the model. In essence, Scissorhands manages the KV cache by storing the pivotal tokens with a higher probability. We validate that Scissorhands reduces the inference memory usage of the KV cache by up to 5X without compromising model quality. We further demonstrate that Scissorhands can be combined with 4-bit quantization, traditionally used to compress model weights, to achieve up to 20X compression.

1 Introduction

LLM deployment is constrained by large, context-dependent KV-cache memory, motivating training-free compression. Scissorhands exploits persistent token importance, supported by empirical attention overlap, to compress the cache while preserving quality.

  • Memory bottleneck: At batch size 128 and sequence length 2048, OPT-175B requires around 950 GB for its KV cache versus around 325 GB for model parameters.The KV cache includes prompts and previously generated tokens.
  • Memory bottleneck: KV-cache compression translates almost linearly into larger inference batches at a fixed sequence length, benefiting high-throughput systems.Hardware memory is largely allocated to the KV cache after model weights and a small buffer.
  • Compression challenge: KV-cache compression remains challenging because the ideal method should require no retraining and support context lengths exceeding 32K tokens.The cache grows linearly with sequence length.
  • Persistence of importance: The Persistence of Importance Hypothesis proposes that tokens with substantial influence at one step will significantly influence future generations.The hypothesis builds on sparse, repetitive attention patterns.
  • Persistence of importance: Over 90% of later-token attention overlaps with tokens heavily attended by earlier tokens in most transformer layers.This empirical pattern motivates forecasting future-important tokens from earlier attention.
  • Scissorhands: Scissorhands maintains a predetermined KV-cache budget by exploiting persistent importance, reducing cache memory up to 5× without degrading model quality.The system is designed for inference without finetuning and is compatible with quantization.

2 Problem Description and Related Work

This section defines autoregressive KV-cache updates and situates KV-cache memory as a deployment bottleneck. Existing efficient-attention and compression methods largely target computation or model weights rather than sequence-length cache memory.

  • Inference workflow: At transformer layer i, key and value caches store embeddings with dimensions determined by batch size, prompt length, and hidden dimension.The cache is represented as Ki, Vi ∈ R^(b×p×d).
  • Inference workflow: LLM inference comprises prompting and token-generation stages, with key/value embeddings stored and updated in the KV cache.The cache avoids repetitive key-value projection computation during future generation steps.
  • Memory breakdown: KV-cache memory depends on model configuration, sequence length, and batch size, while the activation buffer is considerably smaller than model weights and cache.The memory footprint consists of model weights, KV cache, and activation buffer.
  • Memory breakdown: At GPT-3 scale with maximum sequence length 2048, batch size cannot exceed 35 on eight A100 80GB GPUs without offloading.The section connects small batch sizes with limited inference throughput.
  • Related work: Attention-matrix computation has O(n^2) time complexity, where n is sequence length, motivating low-rank and sparsification approaches.These approaches approximate attention output rather than directly addressing autoregressive KV-cache memory.
  • Related work: Quantization and pruning methods mostly reduce model-weight size, while FlexGen quantizes KV values but retains all sequence tokens in CPU memory.FlexGen still loads all attention keys from CPU memory to compute attention scores.

3 The Persistence of Importance Hypothesis

The paper hypothesizes that tokens receiving substantial attention remain influential in future steps, and tests this through repetitive attention patterns and persistence ratios. It then uses this property to motivate budgeted KV-cache compression supported by empirical and theoretical analysis.

  • 3.1 Repetitive Attention Pattern: Attention maps show repetitive patterns: specific tokens receive high attention from multiple positions, while only a few tokens have high scores.The same high-attention positions recur across sampled positions, layers, and text inputs.
  • 3.2 The Persistence of Importance Hypothesis: The Persistence of Importance Hypothesis states that pivotal tokens substantially influential at one previous step will significantly influence a future step.The hypothesis is nontrivial when pivotal tokens form a subset of previous words rather than the entire sentence.
  • 3.2 The Persistence of Importance Hypothesis: Pivotal tokens are defined as tokens receiving an attention score above threshold α from the token at position t.The pivotal-token set S_t contains the tokens meeting this criterion for position t.
  • 3.2 The Persistence of Importance Hypothesis: Persistence ratios generally exceed 95%, while pivotal-token sets remain considerably smaller than sentence length, supporting the hypothesis empirically.The ratio decreases in later transformer layers, and later-half pivotal tokens are mostly included in first-half sets.
  • 3.3 Attention Weights Decides the Pivotal Tokens: The theoretical analysis links persistence to attention-weight behavior, while noting that the guarantee holds only for tokens satisfying a specific large quadratic-form condition.Skip connections can make input and output representations highly aligned, so a larger current attention score can imply a potentially large future score.
  • 3.2 The Persistence of Importance Hypothesis: The proposed inference procedure maintains a KV cache under budget B by updating the cache, compressing it when n > B, and retaining tokens selected by importance records.The compression algorithm tracks low-score-token occurrences over a history window, preserves a recent window, and removes m tokens according to the recorded importance.

4 Sequential Token Generation Under budget

SCISSORHANDS compresses the KV cache under a fixed token budget by retaining influential tokens, using attention history to approximate full-cache attention during sequential generation.

  • Approach: SCISSORHANDS reduces KV-cache memory from the sequence-length dimension without fine-tuning the model.The method maintains the cache under a predetermined budget and provides theoretical analysis of approximation error.
  • Motivation: The persistence-of-importance hypothesis states that tokens contributing strongly to previous outputs will also contribute strongly to future tokens.This persistence supplies the basis for identifying important tokens before later generation steps.
  • Budgeted KV cache: Under budget B, the algorithm maintains a key cache that stores embeddings for at most B previous tokens.The sequential-generation problem is defined as maintaining the cache under a fixed memory budget.
  • Approach: SCISSORHANDS uses attention scores as influence indicators and drops stored non-influential tokens when its fixed buffer is full.Influence is accumulated over a history window to reduce variance, while recent tokens are retained.
  • Algorithm behavior: When m = 1, each iteration drops one lowest-score token, so the cache continuously maintains B tokens.If attention-score rankings remain unchanged, the algorithm removes tokens with the smallest scores.
  • Theoretical analysis: The theoretical error bound scales with 1 − B/Tmax and becomes zero when B = Tmax.The bound concerns the expected difference between tokens generated with the budgeted and original approaches.

5 Empirical Evaluation

Experiments compare SCISSORHANDS-OPT with original OPT on language modeling and few-shot downstream tasks, including evaluations of compression and 4-bit quantization compatibility.

  • Main results: Up to 5× KV-cache compression incurs no model-quality degradation compared with the standard model.The evaluation covers language modeling and downstream tasks, with original OPT as the comparison point.
  • Main results: Accuracy is generally maintained with 15%–30% of the original KV-cache size.For OPT-66B, Winogrande and MathQA retain accuracy even after 5× compression.
  • Language modeling: For OPT-66B, perplexity remains stable until 75% of the original KV-cache size.The reported language-modeling trend becomes flatter as model size increases.
  • Quantization: At 2× compression, adding 4-bit quantization does not introduce compounded errors, including on the more sensitive Hellaswag task.The study explicitly tests compatibility between quantization and SCISSORHANDS.
  • Attention-score ablation: At 3× compression on C4 with OPT-13B, SCISSORHANDS attention scores are almost the same as those from the original KV cache.The change ratio is centered around zero, although a small portion of important tokens may be dropped.

6 Discussion, Limitation, and Future Work

The discussion examines repetitive attention in trained models, limits the evaluation to OPT-66B, and notes unresolved questions about training origins and generation quality.

  • Repetitive attention appears in trained language models but is absent in randomly initialized OPT models.
  • The relationship between repetitive attention patterns and undesirable generations remains an open question.
  • OPT-66B is the largest model evaluated because of academic server limitations.
  • The authors cannot determine exactly how language-model training produces the observed behavior.
  • The authors report no negative impacts while noting that experiments with large models produce carbon dioxide emissions.

7 Conclusion

SCISSORHANDS uses the lasting influence of pivotal tokens to reduce KV-cache memory without compromising LLM performance, and is compatible with quantization.

  • SCISSORHANDS reduces KV-cache memory usage by 5× without compromising LLM performance.
  • SCISSORHANDS leverages the lasting influence of pivotal tokens to reduce KV-cache memory.
  • The method is compatible with quantization techniques for reducing memory in representation and sequence-length dimensions.
  • Figure 5 shows no repetitive attention for a randomly initialized OPT model.

A.1 Repetitive Attention Pattern

The appendix reports repetitive and sparse attention patterns across transformer layers and supports the assumption that x and F(x) have high cosine similarity.

  • A.1 Repetitive Attention Pattern: Repetitive attention patterns and attention sparsity are observed across transformer layers.
  • A.1 Repetitive Attention Pattern: Figure 10 provides empirical evidence for the large cosine-similarity assumption between x and F(x).
  • A.1 Repetitive Attention Pattern: Residual connections and domination by x explain the extremely high cosine similarity between x and F(x).

B.1 Proof of Theorem 3.1

The proof analyzes a simplified single-layer, single-head transformer under normalized-input and similarity assumptions, using lemmas and singular-value bounds.

  • B.1 Proof of Theorem 3.1: The proof considers token generation in a single-layer transformer with single-head attention.
  • B.1 Proof of Theorem 3.1: The model represents token states as row vectors and defines aggregated prior states, attention weights, and the MLP block F.
  • B.1 Proof of Theorem 3.1: The theorem uses largest singular values of WQ, WK, WV, and WO in its assumptions and bounds.
  • B.1 Proof of Theorem 3.1: Inputs are assumed normalized, with ∥x_t∥2 = 1 for all t.
  • B.1 Proof of Theorem 3.1: The proof first establishes lemmas using unit-norm vectors and Cauchy-Schwarz before proceeding to the main argument.

B.2 Proof of Theorem 4.1

The proof bounds the discrepancy between Scissorhands’ sampled transformer and the vanilla transformer under a fixed KV-cache budget. It combines attention and MLP Lipschitz bounds with probabilistic guarantees for retaining high-score tokens.

  • Cache budget: With m = 1, the cache maintains B tokens by dropping one token with the lowest attention score at each iteration.If attention-score rankings remain unchanged, the algorithm continually removes the smallest-scoring tokens.
  • Assumptions: The analysis assumes retained sets contain only tokens among the largest B values of βt,j at every timestep.The theorem also imposes conditions involving singular values, a power-law sampling distribution, and a bounded product of model constants.
  • Proof components: The attention module is analyzed through a Lipschitz bound comparing two token sequences and their resulting outputs.The proof decomposes the attention-output difference and bounds its components using softmax Lipschitzness and the sampling mask.
  • Proof components: The sampled-transformer analysis separately bounds the difference between sampled and vanilla attention outputs using the retained-token set.Tokens outside the retained set have zero sampled attention weight, while retained tokens are analyzed through the corresponding score terms.
  • Proof components: The MLP in (6) is Lipschitz with factor 1 + λ1λ2, using the largest singular values of W1 and W2 and ReLU’s 1-Lipschitz property.This bound is incorporated into the main theorem through a triangle-inequality decomposition of the overall error.
  • Theoretical guarantee: Theorem B.3 gives a high-probability bound on the discrepancy between successive sampled and vanilla hidden states under the power-law score model.The stated expectation bound depends on λ1, λ2, λV, λO, λQ, λK, and Δmax; a union bound supplies the final guarantee.
Loading 2305.17118v2…