Most KV-cache compression applies SVD to the keys alone, or embeds queries and keys jointly. Both miss the obvious target: attention is the product of queries and keys, so that product is what should be approximated.

KQ-SVD (Mila) computes the optimal low-rank decomposition of the attention product directly, in closed form - no training, one calibration pass. The paper proves it is strictly more accurate than key-only SVD, immune to a rescaling failure that breaks EigenAttention (scale keys up and queries down and it quietly degenerates to key-only SVD), and still optimal under grouped-query attention. At inference you store two narrow compressed caches instead of full-width keys and queries.

On LLaMA-2-7B/13B, LLaMA-3-8B, and Mistral-7B, all methods look similar on keys, queries, and values in isolation - but on the attention matrix itself, the thing that determines the output, KQ-SVD cuts error 5-10x versus key-only SVD.

Made a short visual breakdown - one chart per trick. Swipe through.

paper - https://arxiv.org/abs/2512.05916 full summary pdf - https://gist.github.com/vukrosic/75bbdf1277a856a17bedf7dd9f94f91c…

Every Sunday I run a hands-on live AI research with 1 on 1 help: https://skool.com/become-ai-researcher-2669/about…

1 Introduction

Source: https://arxiv.org/html/2512.05916 KQ-SVD: Compressing the KV Cache with Provable Guarantees on Attention Fidelity

Damien LesensBeheshteh T. RakhshanGuillaume Rabusseau

ENS de LyonDIRO, Université de MontréalMilaDIRO, Université de MontréalMila - CIFAR AI Chair

Abstract

The Key–Value (KV) cache is central to the efficiency of transformer-based large language models (LLMs), storing previously computed vectors to accelerate inference. Yet, as sequence length and batch size grow, the cache becomes a major memory bottleneck. Prior compression methods typically apply low-rank decomposition to keys alone or attempt to jointly embed queries and keys, but both approaches neglect that attention fundamentally depends on their inner products. In this work, we prove that such strategies are sub-optimal for approximating the attention matrix. We introduce KQ-SVD, a simple and computationally efficient method that directly performs an optimal low-rank decomposition of the attention matrix via a closed-form solution. By targeting the true source of redundancy, KQ-SVD preserves attention outputs with higher fidelity under compression. Extensive evaluations on LLaMA and Mistral models demonstrate that our approach consistently delivers superior projection quality.

The rise of Large Language Models (LLMs)[touvron2023llama,chaplot2023albert,achiam2023gpt,guo2025deepseek]has expanded AI capabilities beyond earlier models. Transformers[vaswani2017attention]replace recurrence with self-attention, enabling parallelism and improved sequence modeling, but their quadratic memory and computation costs limit long-sequence scalability.

Key-Value (KV) caches are introduced to accelerate autoregressive generation by storing intermediate attention KV vectors, avoiding redundant computation of shared prefixes for each generated token. Although KV caching reduces computational overhead, it substantially increases memory consumption, as the cache size grows linearly with both sequence length and batch size. This trade-off motivates the development of KV cache compression techniques, which are crucial for enabling efficient and cost-effective deployment of LLMs across diverse hardware platforms[fu2024challenges,shi2024keep]. Variants like Multi-Query Attention (MQA)[shazeer2019fast]and Grouped-Query Attention (GQA)[ainslie2023GQA]reduce KV cache size by sharing or grouping query vectors while maintaining performance comparable to full Multi-Head Attention (MHA). However, they may introduce accuracy trade-offs and hardware sensitivity, which can affect performance generalization. Additional approaches, including sparse[zhang2021sparse]and linearized[katharopoulos2020transformers]attention, further reduce computational and memory costs, shaping KV cache optimization strategies.

Another promising line of research exploits the low-rank structure of KV caches to reduce memory overhead. Multi-Head Latent Attention (MLA)[liu2024deepseek,guo2025deepseek]maps tokens into the low-rank latent space and stores these compressed representations in place of the original key and value states. However, using MLA necessitates training the model from the ground up. In contrast, ASVD[yuan2024asvd], LoRC[zhang2024lorc], and Palue[changpalu]apply SVD to key-value parameter matrices without retraining to build low-rank projection modules. A key limitation of these approaches is that they often compress only the keys, neglecting the query-key interaction that underlies attention. EigenAttention[saxena2024eigen]and Zack[zhang2024zack]attempt to address this by incorporating both queries and keys in low-rank decompositions, yet their behavior largely resembles that of SVD-based methods that compress keys alone.

In this work, we address these limitations by introducing KQ-SVD, a compression method that achieves optimal low-rank approximation of the attention matrix efficiently and in closed form. Our method explicitly captures the interactions between queries and keys through their inner products, preserving the fundamental structure of attention. Beyond key-query interactions, we also consider the corresponding interactions between values and the output projection, enhancing the fidelity of the approximation. By leveraging the inherent low-rank structure of KV caches[yu2024effectively,saxena2024eigen], we formulate attention matrix approximation as a principled low-rank decomposition problem. Our theoretical analysis quantifies the error between prior key-only SVD approaches and our optimal method, and shows that methods incorporating both queries and keys can degrade when keys and queries are rescaled by the same factor, effectively behaving like key-only SVD methods. Our contributions can be summarized as follows:

• •We introduce KQ-SVD, an optimal low-rank approximation of the attention matrix capturing key-query interactions.
• •We theoretically quantify the advantages of KQ-SVD over methods based on key low-rank decomposition and SVD on concatenated queries and keys.
• •We show that KQ-SVD is compatible with and also optimal in the Grouped-Query Attention setting.
• •We provide extensive empirical evaluations with LLaMA2-7B, LLaMA2-13B, LLama3-8B and Mistral-7B models on the C4 dataset demonstrating significant advantages of KQ-SVD over existing low rank projection methods.

2Related Works

Low-rank structure of the KV-cache.

Several methods exploit the inherent low-rank structure of cached key–value (KV) matrices to reduce memory footprint. ECKVH[yu2024effectively]compresses the cache by grouping attention heads, performing singular value decomposition (SVD) within each group, and retaining only the dominant singular components. EigenAttention[saxena2024eigen]generalizes this idea by constructing low-rank bases that jointly approximate queries, keys, and values, effectively lowering the dimensionality of KV representations. Q-Filters[godeyq]introduces a training-free variant, projecting keys into a low-rank subspace via SVD to approximate attention scores efficiently with minimal accuracy loss. Moreover,[yu2024effectively]investigates the intrinsic low-rank nature of KV caches and compresses KV heads through careful grouping and SVD-based decomposition. In contrast, Loki[singhania2024loki]adopts a two-stage strategy: it first estimates approximate attention scores in a lower-dimensional space to rank and select the most relevant keys, and then computes exact attention scores using only the selected keys, reducing both memory and computational cost.

KV Weights Compression.

An alternative approach targets the KV weight matrices themselves rather than the cached matrices. LoRC[zhang2024lorc]applies low-rank approximations directly to the key and value weight matrices, achieving compression at the parameter level. Palu[changpalu]follows a similar strategy, jointly compressing key and value weight matrices via SVD. ShadowKV[sun2024shadowkv]introduces a distinct perspective by performing SVD on pre-RoPE key matrices to reduce their dimensionality, demonstrating the versatility of low-rank methods in optimizing KV representations.

Positioning of KQ-SVD.

Although prior methods have made significant strides in compressing KV caches and their weight matrices, they often treat keys and values independently or approximate attention only indirectly, leaving the core query–key interactions underrepresented and leading to sub-optimal low rank approximation of attention matrices. KQ-SVD addresses these limitations by formulating a principled, closed-form low-rank approximation of the full attention matrix.

3Preliminary

In this section, we introduce our notations and present the necessary background on Multi-Head Attention (MHA).

3.1Notations

We use lower case bold letters for vectors (e.g.,𝐚,𝐛\mathbf{a},\mathbf{b}), upper case bold letters for matrices (e.g.,𝐀,𝐁\mathbf{A},\mathbf{B}).𝐀+\mathbf{A}^{+}denotes the Moore–Penrose pseudo-inverse of𝐀\mathbf{A}. Throughout the paper, the singular value decomposition (SVD) of a matrix𝐒∈ℝm×n\mathbf{S}\in\mathbb{R}^{m\times n}is presented by𝐒=𝐔​𝚺​𝐕⊤\mathbf{S}=\mathbf{U}\mathbf{\Sigma}\mathbf{V}^{\top}, with𝐔∈ℝm×n\mathbf{U}\in\mathbb{R}^{m\times n},𝐕∈ℝn×n\mathbf{V}\in\mathbb{R}^{n\times n}matrices with orthonormal columns and𝚺∈ℝn×n\mathbf{\Sigma}\in\mathbb{R}^{n\times n}a diagonal matrix with positive diagonal entries{σi}i=1n\{\sigma_{i}\}_{i=1}^{n}. The columns of𝐔\mathbf{U}and𝐕\mathbf{V}are called respectively the left and right singular vectors of𝐒\mathbf{S}, and theσi\sigma_{i}’s the singular values of𝐒\mathbf{S}, notedσi​(𝐒)\sigma_{i}(\mathbf{S}). The optimal rank-RRapproximation of𝐒\mathbf{S}with respect to the Frobenius norm can be obtained via the SVD by truncating it to keep only the firstRRsingular vectors and singular values:𝐒≈𝐔^​𝚺^​𝐕^⊤{\mathbf{S}}\approx\hat{\mathbf{U}}\hat{\mathbf{\Sigma}}\hat{\mathbf{V}}^{\top}with𝐔^∈ℝm×R\hat{\mathbf{U}}\in\mathbb{R}^{m\times R},𝐕^∈ℝn×R\hat{\mathbf{V}}\in\mathbb{R}^{n\times R}and𝚺^∈ℝR×R\hat{\mathbf{\Sigma}}\in\mathbb{R}^{R\times R}. The column space of𝐒\mathbf{S}is notedℛ​(𝐒)\mathcal{R}(\mathbf{S}).

