arXiv Machine Learning

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

arXiv Machine Learning
Sep 25

A JoLT for the KV cache: Near-Lossless KV Cache Compression via Joint Rank-bit Allocation

The paper introduces JoLT, a training‑free compressor that jointly allocates rank and precision for key‑value (KV) cache compression in long‑context language models. JoLT treats grouped prefill caches as fourth‑order tensors, applies partial Tucker decomposition along token and feature modes, and uses a rotated low‑bit quantizer for residuals, all governed by a single Lagrangian dual under a global byte constraint. Across five models from four architecture families, JoLT achieves 2–3× compression with less than 0.2% perplexity loss, and near‑lossless retrieval accuracy on LLaMA‑3.1‑8B at 64K context up to 3× compression.

By Rahul Krishnan, Volker Schulz
arXiv Machine Learning
Oct 8

Activation-Aware Weight Tensorization: A Calibration-Time Preconditioner for Tensor-Network LLM Compression

Activation-Aware Weight Tensorization (AWT) is a training‑free calibration wrapper that preconditions each weight matrix with a diagonal activation‑derived scale before applying standard Tensor Train (TT) or Tree Tensor Network (TTN) decompositions. By adding only an input‑side elementwise rescaling, AWT improves functional fidelity of post‑training tensor‑network compression across several large language models (Llama 3.1 8B, Ministral 8B, Qwen2.5 7B) at 2–6× compression, closing a significant portion of the perplexity gap to dense baselines and enhancing downstream task performance. The method demonstrates that diagonal preconditioning offers a robustness‑modularity trade‑off, outperforming full‑covariance oracles in many cases while remaining modular and solver‑agnostic.

By Alessandro Beatini, Marco Maronese, Emanuele Rodol\`a
arXiv AI
Oct 12

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

Kronecker Product Attention (KPA) is introduced as a memory‑efficient method for compressing the key‑value cache in autoregressive models. By representing key and value activations as sums of Kronecker products of compact, token‑dependent matrix factors, KPA generalizes Tensor Product Attention and can capture richer block‑structured patterns. Empirical results show that KPA matches or surpasses existing techniques such as MHA, MQA, GQA, and MLA across multiple model scales and benchmarks, achieving up to 10× compression while preserving relative‑position information through a structured rotary embedding.

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