PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rahul Krishnan, Volker Schulz

arXiv:2609.28029v1math.NAcs.CLcs.LG

TL;DR

KV-cache multilinear structure and the effectiveness of standard tensor decompositions had not been systematically studied. The paper measures mode spectra and compares four formats at matched storage; token and feature modes are compressible, while head and layer modes are index-like, and Tucker has the lowest reconstruction error across tested ratios.

  • Problem

    The KV cache can constrain device memory and throughput, but its multilinear structure and the effectiveness of standard tensor decompositions had not been systematically studied.

    Source: §1 · pp. 1–2

  • Method

    The study measures mode-unfolding spectra and compares Tucker, CP, tensor train, and t-SVD at matched storage on caches from two production models.

    Source: §1 · p. 2, Contributions · p. 2

  • Results

    Token and feature modes are compressible while head and grouped-layer modes are index-like, and Tucker has the lowest reconstruction error at every tested ratio from 2× to 5×.

    Source: §8 · p. 14

  • Takeaways & Limitations

    Keys and values need independent rank budgets, and keys should be compressed before the rotary embedding.

    Source: §8 · p. 14

    The reconstruction sweeps use one sequence length and two models, and whether measured token-mode structure persists during incremental generation remains open.

    Source: §8 · p. 14

Abstract

from arXiv · show

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.

1 Introduction

KV caches can limit inference throughput and memory, yet their multilinear structure had not been systematically studied. This paper characterizes cache-mode spectra, compares tensor decompositions, and reports a divide between compressible token and feature modes and index-like head and layer modes.

1 Introduction

  • The KV cache stores prior keys and values for reuse, and its growth with context, batch, layers, and heads can limit device memory and throughput.Existing methods use truncated-SVD low-rank approximations to compress the cache.Source: §1 · p. 1
  • The paper addresses the lack of systematic study of KV-cache multilinear structure and whether standard tensor decompositions work effectively on it.It examines mode-wise compressibility and compares Tucker, CP, tensor train, and t-SVD reconstruction.Source: §1 · pp. 1–2
  • Across two production models, token and feature modes carry compressible structure, while head and grouped-layer modes are index-like.The study measures mode-unfolding spectra and compares four tensor decompositions at matched storage.Source: §1 · p. 2

Contributions

  • The contributions include a mode-wise spectral characterization, matched-storage format and unfolding comparisons, a mode-pinning certificate, and studies of K/V and RoPE spectral differences.The K/V comparison reveals a split not visible in any single format.Source: Contributions · p. 2

2 Background

The background defines the KV cache and its positional treatment, introduces four tensor formats, and specifies the experiments' storage, rank-allocation, and evaluation conventions.

2.1 Autoregressive Transformers

  • The KV cache stores keys and values from previous positions; RoPE rotates keys and queries by position, while values are not rotated.MHA assigns each query head its own KV head; GQA shares a KV head across multiple query heads.Source: §2.1 · p. 2

2.2 Tensor Decomposition formats

  • The paper compares Tucker, CP, tensor train, and t-SVD, which represent tensors using mode factors and a core, rank-one sums, chained cores, or Fourier-domain slice SVDs.Tucker factors are initialized by HOSVD and refined by HOOI; TT uses bond ranks between adjacent modes.Source: §2.2 · p. 3

2.3 Experimental Setup

  • Experiments use per-tensor budgets for the format sweep and joint K/V budgets for the K/V-split study; keys and values receive independent rank vectors.Storage is parameter-count based, full-rank Tucker factors are omitted, and reconstruction quality is mean relative Frobenius error averaged over layers unless stated otherwise.Source: §2.3 · pp. 3–4, §2.2 · p. 3
  • Reconstruction sweeps aggregate three prompt draws across all layers using non-overlapping WikiText-2 validation chunks at sequence length 1024.Caches and decompositions use fp16; layers within one prompt share text and are not treated as independent samples.Source: §2.3 · pp. 3–4

3 Spectral Structure of KV Caches

Mode-unfolding spectra reveal a sharp distinction: head modes resist rank reduction, while token modes and key feature modes show spectral decay; value feature modes are nearly full-rank.

3 Spectral Structure of KV Caches

  • The head mode is index-like: its smallest singular direction alone carries about 28% of energy, preventing practical rank reduction.This holds for both keys and values in the Mistral-7B layer-15 cache.Source: §3 · p. 5, Table 1 · p. 5
  • Token modes are compressible, but keys decay more steeply: rank 225 reaches 10% error for keys versus 576 for values.The last token singular direction contributes less than 10^-5 tail energy.Source: §3 · p. 5, Table 1 · p. 5
  • Feature modes are compressible for keys but nearly index-like for values: rank 126 of 128 is needed for 10% value error.The respective tail energies at rank 127 are about 0.012 for keys and 0.057 for values.Source: §3 · p. 5, Table 1 · p. 5

