KV-Compression Aware Training (KV-CAT)

• KV-CAT is a training framework that incentivizes compressible key-value caches in transformers to mitigate escalating storage and latency costs during long-context processing.
• It incorporates methods like controlled pretraining, learnable orthogonal projections (MatryoshkaKV), and gating mechanisms (Fast KVzip) to achieve over 90% baseline accuracy at high compression rates.
• KV-CAT is applicable to long-context language modeling and edge deployments, offering resource-efficient transformer adaptations for memory-constrained environments.

KV-Compression Aware Training (KV-CAT) refers to a set of methods and training objectives that directly incentivize, during model adaptation, the production of key-value (KV) cache representations in transformers that are highly amenable to downstream compression. As context length increases in LLMs, the linear scaling of storage and latency for the KV cache presents a critical bottleneck. KV-CAT procedures increase the effective compressibility of internal states and/or allow the imposition of compression constraints (eviction or projection) with minimal task degradation, even at aggressive memory budgets. The chief strategies in KV-CAT include (1) controlled modification of the representation space via continual pretraining or side objectives, (2) explicit learnable compressive projections, (3) trainable routing/prediction modules for cache retention and (4) integration with adaptation techniques such as LoRA. These methods are distinct from post-hoc compression, as compressibility becomes an explicit consideration during training.

1. Motivation and Problem Formulation

The memory and compute cost of the KV cache in transformer-based LLMs grows linearly with sequence length and quadratic with feature/layer dimensionality. As LLMs are scaled for practical long-context usage—retrieval, QA, variable tracking, etc.—memory overhead from context preservation becomes a severe obstacle. Classical compression methods—PCA, chunking, head-selection, and token/page-level summarization—operate as post-hoc steps, yielding mixed results depending on the compressibility of the frozen model's internal representations. Theoretical work shows that compressibility is a property of the learned sequence-to-vector mapping, not the input data alone; given the same task, there exist both compressible and non-compressible transformer solutions depending on the training regime (Gelberg et al., 7 May 2026). Thus, training protocols that bias models toward compressible (by deletion or projection) KV representations are necessary for memory-efficient LLMs.

2. Low-Rank Orthogonal Projection with MatryoshkaKV

MatryoshkaKV implements KV-Compression Aware Training by inserting per-layer, per-head learnable orthogonal projections into the cache feature dimension. Each Transformer attention head is augmented with a matrix $U^{(\ell, h)}\in\mathbb{R}^{d\times d}$ with orthogonality enforced via Cayley parameterization. Inference-time compression is achieved by retaining only the first $r$ columns ($U_r$), projecting cached $K,V$ via $K' = K \cdot U_r,\,V' = V \cdot U_r$, and reconstructing the attention output by post-multiplying with $U_r^\top$. The learning objective combines an $L_2$ loss between full-dimension and projected attention outputs with a cross-entropy token modeling loss: $$L = \mathbb{E}_x \| \mathrm{Attn}_\text{full}(x) - \mathrm{Attn}_\text{comp}(x; \{U\}) \|_2^2 + \lambda L_\text{LM}$$ A “Matryoshka” (nested rate) training scheme ensures that the first $r$ columns support optimal subspace choices for all allowed $r$ by sampling compression ranks at each update. Post-training dynamic programming or greedy search allocates per-head rates under a fixed global cache budget.

Empirical evaluations on LLaMA2-7B and Mistral-7B show that MatryoshkaKV retains over 90% base accuracy at 60% compression and over 75% at 37.5% compression, outperforming PCA by a large margin at all compression levels. Continued pretraining (CPT) and LoRA-based supervised fine-tuning (SFT) are both directly supported (Lin et al., 2024).

Method	50% Budget Accuracy (LLaMA2-7B)	37.5% Budget Accuracy
PCA	43.07%	33.14%
MatryoshkaKV	59.11%	56.94%

3. Gated Cache Eviction: Fast KVzip