3.2Background

In transformer architectures, self-attention assigns relative importance to tokens, enabling the model to selectively focus on different segments of the input sequence. For a sequence of token embeddings𝐗∈ℝT×D\mathbf{X}\in\mathbb{R}^{T\times D}, multi-head attention is computed as

MHA​(𝐗)=[𝐇1,…,𝐇h]​𝐖O,\mathrm{MHA}(\mathbf{X})=[\mathbf{H}_{1},\dots,\mathbf{H}_{h}]\mathbf{W}^{O},where

𝐇i=Softmax​(𝐐i​𝐊i⊤d)​𝐕i,\mathbf{H}_{i}=\mathrm{Softmax}\Big(\frac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}}{\sqrt{d}}\Big)\mathbf{V}_{i},with𝐖iQ,𝐖iK,𝐖iV∈ℝD×d\mathbf{W}_{i}^{Q},\mathbf{W}_{i}^{K},\mathbf{W}_{i}^{V}\in\mathbb{R}^{D\times d},d=D/hd=D/h,𝐐i=𝐗𝐖iQ\mathbf{Q}_{i}=\mathbf{X}\mathbf{W}_{i}^{Q},𝐊i=𝐗𝐖iK\mathbf{K}_{i}=\mathbf{X}\mathbf{W}_{i}^{K},𝐕i=𝐗𝐖iV\mathbf{V}_{i}=\mathbf{X}\mathbf{W}_{i}^{V}, and𝐖O∈ℝD×D\mathbf{W}^{O}\in\mathbb{R}^{D\times D}. In masked attention, the upper-diagonal entries of the attention matrix𝐐i​𝐊i⊤\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}are set to−∞-\inftyto prevent a token from attending to future positions.

This computation scales quadratically with the sequence lengthTT. In auto-regressive decoding, previously computed key and value vectors are cached to avoid redundant computation, reducing the per-token cost. Specifically at timeTT, for each head, new key–value pairs are concatenated to the caches𝐊\mathbf{K}and𝐕\mathbf{V}, followed by an attention computation:

𝐊←\displaystyle\mathbf{K}\leftarrowConcat​(𝐊,𝐤T),𝐕←Concat​(𝐕,𝐯T),\displaystyle\text{Concat}(\mathbf{K},\mathbf{k}_{T}),\ \mathbf{V}\leftarrow\text{Concat}(\mathbf{V},\mathbf{\bm{v}}_{T}),𝐡T=Softmax​(𝐪T​𝐊⊤d)​𝐕\displaystyle\mathbf{h}_{T}=\text{Softmax}\Big(\frac{\mathbf{q}_{T}\mathbf{K}^{\top}}{\sqrt{d}}\Big)\mathbf{V}where𝐤T=𝐱T​𝐖K\mathbf{k}_{T}=\mathbf{x}_{T}\mathbf{W}^{K},𝐯T=𝐱T​𝐖V\mathbf{\bm{v}}_{T}=\mathbf{x}_{T}\mathbf{W}^{V},𝐪T=𝐱T​𝐖Q\mathbf{q}_{T}=\mathbf{x}_{T}\mathbf{W}^{Q}. While caching mitigates redundant computation, generating theTT-th token still incurs𝒪​(T)\mathcal{O}(T)cost, and the memory footprint of stored keys and values grows linearly with sequence length. For sufficiently long contexts, this memory requirement becomes a dominant bottleneck, as the cumulative size of the KV cache can surpass the model parameters by several orders of magnitude. In the following sections, we briefly review the K-SVD method[changpalu,yu2024effectively,zhang2024lorc], which compresses key representations using singular value decomposition (SVD), and the Eigen approach[saxena2024eigen], which jointly considers keys and queries by vertically concatenating them and applying SVD to the resulting matrix.

3.3Cache compression with SVD

Recent works[changpalu,zhang2024lorc,chang2025xkv]demonstrate that singular value decomposition (SVD) is a powerful tool for compressing the KV cache in large language models, as the cache exhibits low-rank structure. Let𝐊=𝐔K​𝚺K​𝐕K⊤∈ℝT×d\mathbf{K}=\mathbf{U}_{K}\mathbf{\Sigma}_{K}\mathbf{V}_{K}^{\top}\in\mathbb{R}^{T\times d}be the SVD of the key matrix, and let𝐊~=𝐔^K​𝚺^K​𝐕^K⊤\widetilde{\mathbf{K}}=\hat{\mathbf{U}}_{K}\hat{\mathbf{\Sigma}}_{K}\hat{\mathbf{V}}_{K}^{\top}denote its rank-RRtruncated version. By the Eckart–Young–Mirsky theorem,𝐊~\widetilde{\mathbf{K}}is the best rank-RRapproximation of𝐊\mathbf{K}under the Frobenius norm. In other words, the optimization problem

min𝐏∈ℝd×d⁡‖𝐊𝐏−𝐊‖F2s.t.rank​(𝐏)≤R,\min_{\mathbf{P}\in\mathbb{R}^{d\times d}}\|\mathbf{K}\mathbf{P}-\mathbf{K}\|_{F}^{2}\quad\text{s.t.}\quad\mathrm{rank}(\mathbf{P})\leq R,is solved by𝐏=𝐕^K​𝐕^K⊤\mathbf{P}=\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}, leading to the approximation𝐊~=𝐊​𝐕^K​𝐕^K⊤=𝐔^K​𝚺^K​𝐕^K⊤\widetilde{\mathbf{K}}=\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}=\hat{\mathbf{U}}_{K}\hat{\mathbf{\Sigma}}_{K}\hat{\mathbf{V}}_{K}^{\top}. Applying the same procedure to the value matrix𝐕=𝐔V​𝚺V​𝐕V⊤\mathbf{V}=\mathbf{U}_{V}\mathbf{\Sigma}_{V}\mathbf{V}_{V}^{\top}, we can approximate the attention output as

𝐇~\displaystyle\widetilde{\mathbf{H}}=Softmax​(𝐐​𝐊~⊤/d)​𝐕~\displaystyle=\mathrm{Softmax}(\mathbf{Q}\widetilde{\mathbf{K}}^{\top}/\sqrt{d})\widetilde{\mathbf{V}}=Softmax​(𝐐​𝐕^K​𝐕^K⊤​𝐊⊤/d)​𝐕​𝐕^V​𝐕^V⊤.\displaystyle=\mathrm{Softmax}(\mathbf{Q}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}\mathbf{K}^{\top}/\sqrt{d})\mathbf{V}\hat{\mathbf{V}}_{V}\hat{\mathbf{V}}_{V}^{\top}.This formulation is particularly useful because it allows to store only the compressed caches𝐊​𝐕^K\mathbf{K}\hat{\mathbf{V}}_{K}and𝐕​𝐕^V\mathbf{V}\hat{\mathbf{V}}_{V}in memory. These matrices are of sizeR×TR\times Tinstead ofd×Td\times T, resulting in significant memory savings. At runtime, queries are multiplied by𝐕^K\hat{\mathbf{V}}_{K}, while𝐕^V⊤\hat{\mathbf{V}}_{V}^{\top}can be absorbed into the output projection𝐖O\mathbf{W}_{O}, streamlining computation.

A key advantage of this approach is that the SVD does not need to be computed during token generation. Instead, it can be performed once in a post-training calibration phase. For each layerlland attention headii, we only need to determine a basis𝐕i,l∈ℝd×R\mathbf{V}_{i,l}\in\mathbb{R}^{d\times R}such that𝐊i,l≈𝐊i,l​𝐕i,l​(𝐕i,l)⊤\mathbf{K}_{i,l}\approx\mathbf{K}_{i,l}\mathbf{V}_{i,l}(\mathbf{V}_{i,l})^{\top}with an analogous construction for the values. This basis can be learned from a calibration set of sequences. Specifically, we passnsn_{s}calibration sequences (e.g., sampled from a high-quality dataset such as C4[C4]) through the model. Thekkth sequence produces caches𝐊i,lk\mathbf{K}_{i,l}^{k}and𝐕i,lk\mathbf{V}_{i,l}^{k}for every layerlland headii. These are concatenated to form large cache matrices𝐊i,l=(𝐊i,l1,𝐊i,l2,…,𝐊i,lns)\mathbf{K}_{i,l}=\begin{pmatrix}\mathbf{K}_{i,l}^{1},\mathbf{K}_{i,l}^{2},\dots,\mathbf{K}_{i,l}^{n_{s}}\end{pmatrix}. This aggregated cache provides a representative sample of the key vectors that will appear during inference. Performing SVD on𝐊i,l\mathbf{K}_{i,l}then yields the dominant singular vectors, which form a suitable low-rank basis for compression. The cost of generating calibration caches and computing the SVDs is negligible compared to model training, and is offset by the runtime speedups from cache compression. In the following, we will refer to this method as K-SVD.

