The KV-cache is today the main Achilles' heel of transformers when scaling the context. It grows linearly with the length of the sequence and is stored for each layer and each head.
For example, for a LLaMA-like model with 65B parameters, the KV-cache at 128k tokens occupies ~335 GB of memory. And it's also a pain in terms of time.
🟢 How NVIDIA solved Pain?
However, most optimizations reduce the KV-cache by layers or by heads. Although the main potential is precisely along the token axis: not all of them are really needed by the model.
The first working method of reducing KV by tokens was invented by the authors of KVzip: up to 4× compression with zero quality loss. But in practice, the method turned out to be too slow.
Nvidia took this idea, modified it a bit, and got almost the same result, but practically for free.
They simply train a small model that predicts, based on the hidden state of a token, how important its KV is. It's different for each layer, but it's either a linear model or a two-layer MLP - a maximum of 1–2 matrix multiplications.
And that's it, no expensive operations and recalculations (for comparison: in KVzip, the prompt essentially had to be run twice). Next, the KV pairs with a significance below a specified threshold are simply discarded.
The compute overhead is about 0.02% FLOPs for linear models. On a long context, this is noise against the quadratic attention.
The degradation on benchmarks is about zero, the compression is 3–4×. It's just a fairy tale (though, of course, much still depends on the engine).
Hats off to Nvidia for their excellent work. Everything is in the open source on GitHub | Paper
••••••••••••••••••••••••••••••••••••••
🤖 Data Science, ML & Big Data with @DataXplore