Fast KVzip introduces lightweight sink-attention gating networks, $r$0, per layer to predict the relative importance of each KV slot for selective retention. The gating mechanism computes grouped-query soft attention over hidden states and learnable “sink” vectors, outputting per-head scalar retention scores. Gates are trained post-hoc using precomputed hidden states from the frozen LLM, optimizing a binary cross-entropy objective matching task-agnostic reconstruction targets. No gradients are propagated through the LLM proper, enabling highly efficient training (sub-1 hour on an H100 for Qwen3-8B) (Kim et al., 25 Jan 2026).

During inference, gates select the KV slots to retain under both global budget constraints and required local windows (recency), supporting both prefill and autoregressive decoding scenarios. Empirical results across Qwen and Gemma3 series show near-lossless accuracy while evicting up to 70% of the cache. The Fast KVzip method incurs only ≈1% additional latency relative to baseline, and its gate modules require minimal storage (<0.2 GB for 8B models).

4. Sub-Token Routing and Fine-Grained Query-Aware Compression

KV-Compression Aware Training with sub-token routing partitions the $r$1 vector within each attention head into $r$2 value groups (subspaces), with lightweight routers assigning retention scores to each group. For query-independent training, retention/reconstruction policies are optimized alongside adaptation (e.g., LoRA) to minimize perplexity and auxiliary value reconstruction loss. For query-aware settings, a small predictor MLP allocates a global KV retention budget over pairs of context-token/value-group conditioned on anticipated query relevance. Training uses diagnostic attention weights to label group importance and optimizes the predictor for MSE on these targets (Jiang et al., 23 Apr 2026).

Sub-token routing enables a complementary axis to token-level methods: token-level selection chooses which tokens to preserve, while sub-token routing determines the granularity of internal values per token. Empirical results demonstrate that on Qwen2.5-7B, token-level and sub-token-level routing jointly allow >99.6% MMLU accuracy at 37.5% total KV retention, with independent axes yielding 99.8% (token-level) and 96.3% (sub-token) respectively.

5. Optimization Objectives and Training Protocols

Across implementations, the core of KV-CAT involves augmenting the primary LM objective with auxiliary losses that reward output invariance under deletion or projection of the KV cache. In MatryoshkaKV, the combined loss balances distillation (attention output matching) and language modeling, typically with a 1:3 weighting. Post-hoc gate training, as in Fast KVzip, minimizes cross-entropy with respect to retention/eviction labels derived from model reconstructions. Sub-token routing includes an auxiliary group-reconstruction loss and, in query-aware variants, a group-wise relevance predictor trained to MSE.

Training is typically conducted as continual pretraining on a subset of general text (RedPajama, WikiText-103, or FineWeb-Edu), with SFT protocols for downstream adaptation. Compression rates are not fixed during training, enabling post-hoc budget allocation via search strategies.

6. Empirical Performance and Impact

KV-CAT techniques consistently show that direct optimization for compressibility yields substantial robustness to aggressive cache reduction. MatryoshkaKV preserves 97% of zero-shot benchmark accuracy at 50% KV and 93% at 37.5%. Fast KVzip achieves up to 70% eviction (R=0.3) with <1% drop in QA, retrieval, code generation, and mathematical reasoning metrics. Sub-token routing combined with token-level selection results in near-lossless performance at 37.5% retained KV on language understanding and variable-tracking tasks. Overhead in computation is minimal (≈1% increase in latency or a few ms per inference step), and storage partitioning for gates and projections is highly efficient.

7. Applications and Future Directions

KV-Compression Aware Training has applications in long-context language modeling, high-throughput inference, memory- and latency-sensitive deployments (e.g., retrieval-augmented generation, code completion, QA), and settings requiring dynamic cache adaptation to varying input budgets. A plausible implication is the feasibility of deploying large models on-edge and in resource-constrained environments without severe quality degradation. Directions enabled by KV-CAT include: multi-axis compression combining token, head/layer, and value group selection; task-conditioned or adaptive strategies informed by the query; and integration into dense/sparse model scaling paradigms.

KV-CAT represents an emergent best practice for designing and adapting transformers where cache efficiency is mission-critical, unifying advances in adaptive projections, gating/eviction, and representational training into a flexible and empirically validated framework (Lin et al., 2024, Kim et al., 25 Jan 2026, Jiang et al., 23 Apr 2026).