Transformer key value caches show varied compressible structures

Tensor Decomposition of Transformer Key-Value Caches: Spectral Structure and Format Comparison

Computation and LanguageMachine Learning

Summary

Transformer models keep track of information in a big, four-way table spanning features like attention heads and tokens. The authors studied how this table can be squished down without losing much detail, using methods that break it apart in different ways. They found that some parts, like tokens and features, can be compressed well, but others, such as attention heads and layers, cannot. A particular approach called Tucker decomposition works best at storing the information compactly while keeping accuracy. Also, the ability to compress changes depending on whether the data represents keys or values, and certain model features reduce compressibility.

What this means in practice

Authors

Rahul Krishnan, Volker Schulz

Abstract

The key-value (KV) cache of autoregressive transformers can be viewed as a fourth-order tensor spanning attention heads, tokens, features, and grouped layers. We measure the singular-value spectra of all four mode unfoldings on Mistral-7B-v0.3 and LLaMA-2-13B and compare four standard tensor decompositions: Tucker, CP, tensor train, and t-SVD, at matched storage. The spectra partition the four axes into two classes. The token and feature modes carry low-rank structure, particularly for keys. The head and layer modes are nearly full-rank and resist compression at any practical error level. Among the four decompositions, Tucker achieves the lowest reconstruction error at every compression ratio from $2\times$ to $5\times$, because it can leave the full-rank modes untouched. Comparisons with two-dimensional unfolding baselines show that the preferred representation differs between keys and values: 2D methods achieve lower key error, while four-way Tucker achieves lower value error at matched storage. A mode-pinning theorem certifies the full-rank preservation from the measured spectra alone. Two further spectral properties affect the compressible modes without touching the full-rank ones: values reach a higher error floor than keys at every ratio, and post-RoPE keys lose $41\%$ - $64\%$ of their pre-RoPE compressibility on both models.
View PDFOpen arXiv