Rank selection.The compression rankRRis determined per layer by examining the singular value spectrum. Let{σj}j\{\sigma_{j}\}_{j}denote the singular values of a matrix𝐌\mathbf{M}. For a relative error toleranceϵ\epsilon, we select the smallestRRsuch that

‖𝐌−𝐌~‖F2≤ϵ​‖𝐌‖F2⇔∑j=1Rσj2∑i=1dσj2≥1−ϵ.\|\mathbf{M}-\widetilde{\mathbf{M}}\|_{F}^{2}\leq\epsilon\|\mathbf{M}\|_{F}^{2}\Leftrightarrow\frac{\sum_{j=1}^{R}\sigma_{j}^{2}}{\sum_{i=1}^{d}\sigma_{j}^{2}}\geq 1-\epsilon.Prior studies[yu2024effectively,saxena2024eigen]have shown that KV matrices are indeed approximately low-rank, so substantial compression can be achieved with small error budgetϵ\epsilon. The chosen rank may differ for keys and values depending on their spectra.

3.4Cache compression with Eigen

Other works[saxena2024eigen,zhang2024zack]emphasize that queries should also be considered when compressing key caches. Indeed, by projecting the key cache we also project the query matrix:𝐐​𝐕^K​𝐕^K⊤​𝐊⊤=(𝐐​𝐕^K​𝐕^K⊤)​(𝐊​𝐕^K​𝐕^K⊤)⊤\mathbf{Q}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}\mathbf{K}^{\top}=(\mathbf{Q}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top})(\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top})^{\top}as𝐕^K​𝐕^K⊤\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}is an idempotent matrix. Hence, it makes sense to compute the low rank projection by solving

min𝐒⁡‖𝐊−𝐊𝐒‖F2+‖𝐐−𝐐𝐒‖F2s.t.rank​(𝐒)≤R,\min_{\mathbf{S}}\|\mathbf{K}-\mathbf{K}\mathbf{S}\|_{F}^{2}+\|\mathbf{Q}-\mathbf{Q}\mathbf{S}\|_{F}^{2}\quad\text{s.t.}\quad\mathrm{rank}(\mathbf{S})\leq R,which is equivalent to

min𝐒⁡‖[𝐊𝐐]−[𝐊𝐐]​𝐒‖F2s.t.rank​(𝐒)≤R,\min_{\mathbf{S}}\left\|\begin{bmatrix}\mathbf{K}\\ \mathbf{Q}\end{bmatrix}-\begin{bmatrix}\mathbf{K}\\ \mathbf{Q}\end{bmatrix}\mathbf{S}\right\|_{F}^{2}\quad\text{s.t.}\quad\mathrm{rank}(\mathbf{S})\leq R,so that𝐒\mathbf{S}will approximate queries and keys simultaneously. As the second formulation of the optimization problem shows,𝐒\mathbf{S}can be computed by performing an SVD on the combined matrix[𝐊𝐐]\begin{bmatrix}\mathbf{K}\\ \mathbf{Q}\end{bmatrix}. This approach ensures that the learned projection preserves keys and queries while reducing dimensionality. We refer to this approach as Eigen throughout this paper. The calibration process follows the same procedure as K-SVD: large calibration caches are formed by using a collection of calibration sequences. The only difference is that query matrices from the calibration set are also used in the projection computation.

4Methodology

In this section, we introduce our proposed approach KQ-SVD for KV cache compression. We consider taking the interaction between queries and keys into account. The method views the key and query matrices as a single entity and applies singular value decomposition (SVD) to𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}. We begin by outlining the motivation behind this idea, followed by a detailed explanation of the technique.

4.1Motivation

Existing compression methods based on SVD typically compress keys or jointly embed𝐐/𝐊/𝐕\mathbf{Q}/\mathbf{K}/\mathbf{V}. Theorem1(proof in AppendixA) inspired by[wang2025squat]shows why that can fail: perturbations in𝐊\mathbf{K}are amplified by the inner products𝐐𝐊⊤\mathbf{Q}\mathbf{K}^{\top}and further propagated by the value multiplication.

Theorem 1.

Let𝐗∈ℝT×d\mathbf{X}\in\mathbb{R}^{T\times d}be a sequence of token embeddings,𝐊,𝐐,𝐕∈ℝT×d\mathbf{K},\mathbf{Q},\mathbf{V}\in\mathbb{R}^{T\times d}and

MHA(𝐗)=[Softmax(𝐐i𝐊i⊤/d)𝐕i)]i𝐖O,\mathrm{MHA}(\mathbf{X})=\big[\mathrm{Softmax}(\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}/\sqrt{d})\mathbf{V}_{i})\big]_{i}\mathbf{W}^{O},MHA~(𝐗)=[Softmax(𝐐i𝐊i~⊤/d)𝐕i~)]i𝐖O,\widetilde{\mathrm{MHA}}(\mathbf{X})=\big[\mathrm{Softmax}(\mathbf{Q}_{i}\widetilde{\mathbf{K}_{i}}^{\top}/\sqrt{d})\widetilde{\mathbf{V}_{i}})\big]_{i}\mathbf{W}^{O},whereMHA~​(𝐗),𝐊i~\widetilde{\mathrm{MHA}}(\mathbf{X}),\widetilde{\mathbf{K}_{i}}and𝐕i~\widetilde{\mathbf{V}_{i}}represent the approximation ofMHA​(𝐗),𝐊i\mathrm{MHA}(\mathbf{X}),\mathbf{K}_{i}and𝐕i\mathbf{V}_{i}, respectively. The difference between the actual attention output and the one produced with approximate keys and values is upper bounded as

‖MHA~​(𝐗)−MHA​(𝐗)‖2\displaystyle\|\widetilde{\mathrm{MHA}}(\mathbf{X})-\mathrm{MHA}(\mathbf{X})\|_{2}≤∑i=1h‖𝐕i​𝐖iO‖2d​‖𝐐i​𝐊i⊤−𝐐i​𝐊i~⊤‖2\displaystyle\leq\sum_{i=1}^{h}\frac{\|\mathbf{V}_{i}\mathbf{W}_{i}^{O}\|_{2}}{\sqrt{d}}\|\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}-\mathbf{Q}_{i}\widetilde{\mathbf{K}_{i}}^{\top}\|_{2}+‖𝐕i​𝐖iO−𝐕i~​𝐖iO‖2.\displaystyle+\|\mathbf{V}_{i}\mathbf{W}_{i}^{O}-\widetilde{\mathbf{V}_{i}}\mathbf{W}_{i}^{O}\|_{2}.(1)

The method proposed by this paper stems from the will to minimize this upper bound on the output approximation error. In each attention head, for a key cache𝐊\mathbf{K}and a calibration set of query vectors{𝐪(j)}j∈calibration\{\mathbf{q}^{(j)}\}_{j\in\mathrm{calibration}}(row vectors of sizedd), we want to minimize

∑j∈calibration‖𝐪(j)​𝐊⊤−𝐪(j)​𝐊~⊤‖22,\sum_{j\in\mathrm{calibration}}\|\mathbf{q}^{(j)}\mathbf{K}^{\top}-\mathbf{q}^{(j)}\widetilde{\mathbf{K}}^{\top}\|_{2}^{2},over low-rank𝐊~\widetilde{\mathbf{K}}. This objective is equivalent to a low-rank approximation of the query–key interaction matrix and admits an optimal closed-form solution via the singular value decomposition. Similarly, to compress values we directly optimize the second term of Eq. (Theorem) over low-rank𝐕~\widetilde{\mathbf{V}}for each attention head (see AppendixB).

4.2Proposed Method

In the following, we drop the head indexiifor ease of clarity. Our method addresses cache compression by producing an optimal solution to

min𝐊~⁡‖𝐐𝐊⊤−𝐐​𝐊~⊤‖F2​s.t.​rank​(𝐊~)≤R,\min_{\widetilde{\mathbf{K}}}\|\mathbf{Q}\mathbf{K}^{\top}-\mathbf{Q}\widetilde{\mathbf{K}}^{\top}\|_{F}^{2}\ \text{s.t.}\ \mathrm{rank}(\widetilde{\mathbf{K}})\leq R,with𝐐=(𝐪(j))j∈calibration\mathbf{Q}=\begin{pmatrix}\mathbf{q}^{(j)}\end{pmatrix}_{j\in\text{calibration}}the collection of query vectors from the calibration set. This formulation jointly considers the properties of both𝐊\mathbf{K}and𝐐\mathbf{Q}, while taking into account their interaction through the inner product. The optimization problem can be restated using a projection matrix as

