PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Bahareh Nikpour, Hossein Rajabzadeh, Maryam Dialameh, Vahid Partovi Nia, Parsa Omidi

arXiv:2610.12657v1cs.LGcs.AI

TL;DR

The KV cache grows with sequence length and batch size, making inference-time memory a bottleneck. KPA represents projections with compact Kronecker factors and adapts RoPE to that structure. It matches or exceeds several attention baselines across scales and benchmarks while compressing the cache.

  • Problem

    The KV cache grows with sequence length and batch size, potentially dominating GPU memory and limiting context length, batch size, and deployment scale.

    Source: §1 · p. 1

  • Method

    KPA represents per-token projections with compact matrix factors combined by Kronecker products and applies a matched rotary operator to the factors.

    Source: §1 · pp. 2–3

  • Results

    KPA matches or exceeds TPA, MHA, MQA, GQA, and MLA across nine benchmarks and four scales, with 6.4×–10× KV cache compression.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The tuneable Kronecker head split lets practitioners navigate the compression–accuracy trade-off at design time without architectural changes.

    Source: §4 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

The KV cache can dominate inference memory as context length and batch size grow. KPA represents per-token query, key, and value projections with compact Kronecker factors and pairs this structure with a RoPE operator; experiments report competitive benchmark results with substantial cache compression.

1 Introduction

  • The KV cache stores projected keys and values across layers and tokens, and its growth with sequence length and batch size can constrain context length and deployment scale.The passage identifies inference-time memory, rather than parameter count, as an increasing bottleneck.Source: §1 · p. 1
  • Existing compression approaches trade off representation or decoding machinery: MQA and GQA share keys and values across heads, MLA uses a decoupled RoPE branch, and TPA uses vector-based low-rank factors.The introduction motivates a factorization that can capture structure at both the head and within-head feature levels.Source: §1 · p. 2
  • KPA represents per-token key and value matrices using smaller contextual matrix factors, making their Kronecker products block-structured across head and feature dimensions.Unlike vector-based CP factorization, the Kronecker product can jointly encode structure at both levels; rank-one outer products are a special case.Source: §1 · p. 2
  • Kronecker-Separable RoPE applies rotary position encoding to compact key factors while preserving RoPE’s relative-position property, without an MLA-style decoupled branch.The introduction describes this as applying rotary encoding directly to compact key factors.Source: §1 · p. 3
  • KPA matches or exceeds TPA, MHA, MQA, GQA, and MLA across nine benchmarks at four scales, with 6.4×–10× KV cache compression.The models were trained under identical conditions.Source: §1 · p. 3

2 Kronecker Product Attention

KPA represents token-dependent key and value matrices as sums of Kronecker products of compact factors, preserving head-level and within-head structure while caching only those factors. Its factorization supports factored attention scores and a RoPE-compatible cache, with factor dimensions controlling compression.

2 Kronecker Product Attention

  • KPA factorizes each token’s key and value matrices into sums of Kronecker products of smaller factors for head-level and within-head structure, and caches only those factors.This structure aligns with multi-head attention’s functional organization, unlike vector-based CP decomposition.Source: §2 · p. 3

2.1 Background and Notation

  • The notation defines ⊗ as the Kronecker product and identifies the mixed-product identity as the basis for factored score computation and the RoPE compatibility proof.The passage attributes the identity to Van Loan and Pitsianis (1993).Source: §2.1 · p. 3

2.2 Factorised Key, Value, and Query Representations

  • The query, key, and value slices are normalized sums of contextual Kronecker products, with independent ranks; queries can use larger rank because only key and value factors enter the cache.Both factors depend on the token representation through independent learned linear projections.Source: §2.2 · p. 4
  • A degenerate split makes each Kronecker term a rank-one outer product, exactly recovering TPA [Zhang et al., 2025], while non-degenerate factors have greater representational capacity at matched parameter budgets.The rank of a Kronecker term is the product of the ranks of its two factors.Source: §2.2 · p. 4

2.3 Factored Attention Score Computation

  • Each attention score factors into head-level and feature-level inner products, allowing computation directly from compact factors without materializing the full key tensor.The score is scaled by 1/√dh before softmax in practice.Source: §2.3 · p. 4

2.4 Kronecker-Separable RoPE

  • Kronecker-Separable RoPE rotates only the B factors, preserves dependence on relative offset s − t, and requires no decoupled RoPE branch.The A factors are position-independent and stored as-is; the B factors are pre-rotated before caching.Source: §2.4 · p. 5

2.5 KV Cache Cost and Compression

  • Per-token cache cost is (RK + RV)(h1d1 + h2d2), versus 2h1h2d1d2 for MHA; for fixed h and dh, more asymmetric factor splits increase compression.The factors h1, h2, d1, and d2 control the two-level structure.Source: §2.2 · p. 4, §2.5 · p. 5

3 Experiments

The experiments compare KPA with attention baselines across four model scales and nine benchmarks under identical training conditions, then examine how head-split asymmetry affects performance and compression.

