Kronecker Product Attention: Memory-Efficient KV Compression via Structured Factorization

arXiv cs.LG Papers

Summary

提出 Kronecker Product Attention (KPA),用 Kronecker 积紧凑矩阵因子表示键值激活,实现最高 10× 的 KV 缓存压缩,同时在多个模型规模和基准上匹配或超越 TPA、MQA、GQA 和 MLA 等方法。

arXiv:2610.12657v1 Announce Type: new Abstract: Efficient autoregressive generation is increasingly limited by the memory footprint of the key-value (KV) cache, which grows linearly with sequence length, batch size, and model depth. Existing approaches such as Multi-Query Attention (MQA), Grouped-Query Attention (GQA), Multi-Head Latent Attention (MLA), and Tensor Product Attention (TPA) reduce this cost through different trade-offs, but often sacrifice expressivity, require additional positional-encoding machinery, or impose restrictive low-rank structure. We propose Kronecker Product Attention (KPA), which represents key and value activations as sums of Kronecker products of compact, token-dependent matrix factors. KPA strictly generalizes TPA: rank-one vector-outer-product factorization arises as a degenerate Kronecker split, while non-degenerate matrix factors can represent richer block-structured patterns. We further show that a structured rotary embedding operator matched to the Kronecker representation preserves the relative-position property of RoPE while rotating only the compact factors, removing the need for a decoupled RoPE branch as in MLA. Empirically, KPA matches or exceeds TPA, MHA, MQA, GQA, and MLA across four model scales and nine downstream benchmarks under identical training conditions, while achieving up to 10x KV cache compression. The Kronecker head split is a tuneable hyperparameter that allows practitioners to navigate the compression-accuracy trade-off at design time without architectural changes.
Original Article
View Cached Full Text

Cached at: 10/12/26, 06:41 AM