min𝐒⁡‖𝐊𝐒𝐐⊤−𝐊𝐐⊤‖F2s.t.​rank​(𝐒)≤R,\min_{\mathbf{S}}\|\mathbf{K}\mathbf{S}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2}~~~~~~\text{s.t.}~~~\mathrm{rank}(\mathbf{S})\leq R,where𝐒∈ℝd×d\mathbf{S}\in\mathbb{R}^{d\times d}. To be able to compress the key cache during the attention computation, we write the projection matrix𝐒\mathbf{S}as a product of two matrices, i.e.,𝐒=𝐀𝐁\mathbf{S}=\mathbf{A}\mathbf{B}with𝐀,𝐁∈ℝd×R\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}. The minimization problem tackled by KQ-SVD is thus

min𝐀,𝐁∈ℝd×R⁡‖𝐊𝐀𝐁⊤​𝐐⊤−𝐊𝐐⊤‖F2.\displaystyle\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2}.(2)Given𝐀∗\mathbf{A}^{*}and𝐁∗\mathbf{B}^{*}the solutions of this optimization problem, we store the low-rank matrix𝐊𝐀∗∈ℝT×R\mathbf{K}\mathbf{A}^{*}\in\mathbb{R}^{T\times R}, which allows for compression while maintaining an accurate reconstruction of the attention matrix. The same strategy also applies to the value–output matrices (see AppendixB). In the next section, we state the main theorem which demonstrates that the optimization problem described above admits an optimal solution which can be computed efficiently.

4.3KQ-SVD: Optimal attention factorization

The following theorem establishes the provably optimal low-rank factorization of the key–query matrix, which admits a simple closed-form solution.

Theorem 2.

Let𝐊,𝐐∈ℝT×d\mathbf{K},\mathbf{Q}\in\mathbb{R}^{T\times d}be key and query cache matrices. The optimal solution to the low rank attention approximation problem

min𝐀,𝐁∈ℝd×R⁡‖𝐊𝐀𝐁⊤​𝐐⊤−𝐊𝐐⊤‖F,\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|_{F},is given by

𝐀=𝐊+​𝐔^,𝐁=𝐊T​𝐔^,\mathbf{A}=\mathbf{K}^{+}\hat{\mathbf{U}},\ \ \mathbf{B}=\mathbf{K}^{T}\hat{\mathbf{U}},where𝐔^∈ℝT×R\hat{\mathbf{U}}\in\mathbb{R}^{T\times R}is the matrix having the topRRleft singular vectors of𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}as columns.

Proof.

Observe that𝐊𝐀𝐁⊤​𝐐⊤\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}has rank at mostRR. Hence,

min𝐀,𝐁∈ℝd×R⁡‖𝐊𝐀𝐁⊤​𝐐⊤−𝐊𝐐⊤‖F,\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|_{F},is lower bounded by

min𝐌∈ℝT×T⁡‖𝐌−𝐊𝐐⊤‖F​s.t.​rank​(𝐌)≤R.\min_{\mathbf{M}\in\mathbb{R}^{T\times T}}\|\mathbf{M}-\mathbf{K}\mathbf{Q}^{\top}\|_{F}\ \text{s.t.}\ \mathrm{rank}(\mathbf{M})\leq R.By the Eckart-Young theorem, we know the best rankRRapproximation of𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}is given by its truncated SVD:

𝐊𝐐⊤≃𝐔^​𝚺^​𝐕^⊤,\mathbf{K}\mathbf{Q}^{\top}\simeq\hat{\mathbf{U}}\hat{\mathbf{\Sigma}}\hat{\mathbf{V}}^{\top},where𝐔^∈ℝT×R,𝚺^∈ℝR×R\hat{\mathbf{U}}\in\mathbb{R}^{T\times R},\hat{\mathbf{\Sigma}}\in\mathbb{R}^{R\times R}and𝐕^∈ℝT×R\hat{\mathbf{V}}\in\mathbb{R}^{T\times R}. We will show that𝐊𝐀𝐁⊤​𝐐⊤=𝐔^​𝚺^​𝐕^⊤\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}=\hat{\mathbf{U}}\hat{\mathbf{\Sigma}}\hat{\mathbf{V}}^{\top}from which the optimality of𝐀\mathbf{A}and𝐁\mathbf{B}follows.

Let𝐊=𝐔​𝚺​𝐕\mathbf{K}=\mathbf{U}\mathbf{\Sigma}\mathbf{V}be the full SVD of𝐊\mathbf{K}. Observe that we have the inclusion of (column) spans:ℛ​(𝐔^)⊆ℛ​(𝐔)=ℛ​(𝐊)\mathcal{R}(\hat{\mathbf{U}})\subseteq\mathcal{R}(\mathbf{U})=\mathcal{R}(\mathbf{K}). Since𝐊𝐊+\mathbf{K}\mathbf{K}^{+}is the orthogonal projection ontoℛ​(𝐊)\mathcal{R}(\mathbf{K}), we have𝐊𝐊+​𝐔^=𝐔^\mathbf{K}\mathbf{K}^{+}\hat{\mathbf{U}}=\hat{\mathbf{U}}, hence

𝐊𝐀𝐁⊤​𝐐⊤\displaystyle\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}=𝐊𝐊+​𝐔^​𝐔^⊤​𝐊𝐐⊤=𝐔^​𝐔^⊤​𝐊𝐐⊤\displaystyle=\mathbf{K}\mathbf{K}^{+}\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}\mathbf{K}\mathbf{Q}^{\top}=\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}\mathbf{K}\mathbf{Q}^{\top}=𝐔^​𝐔^⊤​𝐔​𝚺​𝐕⊤=𝐔^​𝚺^​𝐕^⊤.\displaystyle=\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}\mathbf{U}\mathbf{\Sigma}\mathbf{V}^{\top}=\hat{\mathbf{U}}\hat{\mathbf{\Sigma}}\hat{\mathbf{V}}^{\top}.Therefore𝐀=𝐊+​𝐔^\mathbf{A}=\mathbf{K}^{+}\hat{\mathbf{U}}and𝐁=𝐊⊤​𝐔^\mathbf{B}=\mathbf{K}^{\top}\hat{\mathbf{U}}are optimal solutions tomin𝐀,𝐁∈ℝd×R⁡‖𝐊𝐀𝐁⊤​𝐐⊤−𝐊𝐐⊤‖F\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|_{F}. ∎

The Moore-Penrose pseudo inverse of𝐊\mathbf{K}can be expressed through the SVD of𝐊\mathbf{K}as𝐊+=𝐕K​𝚺K−1​𝐔K⊤\mathbf{K}^{+}=\mathbf{V}_{K}\mathbf{\Sigma}^{-1}_{K}\mathbf{U}_{K}^{\top}. The singular value decomposition of𝐊𝐐⊤∈ℝT×T\mathbf{K}\mathbf{Q}^{\top}\in\mathbb{R}^{T\times T}can be computed efficiently as its rank is at mostdd. Indeed, we can first perform an SVD on𝐊=𝐔K​𝚺K​𝐕K⊤\mathbf{K}=\mathbf{U}_{K}\mathbf{\Sigma}_{K}\mathbf{V}_{K}^{\top}and𝐐=𝐔Q​𝚺Q​𝐕Q⊤\mathbf{Q}=\mathbf{U}_{Q}\mathbf{\Sigma}_{Q}\mathbf{V}_{Q}^{\top}, then on thed×dd\times dmatrix𝚺K​𝐕K⊤​𝐕Q​𝚺Q=𝐔′​𝚺′​𝐕′⁣⊤\mathbf{\Sigma}_{K}\mathbf{V}_{K}^{\top}\mathbf{V}_{Q}\mathbf{\Sigma}_{Q}=\mathbf{U}^{\prime}\mathbf{\Sigma}^{\prime}\mathbf{V}^{\prime\top}. The SVD of𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}is then𝐔K​𝐔′​𝚺′​(𝐕Q​𝐕′)⊤=𝐔​𝚺​𝐕⊤\mathbf{U}_{K}\mathbf{U}^{\prime}\mathbf{\Sigma}^{\prime}(\mathbf{V}_{Q}\mathbf{V}^{\prime})^{\top}=\mathbf{U}\mathbf{\Sigma}\mathbf{V}^{\top}with𝐔=𝐔K​𝐔′∈ℝT×d\mathbf{U}=\mathbf{U}_{K}\mathbf{U}^{\prime}\in\mathbb{R}^{T\times d}and𝐕=𝐕Q​𝐕′∈ℝT×d\mathbf{V}=\mathbf{V}_{Q}\mathbf{V}^{\prime}\in\mathbb{R}^{T\times d}. This way the optimal solution provided by Theorem2can be computed efficiently in time𝒪​(T​d2)\mathcal{O}(Td^{2}).

A similar approach is used by KQ-SVD to derive optimal projections for the low rank approximation of the product of the value cache with the output matrix (see AppendixB).

5Theoretical analysis

In this section, we focus on the minimization problem in Eq. (5); the same reasoning applies to value and output matrices (see AppendixB). We first provide an exact formula quantifying the accuracy difference between K-SVD and KQ-SVD. We then exhibit a failure mode of Eigen method[saxena2024eigen]which KQ-SVD circumvent by design. Finally, we show how KQ-SVD also obtains optimal low rank approximation in the Grouped Query Attention (GQA) framework.

