TKV: Online KV Cache Compression via Warmup-Then-Freeze Brand Updates and a Random Sketch Bridge
Abstract
The key-value (KV) cache is the dominant memory bottleneck in long-context autoregressive inference. Existing online compression methods face a tradeoff: Oja's rule (used by OjaKV) is cheap but gives no per-token optimality guarantee and drifts at long context, while Brand's exact rank-1 SVD update is provably optimal but costs O(R³ + TR²) per token — growing with context length. We propose TKV (Tensor KV), which resolves this by running Brand's update only through a bounded warmup window, then freezing the resulting basis and projecting subsequent tokens onto it at O(R·d) per token indefinitely. A Φ-bridge reuses the Gaussian matrix already used for token-importance tracking as the Halko randomized range finder for cold-start initialization, and the update is implemented with XLA-static shapes for GPU/TPU compatibility. On real Mistral-7B activations, a TKV rank pyramid (avg. rank 15) reaches perplexity 20.1 versus an exact-cache baseline of 19.4 (3.6% degradation), while the same pyramid schedule applied to OjaKV degrades to 60× worse than exact; at 128K-token context TKV achieves 8.5× memory compression and lower reconstruction error than OjaKV at every long-context length tested (1K–32K tokens). Importance-weighted cold start further reduces important-token reconstruction error by 90–95% at rank 8–16.