GSRQ: Gain-Shape Residual Quantization for Sub-1-bit KV Cache

The deployment of Large Language Models (LLMs) with extended context windows is increasingly constrained by the linear growth of Key-Value (KV) cache memory. Vector Quantization (VQ), particularly Residual Quantization (RQ), is a promising approach for pushing KV cache storage toward the sub-1-bit regime by progressively encoding residuals with small codebooks. However, most VQ methods still rely on standard $\ell_2$ $K$-means as the core codebook-learning primitive. We identify a subtle high-dimensional issue of this primitive: Euclidean centroid averaging can induce centroid shrinkage, which weakens the angular alignment term in the $\ell_2$ distortion and makes directional preservation harder. To address this issue, we propose Gain-Shape $K$-means (GSKM), a drop-in replacement for $K$-means that improves directional fidelity while matching, and in some regimes improving, $\ell_2$ distortion. We then build Gain-Shape Residual Quantization (GSRQ) by incorporating a weighted extension of GSKM into an RQ pipeline. On LLaMA-3-8B, GSRQ substantially improves over strong KV cache quantization baselines across bit rates. At 1-bit, it improves the average accuracy across LongBench tasks from 11.34 to 33.54, a gain of 22.20 percentage points over VQLLM.

Large language models can read and generate long documents, but doing so requires them to remember information from all previous words. This memory, called the key-value cache, grows linearly with the length of the input and can quickly become too large for efficient use on GPUs. A common way to reduce this memory is to compress groups of numbers into a small set of representative patterns. However, existing methods often learn these representatives by simple averaging. In high-dimensional data, this averaging can shrink the representatives and make them less aligned to the original directions of the data, which is important for preserving model behavior. We propose a new compression method that separates each group representative into two parts-its size and its direction-and refines these two parts separately. This helps preserve directional information while still keeping the memory footprint very small. Our approach improves the quality of compressed key-value caches, especially in extremely low-bit settings. As a result, large language models can handle longer contexts with much lower memory cost while maintaining stronger performance on reasoning and long-document tasks.