5.1Comparing K-SVD and KQ-SVD

In this section, we characterize the optimality gap between KQ-SVD and K-SVD (Section3.3) and derive a closed form expression of their accuracy difference.

Theorem 3.

Letopt=min𝐀,𝐁∈ℝd×R⁡‖𝐊𝐀𝐁⊤​𝐐⊤−𝐊𝐐⊤‖F2\mathrm{opt}=\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2}be the optimal error for the low rank attention approximation problem (achieved by KQ-SVD), and leterrK-SVD=‖𝐊​𝐕^K​𝐕^KT​𝐐T−𝐊𝐐T‖F2\mathrm{err}_{\text{K-SVD}}=\|\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{T}\mathbf{Q}^{T}-\mathbf{K}\mathbf{Q}^{T}\|_{F}^{2}be the error of K-SVD. We have

errK-SVD−opt=∑i=1Rσi​(𝐊𝐐⊤)2−‖𝐊​𝐕^K​𝐕^KT​𝐐T‖F2≥0,\mathrm{err}_{\text{K-SVD}}-\mathrm{opt}=\sum_{i=1}^{R}\sigma_{i}(\mathbf{K}\mathbf{Q}^{\top})^{2}-\|\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{T}\mathbf{Q}^{T}\|^{2}_{F}\geq 0,with equality only if the topRRleft singular vectors of𝐊\mathbf{K}and the topRRof𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}span the same subspace.

Proof.

Let𝐊≈𝐔^K​𝚺^K​𝐕^K⊤\mathbf{K}\approx\hat{\mathbf{U}}_{K}\hat{\mathbf{\Sigma}}_{K}\hat{\mathbf{V}}_{K}^{\top}be the rankRRtruncated SVD of𝐊\mathbf{K}and𝐊𝐐⊤≈𝐔^​𝚺^​𝐕^⊤\mathbf{K}\mathbf{Q}^{\top}\approx\hat{\mathbf{U}}\hat{\mathbf{\Sigma}}\hat{\mathbf{V}}^{\top}be the one of𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}.

By Theorem2,opt=‖𝐔^​𝐔^⊤​𝐊𝐐T−𝐊𝐐T‖F2=∑i=R+1dσi​(𝐊𝐐T)2\mathrm{opt}=\|\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}\mathbf{K}\mathbf{Q}^{T}-\mathbf{K}\mathbf{Q}^{T}\|_{F}^{2}=\sum_{i=R+1}^{d}\sigma_{i}(\mathbf{K}\mathbf{Q}^{T})^{2}.

On the one hand we have

‖𝐊𝐐T‖F2=∑i=1dσi​(𝐊𝐐T)2=∑i=1Rσi​(𝐊𝐐T)2+opt.\displaystyle\|\mathbf{K}\mathbf{Q}^{T}\|_{F}^{2}=\sum_{i=1}^{d}\sigma_{i}(\mathbf{K}\mathbf{Q}^{T})^{2}=\sum_{i=1}^{R}\sigma_{i}(\mathbf{K}\mathbf{Q}^{T})^{2}+\mathrm{opt}. On the other hand,

‖𝐊𝐐⊤‖F2=‖𝐊𝐐⊤−𝐔^K​𝐔^K⊤​𝐊𝐐⊤+𝐔^K​𝐔^K⊤​𝐊𝐐⊤‖F2\displaystyle\|\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2}=\|\mathbf{K}\mathbf{Q}^{\top}-\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}\mathbf{Q}^{\top}+\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2}=‖(𝐈−𝐔^K​𝐔^K⊤)​𝐊𝐐⊤‖F2+‖𝐔^K​𝐔^K⊤​𝐊𝐐⊤‖F2,\displaystyle=\|(\mathbf{I}-\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top})\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2}+\|\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}\mathbf{Q}^{\top}\|_{F}^{2},since(𝐈−𝐔^K​𝐔^K⊤)(\mathbf{I}-\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top})and𝐔^K​𝐔^K⊤\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}are projections on orthogonal subspaces.

Since𝐊​𝐕^K​𝐕^K⊤=𝐔^K​𝐔^K⊤​𝐊\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}=\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}, the left term is exactlyerrK-SVD\mathrm{err}_{\text{K-SVD}}and the second term can be reduced to𝐔^K​𝐔^K⊤​𝐊𝐐⊤=𝐊​𝐕^K​𝐕^K⊤​𝐐⊤\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}\mathbf{Q}^{\top}=\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}\mathbf{Q}^{\top}.

Putting everything together, we get

∑i=1Rσi​(𝐊𝐐T)2+opt=errK-SVD+‖𝐊​𝐕^K​𝐕^K⊤​𝐐⊤‖F2,\sum_{i=1}^{R}\sigma_{i}(\mathbf{K}\mathbf{Q}^{T})^{2}+\mathrm{opt}=\mathrm{err}_{\text{K-SVD}}+\|\mathbf{K}\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}\mathbf{Q}^{\top}\|_{F}^{2},which shows the equality in the theorem.

To show thaterrK-SVD−opt≥0\mathrm{err}_{\text{K-SVD}}-\mathrm{opt}\geq 0, first observe that

errK-SVD−opt=‖𝐔^​𝐔^⊤​𝐊𝐐⊤‖−‖𝐔^K​𝐔^K⊤​𝐊𝐐⊤‖.\mathrm{err}_{\text{K-SVD}}-\mathrm{opt}=\|\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}\mathbf{K}\mathbf{Q}^{\top}\|-\|\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}\mathbf{Q}^{\top}\|.A direct consequence of the Eckart-Young theorem is that, for any matrix𝐌∈ℝT×T\mathbf{M}\in\mathbb{R}^{T\times T}and anyR≤TR\leq T, the solution of

max𝐗∈ℝT×R⁡‖𝐗𝐗⊤​𝐌‖​s.t.𝐗⊤​𝐗=𝐈\max_{\mathbf{X}\in\mathbb{R}^{T\times R}}\|\mathbf{X}\mathbf{X}^{\top}\mathbf{M}\|~~~\text{s.t.}~~\ \mathbf{X}^{\top}\mathbf{X}=\mathbf{I}is obtained by setting the columns of𝐗\mathbf{X}to the topRRleft singular vectors of𝐌\mathbf{M}. Hence,‖𝐔^​𝐔^⊤​𝐊𝐐⊤‖≥‖𝐔^K​𝐔^K⊤​𝐊𝐐⊤‖\|\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}\mathbf{K}\mathbf{Q}^{\top}\|\geq\|\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}\mathbf{K}\mathbf{Q}^{\top}\|and thuserrK-SVD−opt≥0\mathrm{err}_{\text{K-SVD}}-\mathrm{opt}\geq 0, with equality only if𝐔^​𝐔^⊤=𝐔^K​𝐔^K⊤\hat{\mathbf{U}}\hat{\mathbf{U}}^{\top}=\hat{\mathbf{U}}_{K}\hat{\mathbf{U}}_{K}^{\top}, i.e., when the topRRleft singular vectors of𝐊\mathbf{K}and the topRRof𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}span the same subspace. ∎

It is worth observing that equality betweenerrK-SVD\mathrm{err}_{\text{K-SVD}}andopt\mathrm{opt}happensonly whenthe projection onto𝐕^K\hat{\mathbf{V}}_{K}capturesallof the energy (Frobenius norm) in the topRRsingular values of𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}. Since the best rank-RRapproximation in the Frobenius norm is unique, this holds precisely when the subspace spanned by the topRRleft singular vectors of𝐊\mathbf{K}coincides with that spanned by the topRRleft singular vectors of𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}. In other words, equality holdsprecisely when these two subspaces match.

5.2Comparing Eigen and KQ-SVD

We now compare Eigen[saxena2024eigen]with KQ-SVD. Although we do not derive an exact value for the optimality gap of Eigen, we identify a critical limitation of Eigen: the method is highly sensitive to unbalance between the norms of𝐊\mathbf{K}and𝐐\mathbf{Q}. Eigen’s performance can be degraded simply by multiplying𝐊\mathbf{K}by a constantβ\betaand dividing𝐐\mathbf{Q}by the same constant. While this rescaling leaves the attention computation unchanged and does not affect KQ-SVD, it causes Eigen method to behave increasingly like K-SVD as the unbalance between the two norms grows. Theorem4formalizes this intuition.

Theorem 4.

Let𝐊,𝐐∈ℝT×d\mathbf{K},\mathbf{Q}\in\mathbb{R}^{T\times d}, letα=‖𝐐‖F‖𝐊‖F\alpha=\frac{\|\mathbf{Q}\|_{F}}{\|\mathbf{K}\|_{F}}and leterrEigen=‖𝐊​𝐕^Eigen​𝐕^Eigen⊤​𝐐⊤−𝐊𝐐⊤‖\mathrm{err}_{\text{Eigen}}=\|\mathbf{K}\hat{\mathbf{V}}_{\text{Eigen}}\hat{\mathbf{V}}_{\text{Eigen}}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|be the error of Eigen. If there is a non-trivial gap between theRRth and(R+1)(R+1)th singular values of𝐊\mathbf{K}, i.e.σR​(𝐊)>σR+1​(𝐊)\sigma_{R}(\mathbf{K})>\sigma_{R+1}(\mathbf{K}), thenlimα→0errEigen=errK-SVD\lim_{\alpha\to 0}\mathrm{err}_{\text{Eigen}}=\mathrm{err}_{\text{K-SVD}}.

