Researchers introduce JoLT, a method for compressing the key-value (KV) cache in transformer inference by treating it as a third-order tensor. The approach applies partial Tucker decomposition to the token and feature axes while using Johnson-Lindenstrauss rotated low-bit residuals to restore energy lost during truncation.
- A single Lagrangian dual allocates Tucker ranks and residual bit-widths per layer group under one byte budget.
- JoLT achieves 2-3x compression with perplexity, GSM8K accuracy, and RULER retrieval staying within statistical noise of the uncompressed baseline on Mistral-7B-v0.3 and LLaMA-2-13B.
- At 2x compression, relative Frobenius error is 0.009 for keys and 0.006 for values, an order of magnitude lower than cross-layer SVD and 4-bit quantization.
- A randomized-SVD variant named FlashJoLT provides a 5-13x compression-time speedup at matched quality.
This method addresses the dominant memory cost of transformer inference by significantly reducing KV cache size without degrading model performance.