4 Tensor-Format Comparison

At matched storage, Tucker has the lowest reconstruction error across tested ratios on both models. Value errors remain higher than key errors across formats.

4 Tensor-Format Comparison

  • The comparison fits Tucker, CP, tensor train, and t-SVD to pre-RoPE keys and values at T = 1024 over three prompts; CP is omitted on LLaMA for memory reasons.Errors are averaged across layers at each compression ratio.Source: §4 · p. 5
  • At every compression ratio, Tucker has the lowest reconstruction error on both models; Mistral ranks CP, t-SVD, then TT, while LLaMA ranks t-SVD, then TT.The comparisons cover keys and values; CP is omitted for LLaMA because matched-storage fitting required a rank beyond the experiment device's memory.Source: §4 · p. 5, Table 2 · p. 6, Table 3 · p. 6
  • With free rank allocation, Tucker keeps head and feature ranks full and truncates only the token mode.The selected ranks are r1 = nh = 8 and r3 = dh = 128.Source: §4 · p. 5
  • Value errors exceed key errors across formats, and the gap between formats is larger for values than for keys.Figure 2 plots error against compression ratio for both models and both cache types.Source: §4 · pp. 5, 7, Figure 2 · p. 7
  • Tucker leaves the head mode untouched, while t-SVD truncates head-by-token slices and TT pays bond ranks along its mode chain.Tucker can therefore compress the token and feature axes without truncating the head mode.Source: §4 · pp. 5, 7

5 Higher-Order versus 2D Unfoldings

Across the tested representations, keys favor two-dimensional unfoldings while values favor four-way Tucker; Tucker’s shared bases and mode-specific ranks explain why the preferred format differs. Rank ablations and a spectral certificate further show why preserving full head and layer ranks matters.

5.1 Per-head SVD

  • At 2×, per-head SVD has marginally lower key error by 0.002; Tucker wins on keys at 3× and 4× and on values at every ratio.For values, Tucker’s error ranges from 0.155 versus 0.421 at 2× to 0.325 versus 0.604 at 4×.Source: §5.1 · p. 7, Table 4 · p. 8
  • Tucker shares token and feature bases across heads, storing each basis once and leaving more rank budget for those modes at matched storage.The head rank remains full throughout; per-head SVD instead pays factor-storage costs independently for each head.Source: §5.1 · pp. 7–8

5 Higher-Order versus 2D Unfoldings

  • 2D methods achieve lower key error, while four-way Tucker achieves lower value error at matched storage; the K/V split holds on both models.Among tested third-order configurations, Tucker also has the lowest error over 2×–5×.Source: §5 · p. 7, §5.2 · p. 8, Table 5 · p. 8

5.2 Cross-layer comparison: the K/V split

  • On keys, 2D methods win at every ratio, with xKV lowest from 3× onward; on values, Tucker-4D wins at every ratio and beats Tucker-3D.The comparison uses a joint K/V storage budget and fixed groups of four layers; the split also holds on LLaMA.Source: §5.2 · p. 8, Table 5 · p. 8

5.3 Head- and Layer-Rank Ablation

  • Reducing either head or layer rank sharply raises key and value errors at every tested ratio on both models; even rL = 3 nearly triples 2× key error.On Mistral, forcing r1 = nh/4 yields key errors of 0.839, 0.839, and 0.840 as the ratio relaxes from 2× to 4×, and LLaMA’s quarter-head pin is not matched-storage at 2× or 3×.Source: §5.3 · pp. 8–9, Table 6 · p. 9

5.4 A spectral certificate for full-rank preservation

  • A budget-dependent spectral condition certifies that every global minimizer preserves a specified mode at full rank; it depends only on mode spectra, dimensions, and budget.The condition cannot hold when Γk(B) = ∞, and it is sufficient but not necessary for preserving grouped-layer rank.Source: §5.4 · pp. 10, 12
  • When all modes except one are certified full rank, the tail-energy lower and upper bounds coincide; measured Egroup/Lk ratios are 1.00–1.03 in every reported row.The tightness comparison is reported for Mistral at 2× over eight groups.Source: Table 7 · p. 9, §5.4 · p. 12
  • The certificate guarantees full head rank in all tested Mistral configurations and all 224 LLaMA head-mode configurations; grouped-layer ranks remain full even in uncertified value groups.For Mistral’s head mode, the certificate-to-cost ratio ranges from 3.49 to 342.33, with a median of 13.44.Source: §5.4 · p. 12