Proof.

Recall that Eigen approximate𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}with𝐊​𝐕^E​i​g​𝐕^E​i​g⊤​𝐐⊤\mathbf{K}\hat{\mathbf{V}}_{Eig}\hat{\mathbf{V}}_{Eig}^{\top}\mathbf{Q}^{\top}where𝐕^E​i​g\hat{\mathbf{V}}_{Eig}is the matrix with the topRRright singular vectors of[𝐊𝐐]\begin{bmatrix}\mathbf{K}\\ \mathbf{Q}\end{bmatrix}as columns. In the limit where‖𝐐‖F‖𝐊‖F\frac{\|\mathbf{Q}\|_{F}}{\|\mathbf{K}\|_{F}}tends to0, the concatenated matrix[𝐊𝐐]\begin{bmatrix}\mathbf{K}\\ \mathbf{Q}\end{bmatrix}tends to[𝐊𝟎]\begin{bmatrix}\mathbf{K}\\ \mathbf{0}\end{bmatrix}. SinceσR​(𝐊)>σR+1​(𝐊)\sigma_{R}(\mathbf{K})>\sigma_{R+1}(\mathbf{K}), it follows from the Davis-Kahan theorem[davis1970rotation]that, asα\alphatends to0, the space spanned by the topRRright singular vectors of[𝐊𝐐]\begin{bmatrix}\mathbf{K}\\ \mathbf{Q}\end{bmatrix}converges to the one of[𝐊𝟎]\begin{bmatrix}\mathbf{K}\\ \mathbf{0}\end{bmatrix}, and thus of𝐊\mathbf{K}. Hencelimα→0𝐕^E​i​g​𝐕^E​i​g⊤=𝐕^K​𝐕^K⊤\lim_{\alpha\to 0}\hat{\mathbf{V}}_{Eig}\hat{\mathbf{V}}_{Eig}^{\top}=\hat{\mathbf{V}}_{K}\hat{\mathbf{V}}_{K}^{\top}anderrEigen\mathrm{err}_{\text{Eigen}}tends toerrK-SVD\mathrm{err}_{\text{K-SVD}}. ∎

5.3Handling grouped query attention

Standard multi-head attention (MHA) is powerful but slow at inference, while multi-query attention (MQA)[shazeer2019fast]is much faster but can hurt model quality and requires retraining. Grouped-query attention (GQA)[ainslie2023GQA], sits in between MHA and MQA by letting groups of query heads share a key-value head, balancing accuracy and performance. GQA organizes attention/query heads into groups of sizemm. All query heads within a group share the same set of keys and values, enabling more efficient computation without significantly compromising model performance. We have assumed in the presentation of our method that each key head attends to a single query head. We show in the following theorem that simply applying KQ-SVD on the shared key cache and the concatenated query caches provides the optimal approximation with GQA.

Theorem 5.

Given a (shared) key cache matrix𝐊∈ℝT×d\mathbf{K}\in\mathbb{R}^{T\times d}andmmfull column rank query cache matrices𝐐1,⋯,𝐐m∈ℝT×d\mathbf{Q}_{1},\cdots,\mathbf{Q}_{m}\in\mathbb{R}^{T\times d}, solving

min𝐀,𝐁1,⋯,𝐁m∈ℝd×R​∑i=1m‖𝐊𝐀𝐁i⊤​𝐐i⊤−𝐊𝐐i⊤‖F2,\min_{\mathbf{A},\mathbf{B}_{1},\cdots,\mathbf{B}_{m}\in\mathbb{R}^{d\times R}}\sum_{i=1}^{m}\|\mathbf{K}\mathbf{A}\mathbf{B}_{i}^{\top}\mathbf{Q}_{i}^{\top}-\mathbf{K}\mathbf{Q}_{i}^{\top}\|^{2}_{F},is equivalent to solving

min𝐀,𝐁∈ℝd×R⁡‖𝐊𝐀𝐁⊤​𝐐⊤−𝐊𝐐⊤‖F2,\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}^{\top}-\mathbf{K}\mathbf{Q}^{\top}\|^{2}_{F},where𝐐=[𝐐1⊤​𝐐2⊤​⋯​𝐐m⊤]⊤∈ℝm​T×d\mathbf{Q}=[\mathbf{Q}_{1}^{\top}\ \mathbf{Q}_{2}^{\top}\ \cdots\ \mathbf{Q}_{m}^{\top}]^{\top}\in\mathbb{R}^{mT\times d}is the matrix obtained by stacking themmquery matrices.

Proof.

We first show that the solution matrices𝐁i\mathbf{B}_{i}can be chosen to be all equal, i.e., that solving

min𝐀,𝐁1,⋯,𝐁m∈ℝd×R​∑i=1m‖𝐊𝐀𝐁i⊤​𝐐i⊤−𝐊𝐐i⊤‖F2,\min_{\mathbf{A},\mathbf{B}_{1},\cdots,\mathbf{B}_{m}\in\mathbb{R}^{d\times R}}\sum_{i=1}^{m}\|\mathbf{K}\mathbf{A}\mathbf{B}_{i}^{\top}\mathbf{Q}_{i}^{\top}-\mathbf{K}\mathbf{Q}_{i}^{\top}\|^{2}_{F},is equivalent to solving

min𝐀,𝐁∈ℝd×R​∑i=1m‖𝐊𝐀𝐁⊤​𝐐i⊤−𝐊𝐐i⊤‖F2.\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\sum_{i=1}^{m}\|\mathbf{K}\mathbf{A}\mathbf{B}^{\top}\mathbf{Q}_{i}^{\top}-\mathbf{K}\mathbf{Q}_{i}^{\top}\|^{2}_{F}.Indeed, for any𝐀\mathbf{A}and anyii, since𝐐i⊤\mathbf{Q}_{i}^{\top}is full row rank, the minimizers ofmin𝐁i⁡‖𝐊𝐀𝐁i⊤​𝐐i⊤−𝐊𝐐i⊤‖F2\min_{\mathbf{B}_{i}}\|\mathbf{K}\mathbf{A}\mathbf{B}_{i}^{\top}\mathbf{Q}_{i}^{\top}-\mathbf{K}\mathbf{Q}_{i}^{\top}\|^{2}_{F}are the same as the minimizers ofmin𝐁i⁡‖𝐊𝐀𝐁i⊤−𝐊‖F2\min_{\mathbf{B}_{i}}\|\mathbf{K}\mathbf{A}\mathbf{B}_{i}^{\top}-\mathbf{K}\|^{2}_{F}, and are thus independent of𝐐i\mathbf{Q}_{i}. Since this is true for any𝐀\mathbf{A}, and in particular for the optimal one, this shows that all the solutions𝐁i\mathbf{B}_{i}can be taken to be equal.

The result then directly follows from the fact that the squared Frobenius norm of a block matrix is equal to the sum of the squared Frobenius norms of the blocks. ∎

Theorem5states that optimal projections can be computed for models using GQA by stacking query matrices in each group and using KQ-SVD as in the non GQA case. Computing the optimal projection for a head group costs𝒪​(m​T​d2)\mathcal{O}(mTd^{2}), wheremmis the size of the head group, leading to an amortized cost per query head of𝒪​(T​d2)\mathcal{O}(Td^{2}).

6Experiments

Figure 1:Relative Frobenius approximation output error per layer (top) and mean relative errors on𝐊\mathbf{K},𝐐\mathbf{Q},𝐕\mathbf{V},𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}and attention layer output across layers (bottom) for Mistral and Llama models.The theoretical results established in the previous sections require further validation on real cache matrices generated by state-of-the-art LLMs. In this section, we first compare KQ-SVD with K-SVD and Eigen, and then demonstrate the claim of Theorem4using real-world cache data.

6.1Comparing methods

Setup:To validate our theoretical claims, we run experiments on several widely used LLMs. We test two models without grouped-query attention (GQA)—Llama2-7B and Llama2-13B[touvron2023llama]—and two models with GQA—Llama3-8B[grattafiori2024llama]and Mistral-7B-v0.3[jiang2023mistral7b]. All experiments use the C4 dataset[C4], with projections learned on the training split and evaluated on the validation split.

Learning projections:Following the methodology ofsaxena2024eigen, we select 128 training sequences of 2048 tokens each from C4. Each sequence is passed through the model, storing the key, value, and query caches for every layer and attention head (queries are needed by both Eigen and KQ-SVD). For each layer–head pair, we collect the corresponding caches from all 128 sequences and concatenate them, yielding large matrices𝐊,𝐐,𝐕∈ℝThuge×d\mathbf{K},\mathbf{Q},\mathbf{V}\in\mathbb{R}^{T_{\text{huge}}\times d}withThuge=262,144T_{\text{huge}}=262{,}144. Because model context length is limited and attention cost scales quadratically with sequence length, it is not feasible to build these matrices from a single long sequence. Instead, we construct them from multiple shorter ones. Once the large cache matrices are formed, we perform SVD and apply the formulas presented in previous sections to compute low-rank projections.

