Residual vector quantization for KV cache compression in large language model

TL;DR

This paper introduces a residual vector quantization method for compressing key-value caches in large language models, achieving significant memory reduction with minimal performance loss.

Contribution

It adapts residual vector quantization for LLM KV cache compression, simplifying the process and improving efficiency over scalar quantization methods.

Findings

Residual depth of 8 recovers most model performance.

Grouping non-contiguous channels improves compression.

Method achieves 5.5x compression with minimal fine-tuning.

Abstract

KV cache compression methods have mainly relied on scalar quantization techniques to reduce the memory requirements during decoding. In this work, we apply residual vector quantization, which has been widely used for high fidelity audio compression, to compress KV cache in large language models (LLM). We adapt the standard recipe with minimal changes to compress the output of any key or value projection matrix in a pretrained LLM: we scale the vector by its standard deviation, divide channels into groups and then quantize each group with the same residual vector quantizer. We learn the codebook using exponential moving average and there are no other learnable parameters including the input and output projections normally used in a vector quantization set up. We find that a residual depth of 8 recovers most of the performance of the unquantized model. We also find that grouping non-contiguous channels together works better than grouping contiguous channels for compressing key matrix and the method further benefits from a light weight finetuning of LLM together with the quantization. Overall, the proposed technique is competitive with existing quantization methods while being much simpler and results in 5.5x compression compared to half precision.