GSRQ: Gain-Shape Residual Quantization for Sub-1-bit KV Cache
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Language Modeling | C4 (val) | PPL5.8 | 908 | |
| Language Modeling | C4 | Perplexity8.8 | 482 | |
| Physical Interaction Question Answering | PIQA | Accuracy78.56 | 462 | |
| Language Modeling | WikiText-2 | Perplexity8.88 | 205 | |
| Language Modeling | WikiText-2 (val) | Perplexity (BVS)4.98 | 179 | |
| Long-context evaluation | LongBench | Average Score39.39 | 96 | |
| Common Sense Reasoning | HellaSwag | Accuracy76.38 | 85 | |
| Common Sense Reasoning | WinoGrande | Accuracy73.16 | 67 | |
| Reasoning | ARC Challenge | Accuracy59.3 | 44 | |
| Mathematical Reasoning | MATH 500 | Accuracy19 | 22 |