Rank selection:All methods are evaluated at the same rankRR, determined individually for each layer. For a given layer, we analyze the singular value spectra of the key and value matrices, averaged across heads, and choose the smallestRRthat discards no more than anϵ=0.1\epsilon=0.1fraction of the spectral energy. That is, for singular values{σj}j\{\sigma_{j}\}_{j}of𝐊\mathbf{K},∑j=1Rσj2∑i=1dσj2≥1−ϵ\frac{\sum_{j=1}^{R}\sigma_{j}^{2}}{\sum_{i=1}^{d}\sigma_{j}^{2}}\geq 1-\epsilon. This is equivalent to requiring that the relative Frobenius error is at mostϵ\epsilon.

Evaluation:We evaluate the learned projections on 32 validation sequences of 2048 tokens each. For every sequence, we extract the cache matrices (𝐊,𝐐,𝐕\mathbf{K},\mathbf{Q},\mathbf{V}) at each layer and head. Using these matrices, we can simulate attention computations directly, since attention depends only on these three components. Each cache is projected onto its corresponding low-rank subspace to form approximations, and we then compute the approximate Multi-Head Attention output using the standard formulas. For comparison, we also compute the exact (uncompressed) attention output. This enables us to measure the relative error of each method across all matrices of the attention pipeline. Errors are averaged across validation sequences.

Metrics:For a matrix𝐌\mathbf{M}and its approximation𝐌~\widetilde{\mathbf{M}}, we report the relative Frobenius norm errorerrFro=‖𝐌−𝐌~‖F2‖𝐌‖F2\text{err}_{\text{Fro}}=\frac{\|\mathbf{M}-\widetilde{\mathbf{M}}\|_{F}^{2}}{\|\mathbf{M}\|_{F}^{2}}. This error is computed for the key, query, and value matrices, for the attention score matrix𝐊𝐐⊤\mathbf{K}\mathbf{Q}^{\top}, and for the Multi-Head Attention outputMHA​(𝐗)\mathrm{MHA}(\mathbf{X}).

Results:Results are shown in Figure1. For each model, the top plot reports the relative error on the attention output across layers, while the bottom plot reports the averaged errors on the intermediate components. We observe that K-SVD provides the most accurate approximation of the key matrices (as expected from the optimality of SVD), but performs poorly on query matrices, leading to weaker approximations of the attention scores and consequently higher output errors. This effect is more pronounced in GQA models, where sharing the key matrix across queries in a group amplifies approximation errors.

In contrast, Eigen and KQ-SVD achieve comparable accuracy on keys, queries, and values. The key difference lies in the attention score matrix: KQ-SVD consistently delivers higher accuracy, resulting in lower attention errors and outperforming all other methods on all models.

6.2Unbalanced𝐊\mathbf{K}and𝐐\mathbf{Q}matrices

We verify experimentally the claims of Section5.2and analyze the attention approximation error under unbalanced𝐊\mathbf{K}and𝐐\mathbf{Q}matrices.

Set up: We follow the same experimental setup as in the previous section using the C4 train/validation split to learn projections and evaluate their accuracy. The only difference is that cache matrices are scaled to assess the effect of unbalance. Key matrices are multiplied byβ\betaand query matrices are divided byβ\beta. This is equivalent to scaling the projection matrices𝐖K,i\mathbf{W}_{K,i}/𝐖Q,i\mathbf{W}_{Q,i}byβ\beta, as the operation permutes with matrix multiplication (this does not change the output since the query and key matrices are multiplied before any non-linear activation).

Metrics: For each unbalance ratio, we plot the relative attention output error for the three methods, averaged across all layers and validation sequences.

Results: Results are shown in Figure2. As discussed in Section5.2, K-SVD and KQ-SVD are invariant to scaling𝐊\mathbf{K}by a factorβ\betaand dividing𝐐\mathbf{Q}by the same factor, which is confirmed by the constant error observed for these methods. As predicted theoretically, a higher unbalance ratio brings Eigen closer to K-SVD; forβ=10\beta=10, their errors are nearly indistinguishable. This confirms Theorem4and exposes a limitation of Eigen, which underperforms even under modest query-key unbalance.

Figure 2:Llama2-7B relative output approximation error averaged across layers for varying unbalanced factorβ\beta.

7Conclusion

We introduced KQ-SVD, a novel approach for computing low-rank projections for KV cache compression. KQ-SVD is driven by minimizing an upper bound on the attention output approximation error, resulting in an optimization problem that yields the optimal low-rank approximation of the attention matrix. Crucially, this problem admits a closed-form solution that can be computed efficiently. We quantify the advantage of KQ-SVD over previous methods, either through exact error formulas or by highlighting failure modes of prior approaches. Our technique is complementary to popular cache compression methods such as GQA. Experiments validate our theoretical findings and demonstrate that considering the interaction between queries and keys—as KQ-SVD—provides a superior alternative to performing SVD on the key cache or on concatenated keys and queries.

Acknowledgment

This research is supported by the Canadian Institute for Advanced Research (CIFAR AI chair program). This work was completed while Damien Lesens interned at Mila. This work made use of compute resources provided by the Digital Research Alliance of Canada and by Mila (mila.quebec).

KQ-SVD: Compressing the KV Cache with Provable Guarantees on Attention Fidelity (Supplementary Material)

Appendix AProof of Theorem1

Theorem.

Let𝐗∈ℝT×d\mathbf{X}\in\mathbb{R}^{T\times d}be a sequence of token embeddings,𝐊,𝐐,𝐕∈ℝT×d\mathbf{K},\mathbf{Q},\mathbf{V}\in\mathbb{R}^{T\times d}and

MHA(𝐗)=[Softmax(𝐐i𝐊i⊤/d)𝐕i)]i𝐖O,\mathrm{MHA}(\mathbf{X})=\big[\mathrm{Softmax}(\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}/\sqrt{d})\mathbf{V}_{i})\big]_{i}\mathbf{W}^{O},MHA~(𝐗)=[Softmax(𝐐i𝐊i~⊤/d)𝐕i~)]i𝐖O,\widetilde{\mathrm{MHA}}(\mathbf{X})=\big[\mathrm{Softmax}(\mathbf{Q}_{i}\widetilde{\mathbf{K}_{i}}^{\top}/\sqrt{d})\widetilde{\mathbf{V}_{i}})\big]_{i}\mathbf{W}^{O},whereMHA~​(𝐗),𝐊i~\widetilde{\mathrm{MHA}}(\mathbf{X}),\widetilde{\mathbf{K}_{i}}and𝐕i~\widetilde{\mathbf{V}_{i}}represent the approximation ofMHA​(𝐗),𝐊i\mathrm{MHA}(\mathbf{X}),\mathbf{K}_{i}and𝐕i\mathbf{V}_{i}, respectively. The difference between the actual attention output and the one produced with approximate keys and values is upper bounded as

‖MHA~​(𝐗)−MHA​(𝐗)‖2\displaystyle\|\widetilde{\mathrm{MHA}}(\mathbf{X})-\mathrm{MHA}(\mathbf{X})\|_{2}≤∑i=1h‖𝐕i​𝐖iO‖2d​‖𝐐i​𝐊i⊤−𝐐i​𝐊i~⊤‖2\displaystyle\leq\sum_{i=1}^{h}\frac{\|\mathbf{V}_{i}\mathbf{W}_{i}^{O}\|_{2}}{\sqrt{d}}\|\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}-\mathbf{Q}_{i}\widetilde{\mathbf{K}_{i}}^{\top}\|_{2}+‖𝐕i​𝐖iO−𝐕i~​𝐖iO‖2.\displaystyle+\|\mathbf{V}_{i}\mathbf{W}_{i}^{O}-\widetilde{\mathbf{V}_{i}}\mathbf{W}_{i}^{O}\|_{2}.(3)

Proof.

Let𝐖O=[𝐖1O;𝐖2O;⋯;𝐖hO]∈ℝD×D\mathbf{W}^{O}=[\mathbf{W}^{O}_{1};\mathbf{W}^{O}_{2};\cdots;\mathbf{W}^{O}_{h}]\in\mathbb{R}^{D\times D}, where𝐖iO∈ℝd×D\mathbf{W}^{O}_{i}\in\mathbb{R}^{d\times D}are stacked vertically. By the definition of multi-head attention (see Section3.2), we can write

MHA​(𝐗)=[𝐇1,…,𝐇h]​𝐖O=∑i=1h𝐇i​𝐖iO,MHA~​(𝐗)=[𝐇~1,…,𝐇~h]​𝐖O=∑i=1h𝐇~i​𝐖iO.\mathrm{MHA}(\mathbf{X})=[\mathbf{H}_{1},\dots,\mathbf{H}_{h}]\mathbf{W}^{O}=\sum_{i=1}^{h}\mathbf{H}_{i}\mathbf{W}^{O}_{i},\quad\widetilde{\mathrm{MHA}}(\mathbf{X})=[\widetilde{\mathbf{H}}_{1},\dots,\widetilde{\mathbf{H}}_{h}]\mathbf{W}^{O}=\sum_{i=1}^{h}\widetilde{\mathbf{H}}_{i}\mathbf{W}^{O}_{i}.Therefore,