# Kronecker Product Attention: Memory-Efficient KV Compression via Structured Factorization
Source: [https://arxiv.org/html/2610.12657](https://arxiv.org/html/2610.12657)
\\workshoptitle

SLM\-Agents: 1st Workshop on Small Language Models for Agentic Systems

Bahareh NikpourHossein Rajabzadeh††thanks:This work was done while the author was affiliated with Huawei Canada\.Affiliation:University of Waterloo, Waterloo, CanadaEmail:[parsa\.omidi@huawei\.com](mailto:)Maryam Dialameh11footnotemark:1Affiliation:University of Waterloo, Waterloo, CanadaVahid Partovi NiaAffiliation:Huawei Noah’s Ark Labbahar\.nikpour@h\-partners\.com,\{hossein\.rajabzadeh, maryam\.dialameh\}@uwaterloo\.ca,Parsa OmidiAffiliation:Huawei Technologies, Toronto, Canada

###### Abstract

Efficient autoregressive generation is increasingly limited by the memory footprint of the key\-value \(KV\) cache, which grows linearly with sequence length, batch size, and model depth\. Existing approaches such as Multi\-Query Attention \(MQA\), Grouped\-Query Attention \(GQA\), Multi\-Head Latent Attention \(MLA\), and Tensor Product Attention \(TPA\) reduce this cost through different trade\-offs, but often sacrifice expressivity, require additional positional\-encoding machinery, or impose restrictive low\-rank structure\. We propose Kronecker Product Attention \(KPA\), which represents key and value activations as sums of Kronecker products of compact, token\-dependent matrix factors\. KPA strictly generalizes TPA: rank\-one vector\-outer\-product factorization arises as a degenerate Kronecker split, while non\-degenerate matrix factors can represent richer block\-structured patterns\. We further show that a structured rotary embedding operator matched to the Kronecker representation preserves the relative\-position property of RoPE while rotating only the compact factors, removing the need for a decoupled RoPE branch as in MLA\. Empirically, KPA matches or exceeds TPA, MHA, MQA, GQA, and MLA across four model scales and nine downstream benchmarks under identical training conditions, while achieving up to10×10\\timesKV cache compression\. The Kronecker head split is a tuneable hyperparameter that allows practitioners to navigate the compression–accuracy trade\-off at design time without architectural changes\.

## 1Introduction

The deployment of large\-scale Transformer\-based language models is increasingly bottlenecked by inference\-time memory rather than parameter count\. At the core of this bottleneck is the key\-value \(KV\) cache, which stores projected keys and values across every layer and token so they do not need to be recomputed during decoding\. This cache grows linearly with sequence length, batch size, and head dimensionality\. As context windows extend to hundreds of thousands of tokens and batch sizes increase for throughput, the KV cache can dominate total GPU memory consumption, leaving little headroom for model weights or activations and directly limiting the context length, batch size, and deployment scale that practitioners can afford, motivating a broad line of work on KV compression\.

Existing approaches address this along several lines\. Token eviction and sparse attention methods risk discarding relevant tokens\([Zhang et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib30);[Xiao et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib27)\)\. Others propose off\-chip storage of key\-value states, trading memory pressure for increased I/O latency\. At the architectural level, Multi\-Query Attention \(MQA\)[Shazeer \(2019\)](https://arxiv.org/html/2610.12657#bib.bib19)and Grouped\-Query Attention \(GQA\)\([Ainslie et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib1)\)share keys and values across heads at the cost of representational diversity\. Multi\-Head Latent Attention \(MLA\)\([DeepSeek\-AI, 2024](https://arxiv.org/html/2610.12657#bib.bib5)\)avoids this by caching compressed latent representations instead, but requires a decoupled RoPE branch, in which a small set of uncompressed key dimensions carry the rotary embedding and are concatenated to the latent key at each decoding step\. Tensor Product Attention \(TPA\) imposes low\-rank CP \(Canonical Polyadic\) structure directly on KV projections, storing only the low\-dimensional factors\([Zhang et al\., 2025](https://arxiv.org/html/2610.12657#bib.bib29)\)\. However, CP decomposition operates on vectors and cannot jointly encode structure at both the head level and the feature level within each head\. This limitation is important as empirical studies have shown that attention heads specialize into distinct functional roles, with different subsets attending to syntactic, semantic, and positional patterns, suggesting that an effective KV factorization should operate at both levels simultaneously\([Michel et al\., 2019](https://arxiv.org/html/2610.12657#bib.bib13);[Voita et al\., 2019](https://arxiv.org/html/2610.12657#bib.bib25)\)\.

Figure 1:Overview of Multi\-Head Attention \(MHA\), Tensor Product Attention \(TPA\), and Kronecker Product Attention \(KPA, ours\)\. MHA projectsxxdirectly into denseQ,K,VQ,K,V\. TPA and KPA instead factor each projection into two components, combined via an outer product \(×\\times\) for TPA or a Kronecker product \(⊗\\otimes\) for KPA\. RoPE is applied only toQQandKK\(neverVV\), and rotates only theBBfactor in both methods \(Theorem[2\.4](https://arxiv.org/html/2610.12657#S2.Thmtheorem4)formalizes this for KPA\);AAis cached unrotated\. Since TPA’sA,BA,Bare vectors, their outer product is a flat, unstructured matrix; since KPA’sA,BA,Bare matrices, their Kronecker product is block\-structured with a two\-level head/feature hierarchy \(shown by the internal grid lines\)\. OnlyK,VK,Vfactors enter the KV cache;QQis not cached\. The token indexttis omitted throughout except on the finalQt,Kt,VtQ\_\{t\},K\_\{t\},V\_\{t\}, for clarity\.Motivated by this, we use Kronecker product structure to represent per\-token key and value matrices via two smaller contextual matrix factors\. Unlike CP decomposition which operates on vectors, the Kronecker product expands two small matrices into a larger block\-structured matrix with the same shape as a standard key or value slice, making it naturally suited to multi\-head attention where both the head axis and the within\-head feature axis have product structure\. Among structured factorizations, the Kronecker product is distinguished by combining three properties in a single operation: its mixed\-product identity lets inner products between Kronecker\-structured queries and keys factor into products of smaller inner products; it admits a Kronecker\-separable rotary operator that acts on only one compact factor; and a rank\-1 outer product, which is the basic building block of CP decomposition, is a special case, making Kronecker factorization strictly more general than CP at the same parameter budget\. We note that Kronecker structure has also been explored in attention mechanisms for encoder models operating on multi\-dimensional inputs\([Omranpour et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib15)\), reducing training\-time computation rather than KV cache size\. That setting differs fundamentally from autoregressive decoding, which is our focus\.

We introduceKroneckerProductAttention\(KPA\), a structured attention mechanism that applies this factorization to the query, key, and value projections of each attention head\. This yields a family of KV compression schemes parameterized by the factor dimensions, allowing flexible trade\-offs between compression ratio, head structure, and representational capacity\. We provide a theoretical analysis characterizing the expressivity of KPA as a function of rank and factor dimensions, and prove that a modified frequency schedule we call Kronecker\-Separable RoPE integrates seamlessly with the factorized KV representation without materializing the full key or value tensors and without breaking the standard relative\-position property\. To evaluate KPA empirically, we train models at four scales including 124M, 353M, 773M, and 1\.5B parameters on the FineWeb\-Edu corpus\. We compare KPA variants against TPA at matched compression ratios, finding that KPA matches or exceeds TPA across all scales\. Figure 1 summarizes the contrast between MHA, TPA, and KPA\. Our contributions are as follows:

- •We propose Kronecker Product Attention \(KPA\), the first application of Kronecker product factorization to*per\-token, contextual*key and value representations for KV cache compression in autoregressive decoding\. Unlike prior Kronecker attention work that decomposes static attention score matrices for encoder models, KPA operates on input\-dependent factors and directly reduces inference\-time memory\.
- •We introduce Kronecker\-Separable RoPE, which applies rotary position encoding directly to the compact key factors, preserving the standard relative\-position property without a decoupled RoPE branch as in MLA\.
- •We train KPA at four scales \(125M–1\.5B\) on FineWeb\-Edu under identical conditions to TPA, showing KPA matches or exceeds TPA, MHA, MQA, GQA, and MLA across nine benchmarks while achieving6\.4×6\.4\\times–10×10\\timesKV cache compression\.

## 2Kronecker Product Attention

The key\-value cache in standard multi\-head attention stores, for each tokentt, a pair of matrices𝑲t,𝑽t∈ℝh×dh\\bm\{K\}\_\{t\},\\bm\{V\}\_\{t\}\\in\\mathbb\{R\}^\{h\\times d\_\{h\}\}, whose entries are dense and unstructured\. We argue that this is wasteful: attention heads are known to specialise into distinct functional roles\([Michel et al\., 2019](https://arxiv.org/html/2610.12657#bib.bib13);[Voita et al\., 2019](https://arxiv.org/html/2610.12657#bib.bib25)\), suggesting that the*head\-level*dimension of𝑲t\\bm\{K\}\_\{t\}and𝑽t\\bm\{V\}\_\{t\}\(i\.e\., how information is shared or differs across heads\) carries structured, differentiable information\. We hypothesise that the*feature\-level*dimension \(variation across features within a single head\) benefits from an analogous structure, and design KPA\. The proposed KPA makes this two\-level structure explicit by factorising each key and value matrix as a sum of Kronecker products of two smaller matrices, one operating on a coarse partition of the head space, the other on the per\-head feature space \(formalised in Section[2\.2](https://arxiv.org/html/2610.12657#S2.SS2)\), and caching only these factors\. The result is a compression scheme that is structurally aligned with the functional organisation of multi\-head attention in a way that CP decomposition, which operates on vectors rather than matrices, cannot achieve\.

### 2\.1Background and Notation

#### Notation\.

Bold uppercase denotes matrices \(𝑲t,𝑽t,𝑸t,𝑨,𝑩,𝑾K\\bm\{K\}\_\{t\},\\bm\{V\}\_\{t\},\\bm\{Q\}\_\{t\},\\bm\{A\},\\bm\{B\},\\bm\{W\}^\{K\}\); bold lowercase denotes vectors \(𝒒t,𝒌t,𝒗t,𝒙t\\bm\{q\}\_\{t\},\\bm\{k\}\_\{t\},\\bm\{v\}\_\{t\},\\bm\{x\}\_\{t\}\); unbolded lowercase denotes scalars and indices \(h,dh,dmodel,RK,t,ih,d\_\{h\},d\_\{\\mathrm\{model\}\},R\_\{K\},t,i\); and calligraphic uppercase denotes higher\-order tensors \(𝒦,𝒱,𝒬\\mathcal\{K\},\\mathcal\{V\},\\mathcal\{Q\}\)\. We write⊗\\otimesfor the Kronecker product,𝑰n\\bm\{I\}\_\{n\}for then×nn\\times nidentity, and𝑹θ​\(t\)\\bm\{R\}\_\{\\theta\}\(t\)for the RoPE rotation at positiontt\. Subscripts index sequence positions; superscripts identify projections\. We set\[n\]:=\{1,…,n\}\[n\]:=\\\{1,\\ldots,n\\\}, and all vectors are column vectors unless stated otherwise\.

###### Definition 2\.1\(Kronecker product\)\.

For matrices𝐀∈ℝm1×n1\\bm\{A\}\\in\\mathbb\{R\}^\{m\_\{1\}\\times n\_\{1\}\}and𝐁∈ℝm2×n2\\bm\{B\}\\in\\mathbb\{R\}^\{m\_\{2\}\\times n\_\{2\}\}, the Kronecker product𝐀⊗𝐁∈ℝm1​m2×n1​n2\\bm\{A\}\\otimes\\bm\{B\}\\in\\mathbb\{R\}^\{m\_\{1\}m\_\{2\}\\times n\_\{1\}n\_\{2\}\}is defined as

𝑨⊗𝑩=\(a11​𝑩⋯a1​n1​𝑩⋱am1​1​𝑩⋯am1​n1​𝑩\)\.\\bm\{A\}\\otimes\\bm\{B\}=\\begin\{pmatrix\}a\_\{11\}\\bm\{B\}&\\cdots&a\_\{1n\_\{1\}\}\\bm\{B\}\\\\ \\vdots&\\ddots&\\vdots\\\\ a\_\{m\_\{1\}1\}\\bm\{B\}&\\cdots&a\_\{m\_\{1\}n\_\{1\}\}\\bm\{B\}\\end\{pmatrix\}\.

The central algebraic property used in KPA is the mixed\-product identity\([Van Loan and Pitsianis, 1993](https://arxiv.org/html/2610.12657#bib.bib24)\)which states for conformable matrices𝑴,𝑵,𝑷,𝑸\\bm\{M\},\\bm\{N\},\\bm\{P\},\\bm\{Q\}, we can write:

\(𝑴⊗𝑵\)​\(𝑷⊗𝑸\)=\(𝑴​𝑷\)⊗\(𝑵​𝑸\)\.\(\\bm\{M\}\\otimes\\bm\{N\}\)\(\\bm\{P\}\\otimes\\bm\{Q\}\)=\(\\bm\{M\}\\bm\{P\}\)\\otimes\(\\bm\{N\}\\bm\{Q\}\)\.\(1\)
This identity underlies both the factored score computation \(Section[2\.3](https://arxiv.org/html/2610.12657#S2.SS3)\) and the RoPE compatibility proof \(Section[2\.4](https://arxiv.org/html/2610.12657#S2.SS4)\)\.

### 2\.2Factorised Key, Value, and Query Representations

Leth=h1​h2h=h\_\{1\}h\_\{2\}anddh=d1​d2d\_\{h\}=d\_\{1\}d\_\{2\}, whereh1,h2,d1,d2h\_\{1\},h\_\{2\},d\_\{1\},d\_\{2\}are hyperparameters that control the two\-level structure\. For each tokenttwith hidden state𝒙t∈ℝdmodel\\bm\{x\}\_\{t\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{model\}\}\}, KPA represents the key, value, and query slices as normalized sums of contextual Kronecker products:

𝑸t\\displaystyle\\bm\{Q\}\_\{t\}=1RQ​∑r=1RQ𝑨rQ​\(𝒙t\)⊗𝑩rQ​\(𝒙t\)∈ℝh×dh,\\displaystyle=\\frac\{1\}\{R\_\{Q\}\}\\sum\_\{r=1\}^\{R\_\{Q\}\}\\bm\{A\}^\{Q\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\otimes\\bm\{B\}^\{Q\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\;\\in\\mathbb\{R\}^\{h\\times d\_\{h\}\},\(2\)𝑲t\\displaystyle\\bm\{K\}\_\{t\}=1RK​∑r=1RK𝑨rK​\(𝒙t\)⊗𝑩rK​\(𝒙t\)∈ℝh×dh,\\displaystyle=\\frac\{1\}\{R\_\{K\}\}\\sum\_\{r=1\}^\{R\_\{K\}\}\\bm\{A\}^\{K\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\otimes\\bm\{B\}^\{K\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\;\\in\\mathbb\{R\}^\{h\\times d\_\{h\}\},\(3\)𝑽t\\displaystyle\\bm\{V\}\_\{t\}=1RV​∑r=1RV𝑨rV​\(𝒙t\)⊗𝑩rV​\(𝒙t\)∈ℝh×dh,\\displaystyle=\\frac\{1\}\{R\_\{V\}\}\\sum\_\{r=1\}^\{R\_\{V\}\}\\bm\{A\}^\{V\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\otimes\\bm\{B\}^\{V\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\;\\in\\mathbb\{R\}^\{h\\times d\_\{h\}\},\(4\)where𝑨r⋅​\(𝒙t\)∈ℝh1×d1\\bm\{A\}^\{\\cdot\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\in\\mathbb\{R\}^\{h\_\{1\}\\times d\_\{1\}\}operates on the head\-partition subspaceℝh1×d1\\mathbb\{R\}^\{h\_\{1\}\\times d\_\{1\}\}and𝑩r⋅​\(𝒙t\)∈ℝh2×d2\\bm\{B\}^\{\\cdot\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\in\\mathbb\{R\}^\{h\_\{2\}\\times d\_\{2\}\}operates on the within\-head subspaceℝh2×d2\\mathbb\{R\}^\{h\_\{2\}\\times d\_\{2\}\}\. Both factors are fully contextual meaning they depend on𝒙t\\bm\{x\}\_\{t\}through independent learned linear projections as:

𝑨rK​\(𝒙t\)=𝑾rAK​𝒙t,𝑩rK​\(𝒙t\)=𝑾rBK​𝒙t,\\bm\{A\}^\{K\}\_\{r\}\(\\bm\{x\}\_\{t\}\)=\\bm\{W\}^\{A^\{K\}\}\_\{r\}\\,\\bm\{x\}\_\{t\},\\qquad\\bm\{B\}^\{K\}\_\{r\}\(\\bm\{x\}\_\{t\}\)=\\bm\{W\}^\{B^\{K\}\}\_\{r\}\\,\\bm\{x\}\_\{t\},\(5\)with𝑾rAK∈ℝ\(h1​d1\)×dmodel\\bm\{W\}^\{A^\{K\}\}\_\{r\}\\in\\mathbb\{R\}^\{\(h\_\{1\}d\_\{1\}\)\\times d\_\{\\mathrm\{model\}\}\}and𝑾rBK∈ℝ\(h2​d2\)×dmodel\\bm\{W\}^\{B^\{K\}\}\_\{r\}\\in\\mathbb\{R\}^\{\(h\_\{2\}d\_\{2\}\)\\times d\_\{\\mathrm\{model\}\}\}reshaped after multiplication\. Analogous maps define the query and value factors\. In practice, we merge the rank index into aggregate weight matrices𝑾AK∈ℝ\(RK​h1​d1\)×dmodel\\bm\{W\}^\{A^\{K\}\}\\in\\mathbb\{R\}^\{\(R\_\{K\}h\_\{1\}d\_\{1\}\)\\times d\_\{\\mathrm\{model\}\}\}and𝑾BK∈ℝ\(RK​h2​d2\)×dmodel\\bm\{W\}^\{B^\{K\}\}\\in\\mathbb\{R\}^\{\(R\_\{K\}h\_\{2\}d\_\{2\}\)\\times d\_\{\\mathrm\{model\}\}\}, then reshape to𝒜K​\(𝒙t\)∈ℝRK×h1×d1\\mathcal\{A\}^\{K\}\(\\bm\{x\}\_\{t\}\)\\in\\mathbb\{R\}^\{R\_\{K\}\\times h\_\{1\}\\times d\_\{1\}\}andℬK​\(𝒙t\)∈ℝRK×h2×d2\\mathcal\{B\}^\{K\}\(\\bm\{x\}\_\{t\}\)\\in\\mathbb\{R\}^\{R\_\{K\}\\times h\_\{2\}\\times d\_\{2\}\}, enabling a single batched projection per token\.

The ranksRQ,RK,RVR\_\{Q\},R\_\{K\},R\_\{V\}are independent\. Since queries are never cached, one may freely setRQ≫RK,RVR\_\{Q\}\\gg R\_\{K\},R\_\{V\}to enrich query expressivity at no memory cost\. Only the𝑲\\bm\{K\}and𝑽\\bm\{V\}factors enter the cache\.

It should be noted that settingh1=hh\_\{1\}=h,d1=1d\_\{1\}=1,h2=1h\_\{2\}=1,d2=dhd\_\{2\}=d\_\{h\}makes𝑨r∈ℝh×1\\bm\{A\}\_\{r\}\\in\\mathbb\{R\}^\{h\\times 1\}and𝑩r∈ℝ1×dh\\bm\{B\}\_\{r\}\\in\\mathbb\{R\}^\{1\\times d\_\{h\}\}vectors, so each Kronecker term reduces to a rank\-one outer product, exactly recovering TPA\([Zhang et al\., 2025](https://arxiv.org/html/2610.12657#bib.bib29)\)\. Non\-degenerate Kronecker factors are strictly richer:rank⁡\(𝑨r⊗𝑩r\)=rank⁡\(𝑨r\)⋅rank⁡\(𝑩r\)\\mathrm\{rank\}\(\\bm\{A\}\_\{r\}\\otimes\\bm\{B\}\_\{r\}\)=\\mathrm\{rank\}\(\\bm\{A\}\_\{r\}\)\\cdot\\mathrm\{rank\}\(\\bm\{B\}\_\{r\}\), giving KPA strictly greater representational capacity at matched parameter budgets\.

### 2\.3Factored Attention Score Computation

Heads in KPA are indexed by pairs\(i1,i2\)\(i\_\{1\},i\_\{2\}\)withi1∈\[h1\]i\_\{1\}\\in\[h\_\{1\}\],i2∈\[h2\]i\_\{2\}\\in\[h\_\{2\}\]\. Denoting thei1i\_\{1\}\-th row of𝑨rQ\\bm\{A\}^\{Q\}\_\{r\}as\(𝑨rQ\)i1∈ℝd1\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\in\\mathbb\{R\}^\{d\_\{1\}\}, the mixed\-product identity \([1](https://arxiv.org/html/2610.12657#S2.E1)\) gives an exact factorisation of the attention score \(scaled by1/dh1/\\sqrt\{d\_\{h\}\}before softmax in practice\):

###### Proposition 2\.2\(Factored Score Computation\)\.

The attention score for head\(i1,i2\)\(i\_\{1\},i\_\{2\}\)at positionsttandsssatisfies

score\(i1,i2\)​\(t,s\)=1RQ​RK​∑r=1RQ∑r′=1RK\(𝑨rQ\)i1⊤​\(𝑨r′K\)i1⏟head\-level​\(𝑩rQ\)i2⊤​\(𝑩r′K\)i2⏟feature\-level\.\\mathrm\{score\}\_\{\(i\_\{1\},i\_\{2\}\)\}\(t,s\)=\\frac\{1\}\{R\_\{Q\}R\_\{K\}\}\\sum\_\{r=1\}^\{R\_\{Q\}\}\\sum\_\{r^\{\\prime\}=1\}^\{R\_\{K\}\}\\underbrace\{\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}^\{\\top\}\(\\bm\{A\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{1\}\}\}\_\{\\text\{head\-level\}\}\\underbrace\{\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}^\{\\top\}\(\\bm\{B\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{2\}\}\}\_\{\\text\{feature\-level\}\}\.\(6\)

Every attention score decomposes into a product of a*head\-level*and a*feature\-level*inner product, reflecting the two\-level structure of KPA\. Proposition[2\.2](https://arxiv.org/html/2610.12657#S2.Thmtheorem2)shows that scores can in principle be computed directly from the compact factors without materializing the full key tensor, analogously to factored score computation in TPA\([Zhang et al\., 2025](https://arxiv.org/html/2610.12657#bib.bib29)\)\.

### 2\.4Kronecker\-Separable RoPE

A natural concern with any structured KV factorisation is compatibility with rotary positional embeddings \(RoPE\)\([Su et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib21)\)\. MLA resolves this by maintaining a small set of non\-compressed KV dimensions that carry the rotary encoding, introducing a decoupled RoPE branch at every decoding step\. We show that KPA requires no such branch\.

###### Definition 2\.3\(Kronecker\-Separable RoPE\)\.

Let𝐈d1\\bm\{I\}\_\{d\_\{1\}\}be thed1×d1d\_\{1\}\\times d\_\{1\}identity and𝐑θ\(2\)​\(t\)∈ℝd2×d2\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)\\in\\mathbb\{R\}^\{d\_\{2\}\\times d\_\{2\}\}be a RoPE rotation matrix with frequencies calibrated tod2d\_\{2\}:θj=10000−2j/d2\\theta\_\{j\}=10000^\{\-2j/d\_\{2\}\},j=0,…,d2/2−1j=0,\\ldots,d\_\{2\}/2\-1\. The*Kronecker\-Separable RoPE*operator onℝdh\\mathbb\{R\}^\{d\_\{h\}\}is𝐑KPA​\(t\)=𝐈d1⊗𝐑θ\(2\)​\(t\)\\bm\{R\}\_\{\\mathrm\{KPA\}\}\(t\)=\\bm\{I\}\_\{d\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)\.

###### Theorem 2\.4\(RoPE Compatibility of KPA\)\.

Under Kronecker\-Separable RoPE, applying𝐈d1⊗𝐑θ\(2\)​\(⋅\)\\bm\{I\}\_\{d\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(\\cdot\)to the query or key vector of head\(i1,i2\)\(i\_\{1\},i\_\{2\}\)is equivalent to rotating only the𝐁\\bm\{B\}factor:

\(𝑰d1⊗𝑹θ\(2\)​\(t\)\)​\(\(𝑨rQ\)i1⊗\(𝑩rQ\)i2\)\\displaystyle\\bigl\(\\bm\{I\}\_\{d\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)\\bigr\)\\bigl\(\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\otimes\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}\\bigr\)=\(𝑨rQ\)i1⊗𝑹θ\(2\)​\(t\)​\(𝑩rQ\)i2,\\displaystyle=\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\},\(7\)\(𝑰d1⊗𝑹θ\(2\)​\(s\)\)​\(\(𝑨r′K\)i1⊗\(𝑩r′K\)i2\)\\displaystyle\\bigl\(\\bm\{I\}\_\{d\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(s\)\\bigr\)\\bigl\(\(\\bm\{A\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{1\}\}\\otimes\(\\bm\{B\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{2\}\}\\bigr\)=\(𝑨r′K\)i1⊗𝑹θ\(2\)​\(s\)​\(𝑩r′K\)i2\.\\displaystyle=\(\\bm\{A\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(s\)\(\\bm\{B\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{2\}\}\.\(8\)Moreover, the attention score depends only on the relative offsets−ts\-t:

score\(i1,i2\)​\(t,s\)=1RQ​RK​∑r,r′\[\(𝑨rQ\)i1⊤​\(𝑨r′K\)i1\]⋅\[\(𝑩rQ\)i2⊤​𝑹θ\(2\)​\(s−t\)​\(𝑩r′K\)i2\]\.\\mathrm\{score\}\_\{\(i\_\{1\},i\_\{2\}\)\}\(t,s\)=\\frac\{1\}\{R\_\{Q\}R\_\{K\}\}\\sum\_\{r,r^\{\\prime\}\}\\Bigl\[\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}^\{\\top\}\(\\bm\{A\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{1\}\}\\Bigr\]\\cdot\\Bigl\[\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}^\{\\top\}\\,\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(s\-t\)\\,\(\\bm\{B\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{2\}\}\\Bigr\]\.\(9\)

Theorem[2\.4](https://arxiv.org/html/2610.12657#S2.Thmtheorem4)has two practical consequences: \(i\) only the𝑩K\\bm\{B\}^\{K\}factors need to be pre\-rotated before caching; the𝑨K\\bm\{A\}^\{K\}factors are position\-independent and stored as\-is, and \(ii\) the𝑩\\bm\{B\}factor structure is preserved by the rotation, so positional encoding need not require reconstructing the full key tensor\. Unlike MLA, no decoupled RoPE branch is required\.

### 2\.5KV Cache Cost and Compression

During autoregressive decoding, KPA caches the key and value factors per tokentt:

\{𝑨rK​\(𝒙t\),𝑩rK​\(𝒙t\)\}r=1RK∪\{𝑨rV​\(𝒙t\),𝑩rV​\(𝒙t\)\}r=1RV,\\\{\\bm\{A\}^\{K\}\_\{r\}\(\\bm\{x\}\_\{t\}\),\\,\\bm\{B\}^\{K\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\\}\_\{r=1\}^\{R\_\{K\}\}\\;\\cup\\;\\\{\\bm\{A\}^\{V\}\_\{r\}\(\\bm\{x\}\_\{t\}\),\\,\\bm\{B\}^\{V\}\_\{r\}\(\\bm\{x\}\_\{t\}\)\\\}\_\{r=1\}^\{R\_\{V\}\},where𝑨rK∈ℝh1×d1\\bm\{A\}^\{K\}\_\{r\}\\in\\mathbb\{R\}^\{h\_\{1\}\\times d\_\{1\}\}and𝑩rK∈ℝh2×d2\\bm\{B\}^\{K\}\_\{r\}\\in\\mathbb\{R\}^\{h\_\{2\}\\times d\_\{2\}\}\(pre\-rotated via Kronecker\-Separable RoPE\)\. The per\-token cache cost is\(RK\+RV\)​\(h1​d1\+h2​d2\)\(R\_\{K\}\+R\_\{V\}\)\(h\_\{1\}d\_\{1\}\+h\_\{2\}d\_\{2\}\), compared with2​h​dh=2​h1​h2​d1​d22hd\_\{h\}=2h\_\{1\}h\_\{2\}d\_\{1\}d\_\{2\}for MHA\. The compression factor is:

CKPA=2​h1​h2​d1​d2\(RK\+RV\)​\(h1​d1\+h2​d2\)\.C\_\{\\mathrm\{KPA\}\}=\\frac\{2h\_\{1\}h\_\{2\}d\_\{1\}d\_\{2\}\}\{\(R\_\{K\}\+R\_\{V\}\)\(h\_\{1\}d\_\{1\}\+h\_\{2\}d\_\{2\}\)\}\.\(10\)For fixedhhanddhd\_\{h\}, compression increases with factor asymmetry: a more unequal split ofhhordhd\_\{h\}makesh1​d1\+h2​d2h\_\{1\}d\_\{1\}\+h\_\{2\}d\_\{2\}smaller relative toh1​h2​d1​d2h\_\{1\}h\_\{2\}d\_\{1\}d\_\{2\}\.

## 3Experiments

### 3\.1Experimental Setup

We train KPA models at four scales \(125M, 353M, 773M, 1\.5B\) using a Llama\-inspired architecture with SwiGLU and RoPE on top of nanoGPT\([Karpathy, 2022](https://arxiv.org/html/2610.12657#bib.bib10)\), trained on FineWeb\-Edu for∼\\sim50B tokens withRK=RV=2R\_\{K\}\{=\}R\_\{V\}\{=\}2andRQ=6R\_\{Q\}\{=\}6\. We compare against MHA, MQA, GQA, MLA, and TPA baselines from[Zhang et al\. \(2025\)](https://arxiv.org/html/2610.12657#bib.bib29)under identical conditions, evaluating zero\-shot on nine benchmarks usinglm\-evaluation\-harness\([Gao et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib8)\)\. Full training and evaluation details are in Appendix C\.

### 3\.2Main Results

Tables[1](https://arxiv.org/html/2610.12657#S3.T1)and[2](https://arxiv.org/html/2610.12657#S3.T2)report zero\-shot benchmark performance at 353M and 773M scales; 1\.5B and 125M results appear in Appendix D\.

At 353M, all three KPA configurations outperform TPA\. KPAh=\(8×2\)h\{=\}\(8\{\\times\}2\)achieves the best average at6\.4×6\.4\\timescompression, matching TPA’s ratio while exceeding its accuracy by 1\.49pp, with the largest gains on ARC\-Easy and SciQ\. KPAh=\(4×4\)h\{=\}\(4\{\\times\}4\)outperforms TPA at strictly higher compression \(8×8\\times\), demonstrating that increasing compression need not sacrifice accuracy\. At 773M, KPAh=\(10×2\)h\{=\}\(10\{\\times\}2\)outperforms TPA, MLA, and all other baselines at6\.7×6\.7\\timescompression, with particularly strong gains on ARC\-Easy, BoolQ, and SciQ\. Notably, KPA surpasses MLA despite requiring no decoupled RoPE branch\. KPAh=\(4×5\)h\{=\}\(4\{\\times\}5\)achieves better accuracy than TPA at higher compression \(8\.9×8\.9\\times\)\.

In both cases, the best\-performing KPA variant uses an asymmetric split with largerh1h\_\{1\}, consistent with the ablation findings in Section[3\.3](https://arxiv.org/html/2610.12657#S3.SS3)\. KPA improves monotonically across scales and matches or exceeds TPA at comparable or higher compression ratios at every scale\.

### 3\.3Ablation on Head\-Split Asymmetry

We ablate the effect of the Kronecker head splith1×h2=hh\_\{1\}\\times h\_\{2\}=h, holdingd1=d2=8d\_\{1\}=d\_\{2\}=8andRK=RV=2R\_\{K\}=R\_\{V\}=2fixed\. At 353M we compare symmetrich=\(4×4\)h\{=\}\(4\{\\times\}4\), asymmetrich=\(2×8\)h\{=\}\(2\{\\times\}8\), and reversed asymmetrich=\(8×2\)h\{=\}\(8\{\\times\}2\); at 773M we compareh=\(4×5\)h\{=\}\(4\{\\times\}5\),h=\(2×10\)h\{=\}\(2\{\\times\}10\), andh=\(10×2\)h\{=\}\(10\{\\times\}2\)\. Tables[1](https://arxiv.org/html/2610.12657#S3.T1)and[2](https://arxiv.org/html/2610.12657#S3.T2)reveal a consistent pattern across scales\. Asymmetric splits with largerh1h\_\{1\}outperform both symmetric KPA and TPA at both scales\. At 353M,h=\(8×2\)h\{=\}\(8\{\\times\}2\)achieves the best average, and at 773M,h=\(10×2\)h\{=\}\(10\{\\times\}2\)outperforms all baselines\. We attribute this to the𝑨\\bm\{A\}\-factor’s capacity to encode joint interactions acrossh1h\_\{1\}heads simultaneously in feature space\. A largerh1h\_\{1\}expands this capacity, whereas the𝑩\\bm\{B\}\-factor refines within\-head feature structure independently, which appears less critical\. This is distinct from MHA, where head projections are fully independent by construction\. Symmetric splits trade some accuracy for higher compression \(8×8\\timesto8\.9×8\.9\\times\), offering a favourable operating point when memory is the primary constraint\. The head splith1×h2h\_\{1\}\\times h\_\{2\}is a tuneable hyperparameter that allows practitioners to navigate the compression–accuracy trade\-off at design time, with no other architectural changes required\.

Table 1:Zero\-shot evaluation on nine benchmarks \(FineWeb\-Edu,∼\\sim50B tokens\), 353M\-parameter models\. Best score per column inbold\. HellaSw\. = HellaSwag, W\.G\. = WinoGrande\.MethodARC\-EARC\-CBoolQHellaSw\.OBQAPIQAW\.G\.MMLUSciQAvg\.MHA59\.5129\.5259\.6045\.6834\.2068\.8253\.4323\.3376\.9050\.11MQA57\.6231\.9159\.4545\.6935\.4069\.3153\.5126\.4774\.6050\.44GQA58\.6731\.4858\.2945\.4535\.2068\.5054\.4624\.5876\.5050\.35MLA56\.6529\.5257\.8346\.0534\.6069\.4252\.8024\.6279\.7050\.13TPA58\.3831\.5759\.3946\.8337\.0070\.0254\.0625\.5279\.9051\.41KPAh=\(4×4\)h\{=\}\(4\{\\times\}4\)64\.2330\.7259\.6644\.0934\.4070\.0854\.3826\.7985\.3052\.18KPAh=\(2×8\)h\{=\}\(2\{\\times\}8\)64\.1531\.4857\.2145\.9336\.2070\.1854\.9226\.3183\.9052\.25KPAh=\(8×2\)h\{=\}\(8\{\\times\}2\)66\.1633\.0260\.3745\.1135\.8071\.2953\.1226\.7984\.4052\.90Table 2:Zero\-shot evaluation on nine benchmarks \(FineWeb\-Edu,∼\\sim50B tokens\), 773M\-parameter models\. Best score per column inbold\. HellaSw\. = HellaSwag, W\.G\. = WinoGrande\.MethodARC\-EARC\-CBoolQHellaSw\.OBQAPIQAW\.G\.MMLUSciQAvg\.MHA59\.9333\.6261\.9350\.6336\.0071\.0655\.4122\.8781\.2052\.52MQA60\.7333\.6257\.3450\.0937\.0069\.9755\.4925\.3079\.6052\.13GQA61\.6634\.3058\.7249\.8538\.4071\.1653\.7525\.2377\.6052\.30MLA63\.5532\.8560\.9551\.7238\.8070\.5155\.0124\.5581\.9053\.32TPA63\.2235\.5860\.0351\.2636\.8071\.4455\.5624\.7779\.6053\.14KPAh=\(4×5\)h\{=\}\(4\{\\times\}5\)68\.5635\.9260\.0950\.2836\.8072\.8555\.3325\.2685\.9054\.55KPAh=\(2×10\)h\{=\}\(2\{\\times\}10\)68\.3134\.9855\.3250\.9138\.0071\.0655\.6424\.6685\.3053\.80KPAh=\(10×2\)h\{=\}\(10\{\\times\}2\)68\.8936\.0162\.0350\.3839\.4072\.8555\.9824\.5286\.0055\.12

## 4Conclusion

We introduced Kronecker Product Attention \(KPA\), a structured factorization of key and value projections that strictly generalizes TPA and integrates natively with RoPE without requiring a decoupled branch as in MLA\. KPA matches or outperforms TPA, MHA, MQA, GQA, and MLA across all evaluated scales under identical training conditions, while achieving substantial KV cache compression\. The Kronecker head split is a tuneable hyperparameter that allows practitioners to navigate the compression–accuracy trade\-off at design time without architectural changes\. By reducing the KV cache footprint, KPA directly addresses the memory\-bandwidth bottleneck that limits language models deployment on resource\-constrained devices such as smartphones, robotics platforms, and embedded systems, making it a natural fit for agentic systems that must plan and act under strict memory budgets without relying on cloud infrastructure\. The same compression benefits apply during fine\-tuning for tool use and multi\-step reasoning, where longer context windows are critical but memory is scarce\. Developing an efficient factored decoding kernel and evaluating KPA on longer contexts where cache compression is most practically impactful are natural next steps\.

## References

- Ainslie et al\. \[2023\]Joshua Ainslie, James Lee\-Thorp, Michiel de Jong, Yury Zelaski, Sumit Sanghai, and Yuhuai Fu\.Gqa: Training generalized multi\-query transformer models from multi\-head checkpoints\.In*Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing*, pages 4895–4901, 2023\.
- Chen et al\. \[2023\]Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian\.Extending context window of large language models via positional interpolation\.2023\.
- Choromanski et al\. \[2021\]Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller\.Rethinking attention with Performers\.In*International Conference on Learning Representations \(ICLR\)*, 2021\.
- Dao et al\. \[2022\]Tri Dao, Beidi Chen, Nimit S\. Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré\.Monarch: Expressive structured matrices for efficient and accurate training\.In*International Conference on Machine Learning \(ICML\)*, 2022\.
- DeepSeek\-AI \[2024\]DeepSeek\-AI\.DeepSeek\-V2: A strong, economical, and efficient mixture\-of\-experts language model\.2024\.
- Edalati et al\. \[2022a\]Ali Edalati, Marzieh Tahaei, Ivan Kobyzev, Vahid Partovi Nia, James J\. Clark, and Mehdi Rezagholizadeh\.KronA: Parameter efficient tuning with Kronecker adapter\.*arXiv preprint arXiv:2212\.10650*, 2022a\.
- Edalati et al\. \[2022b\]Ali Edalati, Marzieh Tahaei, Ahmad Rashid, Vahid Partovi Nia, James J Clark, and Mehdi Rezagholizadeh\.Kronecker decomposition for GPT compression\.In*ACL*, 2022b\.
- Gao et al\. \[2023\]Leo Gao, Jonathan Tow, Stella Biderman, et al\.A framework for few\-shot language model evaluation\.[https://github\.com/EleutherAI/lm\-evaluation\-harness](https://github.com/EleutherAI/lm-evaluation-harness), 2023\.
- Hooper et al\. \[2024\]Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W\. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami\.KVQuant: Towards 10 million context length LLM inference with KV cache quantization\.In*Advances in Neural Information Processing Systems \(NeurIPS\)*, 2024\.
- Karpathy \[2022\]Andrej Karpathy\.nanogpt\.[https://github\.com/karpathy/nanoGPT](https://github.com/karpathy/nanoGPT), 2022\.
- Kwon et al\. \[2023\]Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E\. Gonzalez, Hao Zhang, and Ion Stoica\.Efficient memory management for large language model serving with PagedAttention\.In*ACM Symposium on Operating Systems Principles \(SOSP\)*, 2023\.
- Liu et al\. \[2024\]Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu\.KIVI: A tuning\-free asymmetric 2bit quantization for KV cache\.In*International Conference on Machine Learning \(ICML\)*, 2024\.
- Michel et al\. \[2019\]Paul Michel, Omer Levy, and Graham Neubig\.Are sixteen heads really better than one?In*Advances in Neural Information Processing Systems*, volume 32, 2019\.
- Novikov et al\. \[2015\]Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov\.Tensorizing neural networks\.In*Advances in Neural Information Processing Systems \(NeurIPS\)*, 2015\.
- Omranpour et al\. \[2024\]Soroush Omranpour, Guillaume Rabusseau, and Reihaneh Rabbany\.Higher order transformers: Efficient attention mechanism for tensor structured data\.2024\.
- Penedo et al\. \[2024\]Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf\.The fineweb datasets: Decanting the web for the finest text data at scale\.*arXiv preprint arXiv:2406\.17557*, 2024\.
- Peng et al\. \[2024\]Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole\.YaRN: Efficient context window extension of large language models\.In*International Conference on Learning Representations \(ICLR\)*, 2024\.
- Press et al\. \[2022\]Ofir Press, Noah A\. Smith, and Mike Lewis\.Train short, test long: Attention with linear biases enables input length extrapolation\.In*International Conference on Learning Representations \(ICLR\)*, 2022\.
- Shazeer \[2019\]Noam Shazeer\.Fast transformer decoding: One write\-head is all you need\.2019\.
- Shazeer \[2020\]Noam Shazeer\.Glu variants improve transformer\.*arXiv preprint arXiv:2002\.05202*, 2020\.
- Su et al\. \[2024\]Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu\.Roformer: Enhanced transformer with rotary position embedding\.*Neurocomputing*, 568:127063, 2024\.
- Sun et al\. \[2023\]Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei\.A length\-extrapolatable transformer\.In*Annual Meeting of the Association for Computational Linguistics \(ACL\)*, 2023\.
- Tahaei et al\. \[2022\]Marzieh Tahaei, Ella Charlaix, Vahid Partovi Nia, Ali Ghodsi, and Mehdi Rezagholizadeh\.KroneckerBERT: Significant compression of pre\-trained language models through Kronecker decomposition and knowledge distillation\.In*Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies \(NAACL\)*, pages 2116–2127, 2022\.
- Van Loan and Pitsianis \[1993\]Charles F\. Van Loan and Nikos Pitsianis\.Approximation with kronecker products\.In*Linear Algebra for Large Scale and Real\-Time Applications*, pages 293–314\. Springer, 1993\.
- Voita et al\. \[2019\]Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov\.Analyzing multi\-head self\-attention: Specialized heads do the heavy lifting, the rest can be pruned\.In*Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics*, pages 5797–5808, 2019\.
- Wang et al\. \[2020\]Sinong Wang, Belinda Z\. Li, Madian Khabsa, Han Fang, and Hao Ma\.Linformer: Self\-attention with linear complexity\.2020\.
- Xiao et al\. \[2023\]Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis\.Efficient streaming language models with attention sinks\.*arXiv preprint arXiv:2309\.17453*, 2023\.
- Xiao et al\. \[2024\]Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis\.Efficient streaming language models with attention sinks\.In*International Conference on Learning Representations \(ICLR\)*, 2024\.
- Zhang et al\. \[2025\]Yifan Zhang, Yifeng Liu, et al\.Tensor product attention is all you need\.In*Advances in Neural Information Processing Systems*, 2025\.Spotlight\.
- Zhang et al\. \[2023\]Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen\.H2O: Heavy\-hitter oracle for efficient generative inference of large language models\.In*Advances in Neural Information Processing Systems \(NeurIPS\)*, 2023\.

## Appendix AProofs

### A\.1Proof of Proposition[2\.2](https://arxiv.org/html/2610.12657#S2.Thmtheorem2)

###### Proof\.

For head\(i1,i2\)\(i\_\{1\},i\_\{2\}\), the per\-head query and key vectors are

𝒒t,\(i1,i2\)=1RQ​∑r\(𝑨rQ\)i1⊗\(𝑩rQ\)i2,𝒌s,\(i1,i2\)=1RK​∑r′\(𝑨r′K\)i1⊗\(𝑩r′K\)i2\.\\bm\{q\}\_\{t,\(i\_\{1\},i\_\{2\}\)\}=\\frac\{1\}\{R\_\{Q\}\}\\sum\_\{r\}\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\otimes\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\},\\qquad\\bm\{k\}\_\{s,\(i\_\{1\},i\_\{2\}\)\}=\\frac\{1\}\{R\_\{K\}\}\\sum\_\{r^\{\\prime\}\}\(\\bm\{A\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{1\}\}\\otimes\(\\bm\{B\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{2\}\}\.
Taking the dot product and applying the mixed\-product identity

\(𝒂⊗𝒃\)⊤​\(𝒄⊗𝒅\)=\(𝒂⊤​𝒄\)​\(𝒃⊤​𝒅\),\(\\bm\{a\}\\otimes\\bm\{b\}\)^\{\\top\}\(\\bm\{c\}\\otimes\\bm\{d\}\)=\(\\bm\{a\}^\{\\top\}\\bm\{c\}\)\(\\bm\{b\}^\{\\top\}\\bm\{d\}\),we obtain

𝒒t,\(i1,i2\)⊤​𝒌s,\(i1,i2\)=1RQ​RK​∑r,r′\[\(𝑨rQ\)i1⊤​\(𝑨r′K\)i1\]​\[\(𝑩rQ\)i2⊤​\(𝑩r′K\)i2\]\.∎\\bm\{q\}\_\{t,\(i\_\{1\},i\_\{2\}\)\}^\{\\top\}\\bm\{k\}\_\{s,\(i\_\{1\},i\_\{2\}\)\}=\\frac\{1\}\{R\_\{Q\}R\_\{K\}\}\\sum\_\{r,r^\{\\prime\}\}\\bigl\[\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}^\{\\top\}\(\\bm\{A\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{1\}\}\\bigr\]\\bigl\[\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}^\{\\top\}\(\\bm\{B\}^\{K\}\_\{r^\{\\prime\}\}\)\_\{i\_\{2\}\}\\bigr\]\.\\qed

### A\.2Proof of Theorem[2\.4](https://arxiv.org/html/2610.12657#S2.Thmtheorem4)

###### Proof\.

Applying𝑰d1⊗𝑹θ\(2\)​\(t\)\\bm\{I\}\_\{d\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)to\(𝑨rQ\)i1⊗\(𝑩rQ\)i2\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\otimes\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}and using the mixed\-product identity gives

\(𝑰d1⊗𝑹θ\(2\)​\(t\)\)​\(\(𝑨rQ\)i1⊗\(𝑩rQ\)i2\)=\(𝑨rQ\)i1⊗𝑹θ\(2\)​\(t\)​\(𝑩rQ\)i2\.\(\\bm\{I\}\_\{d\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)\)\\bigl\(\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\otimes\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}\\bigr\)=\(\\bm\{A\}^\{Q\}\_\{r\}\)\_\{i\_\{1\}\}\\otimes\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)\(\\bm\{B\}^\{Q\}\_\{r\}\)\_\{i\_\{2\}\}\.
Summing overrrand averaging by1/RQ1/R\_\{Q\}preserves the factored form\.

The relative\-position identity then follows from

𝑹θ\(2\)​\(t\)⊤​𝑹θ\(2\)​\(s\)=𝑹θ\(2\)​\(s−t\),\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(t\)^\{\\top\}\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(s\)=\\bm\{R\}^\{\(2\)\}\_\{\\theta\}\(s\-t\),which is the standard RoPE relative\-position property restricted to thed2d\_\{2\}\-dimensional subspace\. ∎

## Appendix BKV Cache Compression

During autoregressive decoding, KPA caches the factor matrices per past token at a cost of\(RK\+RV\)​\(h1​d1\+h2​d2\)\(R\_\{K\}\+R\_\{V\}\)\(h\_\{1\}d\_\{1\}\+h\_\{2\}d\_\{2\}\)per token, compared with2​h1​h2​d1​d22h\_\{1\}h\_\{2\}d\_\{1\}d\_\{2\}for MHA\. The compression factor is

CKPA=2​h1​h2​d1​d2\(RK\+RV\)​\(h1​d1\+h2​d2\)\.C\_\{\\mathrm\{KPA\}\}=\\frac\{2\\,h\_\{1\}h\_\{2\}d\_\{1\}d\_\{2\}\}\{\(R\_\{K\}\+R\_\{V\}\)\(h\_\{1\}d\_\{1\}\+h\_\{2\}d\_\{2\}\)\}\.\(11\)
Table[3](https://arxiv.org/html/2610.12657#A2.T3)reportsCKPAC\_\{\\mathrm\{KPA\}\}for all evaluated configurations withRK=RV=2R\_\{K\}=R\_\{V\}=2andd1=d2=8d\_\{1\}=d\_\{2\}=8\.

Table 3:KV cache compression factors for evaluated KPA configurations\.
## Appendix CTraining and Evaluation Details

#### Architecture and data\.

All models use the nanoGPT codebase\[[Karpathy, 2022](https://arxiv.org/html/2610.12657#bib.bib10)\]with SwiGLU activations\[[Shazeer, 2020](https://arxiv.org/html/2610.12657#bib.bib20)\]and RoPE embeddings\[[Su et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib21)\]\. We fixd1=d2=8d\_\{1\}=d\_\{2\}=8and varyh1,h2h\_\{1\},h\_\{2\}across configurations to study head\-split asymmetry\. Models are trained from scratch on FineWeb\-Edu\-100B\[[Penedo et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib16)\]for 50,000 iterations with a batch size yielding approximately 1M tokens per iteration, totalling∼\\sim50B tokens\.

#### Optimisation\.

We use AdamW with peak learning rate6×10−46\\times 10^\{\-4\}\(3×10−43\\times 10^\{\-4\}at 773M and 1\.5B scales\), cosine decay to one\-tenth of the peak, 2,000 warmup steps, and weight decay 0\.1\. All models train in bfloat16 on NVIDIA H100 GPUs\.

#### Evaluation\.

We evaluate zero\-shot performance on nine benchmarks usinglm\-evaluation\-harness\[[Gao et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib8)\]: ARC\-Easy, ARC\-Challenge, BoolQ, HellaSwag, OpenbookQA, PIQA, WinoGrande, MMLU, and SciQ\. We report accuracy\-normalized scores for ARC\-E, ARC\-C, HellaSwag, OBQA, PIQA, and SciQ, and standard accuracy for BoolQ, WinoGrande, and MMLU\. The primary scalar metric is the unweighted average across all nine tasks\.

## Appendix DAdditional Experimental Results

### D\.11\.5B Scale

Sinceh=25h\{=\}25admits only the symmetric Kronecker splith1=h2=5h\_\{1\}\{=\}h\_\{2\}\{=\}5, a single KPA configuration is evaluated at the 1\.5B scale\. KPA outperforms TPA and all other baselines on eight of nine benchmarks, with the only exception being MMLU where MHA achieves a slightly higher score\. The gains are most pronounced on ARC\-Easy, SciQ, and BoolQ, consistent with the pattern observed at 353M and 773M scales\. This confirms that KPA scales favourably with model size even when the head split is constrained to a single symmetric factorization\.

### D\.2125M Scale

At 125M scale, all KPA configurations outperform TPA \(46\.21%\)\. Consistent with the pattern at larger scales, asymmetric splits with largerh1h\_\{1\}perform better:h=\(6×2\)h\{=\}\(6\{\\times\}2\)achieves the best average, followed byh=\(4×3\)h\{=\}\(4\{\\times\}3\),h=\(3×4\)h\{=\}\(3\{\\times\}4\), andh=\(2×6\)h\{=\}\(2\{\\times\}6\)\. The gains over TPA are smaller than at 353M and 773M, which is expected as smaller models have less capacity to benefit from richer factorization\. Nevertheless, KPA consistently outperforms all baselines including MHA, MQA, GQA, MLA, and TPA across all four configurations\.

Table 4:Zero\-shot evaluation of 1\.5B\-parameter models \(FineWeb\-Edu,∼\\sim50B tokens\)\. Baselines from[Zhang et al\. \[2025\]](https://arxiv.org/html/2610.12657#bib.bib29)\. Best score in each column isbold\.Table 5:Zero\-shot evaluation of 125M\-parameter models \(FineWeb\-Edu,∼\\sim50B tokens\)\. Baselines from[Zhang et al\. \[2025\]](https://arxiv.org/html/2610.12657#bib.bib29)\. Best score in each column isbold\.

## Appendix ERelated Work

### E\.1Head sharing

MQA\[[Shazeer, 2019](https://arxiv.org/html/2610.12657#bib.bib19)\]reduces the KV cache by sharing a single key and value head across all query heads, and GQA\[[Ainslie et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib1)\]generalizes this by grouping query heads to share a smaller set of key\-value heads, recovering a spectrum between MQA and full MHA\. MLA\[[DeepSeek\-AI, 2024](https://arxiv.org/html/2610.12657#bib.bib5)\]instead projects keys and values into a low\-dimensional latent space before caching, achieving stronger compression, but requires a decoupled RoPE branch to handle positional encoding\. KPA avoids head sharing entirely and eliminates the decoupled RoPE branch by exploiting the Kronecker structure of its projections\.

### E\.2Low\-rank and structured attention

Linformer\[[Wang et al\., 2020](https://arxiv.org/html/2610.12657#bib.bib26)\]approximates the attention matrix via low\-rank projection of keys and values along the sequence dimension, while Performer\[[Choromanski et al\., 2021](https://arxiv.org/html/2610.12657#bib.bib3)\]approximates the softmax kernel with random features\. These methods target the quadratic complexity of attention rather than KV cache size and do not apply to autoregressive decoding\. Monarch\[[Dao et al\., 2022](https://arxiv.org/html/2610.12657#bib.bib4)\]and Tensorizing\-NN\[[Novikov et al\., 2015](https://arxiv.org/html/2610.12657#bib.bib14)\]apply structured matrix factorizations to weight matrices for parameter efficiency, but do not specifically target KV cache compression\. TPA\[[Zhang et al\., 2025](https://arxiv.org/html/2610.12657#bib.bib29)\]is the closest prior work, applying CP decomposition to KV projections; KPA strictly generalizes TPA by replacing rank\-1 vector outer products with Kronecker matrix factors\.

### E\.3KV compression, eviction, and quantization

H2O\[[Zhang et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib30)\]and StreamingLLM\[[Xiao et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib28)\]reduce cache size by evicting tokens deemed less important based on attention scores, but risk discarding tokens that become relevant later\. KIVI\[[Liu et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib12)\]and KVQuant\[[Hooper et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib9)\]compress the cache through low\-bit quantization, trading precision for memory\. PagedAttention\[[Kwon et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib11)\]improves memory management through virtual paging but does not reduce the total cache footprint\. KPA takes a complementary architectural approach, reducing the cache at the representation level before any eviction or quantization is applied\.

### E\.4Positional encoding

RoPE\[[Su et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib21)\]encodes relative position by rotating query and key vectors and has become the dominant positional scheme in modern language models\. ALiBi\[[Press et al\., 2022](https://arxiv.org/html/2610.12657#bib.bib18)\]instead adds position\-dependent biases to attention logits, avoiding explicit positional embeddings\. Extensions such as YaRN\[[Peng et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib17)\], xPos\[[Sun et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib22)\], and PI\[[Chen et al\., 2023](https://arxiv.org/html/2610.12657#bib.bib2)\]adapt RoPE for longer contexts through frequency scaling or interpolation\. KPA introduces Kronecker\-Separable RoPE, which applies rotary embeddings directly to the factorized representation by rotating only the feature\-level factor, preserving the standard relative\-position property without materializing full key tensors or requiring a decoupled positional branch\.

### E\.5Kronecker factorization in transformers\.

Kronecker decomposition has been applied to transformers in several distinct settings, all targeting static weight compression rather than inference\-time KV cache reduction\. KroneckerBERT\[[Tahaei et al\., 2022](https://arxiv.org/html/2610.12657#bib.bib23)\]compresses BERT weight matrices via Kronecker decomposition with knowledge distillation\. KnGPT2\[[Edalati et al\., 2022b](https://arxiv.org/html/2610.12657#bib.bib7)\]applies the same idea to GPT\-2, replacing linear layer weights with Kronecker\-structured matrices and using intermediate layer distillation to recover accuracy\. KronA\[[Edalati et al\., 2022a](https://arxiv.org/html/2610.12657#bib.bib6)\]uses Kronecker products as a drop\-in replacement for LoRA adapters in parameter\-efficient fine\-tuning, achieving a better expressivity\-to\-parameter trade\-off than low\-rank decomposition\. HOT\[[Omranpour et al\., 2024](https://arxiv.org/html/2610.12657#bib.bib15)\]decomposes the attention score matrix as a sum of Kronecker products to reduce training\-time computation for multi\-dimensional encoder inputs\. KPA differs from all of these: it factorizes dynamic, per\-token key and value activations rather than static weight matrices or score matrices, targeting inference\-time memory reduction in autoregressive decoding, a setting none of the above address\.

Similar Articles

A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention

arXiv cs.LG

IBM Research and UIUC researchers propose Universal Attention, a self-pruning transformer framework combining first- and second-order key statistics into expressive forget gates that preserve RoPE and Softmax attention while achieving state-of-the-art 10x KV-cache compression and 25x compression at 16k context length.

Thought-Aware KV Cache Compaction for Reasoning via Adaptive Attention Matching

arXiv cs.CL

This paper introduces Thought-Aware Attention Matching (TAM), a method for compressing KV caches during long chain-of-thought reasoning by segmenting reasoning traces, adaptively allocating compression budgets, and protecting pivotal tokens. Experiments on AIME 2024 and MATH-500 with Qwen3-4B show significant memory reduction while maintaining accuracy.

KV Cache Compression Through the Lens of Transform Coding

arXiv cs.LG

The paper proposes Attention-Aware Transform Coding (AATC) for compressing KV caches in large language models, achieving near-lossless accuracy at around 5.8x compression by minimizing distortion through attention mechanisms.