3.1 Experimental Setup

  • Models at 125M, 353M, 773M, and 1.5B were trained on FineWeb-Edu for ∼50B tokens and evaluated zero-shot on nine benchmarks against MHA, MQA, GQA, MLA, and TPA.The baselines were trained under identical conditions.Source: §3.1 · p. 5

3.2 Main Results

  • At 353M, KPA h=(8×2) exceeds TPA by 1.49pp at 6.4× compression; at 773M, h=(10×2) outperforms TPA, MLA, and other baselines at 6.7× compression.KPA also improves monotonically across scales and matches or exceeds TPA at comparable or higher compression ratios at every scale.Source: §3.2 · pp. 5–6, Table 1 · p. 6, Table 2 · p. 6

3.3 Ablation on Head-Split Asymmetry

  • With d1=d2=8 and RK=RV=2 fixed, asymmetric splits with larger h1 outperform symmetric KPA and TPA at both scales.The authors attribute this to the A-factor encoding joint interactions across h1 heads; symmetric splits offer higher compression (8× to 8.9×) when memory is the primary constraint.Source: §3.3 · p. 6

4 Conclusion

KPA factorizes key and value projections, integrates with RoPE without an MLA-style decoupled branch, and matches or outperforms the evaluated attention baselines while compressing the KV cache.

  • KPA strictly generalizes TPA and integrates natively with RoPE without requiring a decoupled branch as in MLA.The conclusion also identifies the head split as a tunable compression–accuracy trade-off and names efficient decoding kernels and longer-context evaluation as next steps.Source: §4 · p. 6

A.1 Proof of Proposition 2.2

The proof introduces per-head query and key vectors, then uses their dot product and the mixed-product identity.

  • The proof begins with the per-head query and key vectors for head (i1, i2).Source: §A.1 · p. 9
  • The next proof step takes the dot product of those vectors.Source: §A.1 · p. 9
  • The dot-product step applies the mixed-product identity.Source: §A.1 · p. 9

A.2 Proof of Theorem 2.4

The proof preserves a factored form under rank summation and establishes the relative-position identity in a restricted feature subspace; the appendix also reports per-token cache costs and compression factors.

A.2 Proof of Theorem 2.4

  • Summing over r and averaging by 1/RQ preserves the factored form, and the relative-position identity holds in the d2-dimensional subspace.Source: §A.2 · p. 9

B KV Cache Compression

C Training and Evaluation Details

The experiments use a shared training setup and evaluate zero-shot performance on nine benchmarks, with an unweighted average as the primary scalar metric.

Architecture and data.

  • Models fix d1 = d2 = 8, vary h1 and h2, and train from scratch on FineWeb-Edu-100B for 50,000 iterations, totaling about 50B tokens.Each iteration processes approximately 1M tokens.Source: C Training and Evaluation Details · p. 9

Optimisation.

  • Training uses AdamW, cosine learning-rate decay, 2,000 warmup steps, weight decay 0.1, and bfloat16 on NVIDIA H100 GPUs.Peak learning rate is 6 × 10−4, reduced to 3 × 10−4 at 773M and 1.5B scales.Source: C Training and Evaluation Details · p. 10

Evaluation.

  • Zero-shot evaluation covers nine benchmarks; the primary scalar metric is their unweighted average.The suite includes ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, OpenbookQA, PIQA, WinoGrande, MMLU, and SciQ.Source: C Training and Evaluation Details · p. 10

D.1 1.5B Scale

At 1.5B parameters, KPA performs favorably against the baselines despite the head-count constraint allowing only one symmetric split.

D.1 1.5B Scale

  • Eight of nine benchmarks: KPA outperforms TPA and all other baselines, while MHA scores slightly higher on MMLU.The 1.5B configuration is evaluated under the sole symmetric split h1 = h2 = 5 allowed by h = 25.Source: §D.1 · p. 10
  • Gains are most pronounced on ARC-Easy, SciQ, and BoolQ, consistent with the pattern at 353M and 773M scales.Source: §D.1 · p. 10
  • The single symmetric factorization at 1.5B still supports favorable scaling against the baselines.Source: §D.1 · p. 10

D.2 125M Scale

At 125M, all KPA configurations outperform the baselines, with the strongest average from an asymmetric split; the related-work comparisons distinguish KPA's representation-level approach from prior methods.

D.2 125M Scale

  • At 125M, every KPA configuration beats TPA (46.21%) and all baselines; h=(6×2) leads, while gains are smaller than at 353M and 773M.Source: Table 5 · p. 10, §D.2 · p. 10

E.1 Head sharing

E.2 Low-rank and structured attention

E.3 KV compression, eviction, and quantization

  • KPA compresses the cache at the representation level, complementing eviction methods such as H2O and quantization methods such as KIVI [Liu et al., 2024].Source: §E.3 · p. 11

E.4 Positional encoding

  • Kronecker-Separable RoPE rotates only the feature-level factor, preserving relative position without materializing full key tensors or requiring a decoupled positional branch.Source: §E.4 · p. 11

E.5 Kronecker factorization in transformers.

  • Unlike prior Kronecker transformer methods targeting static weights or score matrices, KPA factorizes dynamic per-token key and value activations for inference-time memory reduction.Source: §E.5 · p. 11

Related papers

Loading 2610.12657v1…