‖MHA~​(𝐗)−MHA​(𝐗)‖2=‖∑i=1h(𝐇~i−𝐇i)​𝐖iO‖2≤∑i=1h‖(𝐇~i−𝐇i)​𝐖iO‖2.\|\widetilde{\mathrm{MHA}}(\mathbf{X})-\mathrm{MHA}(\mathbf{X})\|_{2}=\left\|\sum_{i=1}^{h}(\widetilde{\mathbf{H}}_{i}-\mathbf{H}_{i})\mathbf{W}^{O}_{i}\right\|_{2}\leq\sum_{i=1}^{h}\|(\widetilde{\mathbf{H}}_{i}-\mathbf{H}_{i})\mathbf{W}^{O}_{i}\|_{2}. For each headi∈{1,…,h}i\in\{1,\ldots,h\},

‖(𝐇~i−𝐇i)​𝐖iO‖2\displaystyle\|(\widetilde{\mathbf{H}}_{i}-\mathbf{H}_{i})\mathbf{W}^{O}_{i}\|_{2}=‖(Softmax​(𝐐i​𝐊i⊤d)​𝐕i−Softmax​(𝐐i​𝐊~i⊤d)​𝐕~i)​𝐖iO‖2\displaystyle=\big\|\big(\text{Softmax}(\tfrac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}}{\sqrt{d}})\mathbf{V}_{i}-\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})\widetilde{\mathbf{V}}_{i}\big)\mathbf{W}^{O}_{i}\big\|_{2}≤‖(Softmax​(𝐐i​𝐊i⊤d)−Softmax​(𝐐i​𝐊~i⊤d))​𝐕i​𝐖iO‖2+‖Softmax​(𝐐i​𝐊~i⊤d)​(𝐕i−𝐕~i)​𝐖iO‖2.\displaystyle\leq\big\|\big(\text{Softmax}(\tfrac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}}{\sqrt{d}})-\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})\big)\mathbf{V}_{i}\mathbf{W}^{O}_{i}\big\|_{2}+\big\|\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})(\mathbf{V}_{i}-\widetilde{\mathbf{V}}_{i})\mathbf{W}^{O}_{i}\big\|_{2}.(4) For the first term in (4), applying the submultiplicative property of the 2-norm gives

‖(Softmax​(𝐐i​𝐊i⊤d)−Softmax​(𝐐i​𝐊~i⊤d))​𝐕i​𝐖iO‖2≤‖Softmax​(𝐐i​𝐊i⊤d)−Softmax​(𝐐i​𝐊~i⊤d)‖2​‖𝐕i​𝐖iO‖2.\big\|\big(\text{Softmax}(\tfrac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}}{\sqrt{d}})-\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})\big)\mathbf{V}_{i}\mathbf{W}^{O}_{i}\big\|_{2}\leq\|\text{Softmax}(\tfrac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}}{\sqrt{d}})-\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})\|_{2}\,\|\mathbf{V}_{i}\mathbf{W}^{O}_{i}\|_{2}.The factor‖𝐕i​𝐖iO‖2\|\mathbf{V}_{i}\mathbf{W}^{O}_{i}\|_{2}acts as an amplification term, capturing how sensitivity in the value projections may magnify small perturbations in the attention weights—this term is typically not directly controllable in practice.

Since the Softmax function is12\tfrac{1}{2}-Lipschitz continuous (see Appendix A.4[alghamdi2022beyond]), we have

‖Softmax​(𝐐i​𝐊i⊤d)−Softmax​(𝐐i​𝐊~i⊤d)‖2≤1d​‖𝐐i​𝐊i⊤−𝐐i​𝐊~i⊤‖2.\|\text{Softmax}(\tfrac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}}{\sqrt{d}})-\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})\|_{2}\leq\frac{1}{\sqrt{d}}\|\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}-\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}\|_{2}. For the second term in (4), note that‖Softmax​(⋅)‖1=1\|\text{Softmax}(\cdot)\|_{1}=1, implying‖Softmax​(⋅)‖2≤1\|\text{Softmax}(\cdot)\|_{2}\leq 1. Therefore,

‖Softmax​(𝐐i​𝐊~i⊤d)​(𝐕i−𝐕~i)​𝐖iO‖2≤‖𝐕i​𝐖iO−𝐕~i​𝐖iO‖2.\big\|\text{Softmax}(\tfrac{\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}}{\sqrt{d}})(\mathbf{V}_{i}-\widetilde{\mathbf{V}}_{i})\mathbf{W}^{O}_{i}\big\|_{2}\leq\|\mathbf{V}_{i}\mathbf{W}^{O}_{i}-\widetilde{\mathbf{V}}_{i}\mathbf{W}^{O}_{i}\|_{2}. Combining the two bounds yields

‖(𝐇~i−𝐇i)​𝐖iO‖2≤‖𝐕i​𝐖iO‖2d​‖𝐐i​𝐊i⊤−𝐐i​𝐊~i⊤‖2+‖𝐕i​𝐖iO−𝐕~i​𝐖iO‖2.\|(\widetilde{\mathbf{H}}_{i}-\mathbf{H}_{i})\mathbf{W}^{O}_{i}\|_{2}\leq\frac{\|\mathbf{V}_{i}\mathbf{W}^{O}_{i}\|_{2}}{\sqrt{d}}\|\mathbf{Q}_{i}\mathbf{K}_{i}^{\top}-\mathbf{Q}_{i}\widetilde{\mathbf{K}}_{i}^{\top}\|_{2}+\|\mathbf{V}_{i}\mathbf{W}^{O}_{i}-\widetilde{\mathbf{V}}_{i}\mathbf{W}^{O}_{i}\|_{2}.Summing over all headsi=1,…,hi=1,\ldots,hgives the desired bound (Theorem), completing the proof. ∎

Appendix BValue-Output Projection

In this section, we examine the interaction between the value representations and the output projection. While the main analysis of this paper focuses on the relationship between keys and queries, the same reasoning naturally extends to values𝐕i∈ℝT×d\mathbf{V}_{i}\in\mathbb{R}^{T\times d}and the output matrix𝐖iO∈ℝd×D\mathbf{W}_{i}^{O}\in\mathbb{R}^{d\times D}. To further tighten the upper bound established in Theorem1, we aim to minimize the second term in the summation, which leads to the following optimization problem:

min𝐕~⁡‖𝐕𝐖O−𝐕~​𝐖O‖F2s.t.​rank​(𝐕~)≤R,\min_{\widetilde{\mathbf{V}}}\,\big\|\mathbf{V}\mathbf{W}^{O}-\widetilde{\mathbf{V}}\mathbf{W}^{O}\big\|_{F}^{2}~~~~\text{s.t.}~~~\mathrm{rank}(\widetilde{\mathbf{V}})\leq R,where for simplicity we drop the subscript. Likewise for the keys and queries case, the optimization problem can be restated using a projection matrix𝐒∈ℝd×d\mathbf{S}\in\mathbb{R}^{d\times d}such that

min𝐒⁡‖𝐕𝐒𝐖O−𝐕𝐖O‖F2s.t.​rank​(𝐒)≤R,\min_{\mathbf{S}}\|\mathbf{V}\mathbf{S}\mathbf{W}^{O}-\mathbf{V}\mathbf{W}^{O}\|_{F}^{2}~~~~~~\text{s.t.}~~~\mathrm{rank}(\mathbf{S})\leq R,where𝐒∈ℝd×d\mathbf{S}\in\mathbb{R}^{d\times d}where we write the projection matrix𝐒\mathbf{S}as a product of two matrices, i.e.,𝐒=𝐀𝐁⊤\mathbf{S}=\mathbf{A}\mathbf{B}^{\top}with𝐀,𝐁∈ℝd×R\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}. The minimization problem tackled by KQ-SVD for the value and output projections is

min𝐀,𝐁∈ℝd×R⁡‖𝐕𝐀𝐁⊤​𝐖O−𝐕𝐖O‖F2.\displaystyle\min_{\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times R}}\|\mathbf{V}\mathbf{A}\mathbf{B}^{\top}\mathbf{W}^{O}-\mathbf{V}\mathbf{W}^{O}\|_{F}^{2}.(5) Theorem2applies, in a similar fashion to the case of values and outputs.

Appendix CPractical settings

Code:The code used for experiments presented in Section6is available athttps://github.com/DamienLesens/KQ-SVD. We used Pytorch[pytorch]and the Hugging Facetransformerslibrary[wolf2020huggingfacestransformersstateoftheartnatural]. For reproducibility, we fixed random seeds equals to 0.

Hardware:All experiments were conducted on NVIDIA Tesla V100-SXM2-32GB GPUs with CUDA acceleration. The primary compute nodes were Intel Xeon E5-2698 v4 @ 2.20GHz (503GB RAM).