6 Structural Properties

The cache’s compressible modes show a consistent key–value asymmetry: values have flatter spectra and higher reconstruction error than keys. This difference holds across decompositions and motivates separate rank choices for keys and values.

6 Structural Properties

  • The K/V asymmetry is identified as a spectral property of the cache, distinct from the choice of decomposition format.The paper frames this asymmetry alongside the pre/post rotary-embedding penalty as a property affecting low-rank approximability.Source: §6 · p. 12, §6.1 · p. 12

6.1 The K/V asymmetry

  • At matched compression ratios, value error is consistently higher than key error; the median value-to-key error ratio is approximately 2.4, ranging from 1.7 to 3.2.The pattern holds across formats, layers, prompt draws, models, and ratios.Source: §6.1 · p. 12
  • The value projection WV has a comparatively flat spectrum, while WK drops steeply; on Mistral, their σ1/σmin ratios are approximately 6 and 36, respectively.WV also requires higher rank than WK to capture 90% of its energy on nearly every layer.Source: §6.1 · p. 13
  • Values and keys share compressible token and feature axes, but values have a flatter spectrum, so they require independent rank budgets.The paper attributes the flatter value spectrum to the different roles of WV and WK: values contribute to attention outputs, while keys determine softmax scores.Source: §6.1 · pp. 12–13

7 Rotary Position Embedding and the Pre/Post Penalty

The paper compares key compression before and after RoPE, which preserves approximation distances but changes rank efficiency. Across both models, post-RoPE keys have higher error, consistent with rotation flattening token- and feature-mode spectra.

7 Rotary Position Embedding and the Pre/Post Penalty

  • RoPE preserves Frobenius distances to a fixed approximation but can reduce rank efficiency by flattening token- and feature-mode spectra.The position-dependent orthogonal rotation mixes coordinates within each pair.Source: §7 · p. 13
  • Post-RoPE key error is 41%–64% higher than pre-RoPE error across both models and five compression ratios.A frozen-allocation control produces even larger gaps; the paper attributes the penalty to rotation rather than allocator choices.Source: §7 · p. 13, Table 8 · p. 13
  • On Mistral, top-8 token-mode energy share falls from 82.5% to 65.3%, while feature-mode share falls from 38.9% to 33.2%.The changes are consistent across layer groups; Figure 3 reports group means and ranges.Source: §7 · p. 13, Figure 3 · p. 13
  • The token-mode flattening is larger because position-dependent angles vary across the sequence, spreading energy beyond leading token-mode components.The feature spectrum also flattens, shrinking keys’ compression advantage over values in that mode.Source: §7 · p. 14

8 Conclusions

The measured spectra divide the cache into compressible token and feature modes and index-like head and grouped-layer modes. Tucker performs best among the compared decompositions, while value spectra and post-RoPE keys present additional compression challenges.

8 Conclusions

  • Across target ratios of 2×–5× on both models, Tucker has the lowest reconstruction error because it leaves index-like modes untouched and allocates budget to compressible modes.Theorem 5.1 certifies the mode partition from the spectra alone for every tested head-mode configuration.Source: §8 · p. 14
  • Values have a higher reconstruction-error floor than keys, and post-RoPE keys lose 41%–64% of their pre-RoPE compressibility on both models.The paper attributes these differences to the flatter value-projection spectrum and rotation-induced flattening of token- and feature-mode spectra; it concludes that keys and values need independent rank budgets and keys should be compressed before RoPE.Source: §8 · p. 14
  • The reconstruction sweeps use only T = 1024, spectral analysis covers two models, and persistence of filled-prompt token structure during incremental generation remains open.Source: §8 · p. 14

Reproducibility Statement

Code and data for reproducing the paper’s results are available at the stated project repository.

Reproducibility Statement

AI Use Statement

The authors used generative AI for coding, manuscript drafting, and related-work and citation checks, then reviewed and verified all AI-assisted code and text and accepted responsibility for the paper’s final content.

AI Use Statement

  • Generative AI assisted with software implementation, manuscript drafting, and retrieving related work and checking citation records against publisher metadata.Source: AI Use Statement · p. 14
  • The authors reviewed and verified all AI-assisted code and text.Source: AI Use Statement · p. 14
  • The authors take full responsibility for the paper’s final content.Source: AI Use Statement · p. 14
Loading 2609.28029v1…