🤖 AI Summary
This work addresses the limitation of existing KV cache quantization methods, which, despite reducing storage, still require dequantization to high-precision floating-point numbers and thus fail to alleviate the memory bandwidth bottleneck in attention computation. To overcome this, the paper introduces— for the first time—product quantization and asymmetric distance computation from vector retrieval into the Transformer attention mechanism. By leveraging subspace decomposition, codebook learning, and lookup-table-based operations, the approach transforms memory-bound attention into compute-bound without modifying the model architecture or requiring retraining. Evaluated on GPT-2, the method achieves 95.7% output fidelity at a 64× compression ratio (95.0% at 32×) and maintains a Spearman rank correlation coefficient ρ > 0.95 in attention scores. Both theoretical analysis and experiments confirm its effectiveness for sequence lengths up to 1024.
📝 Abstract
Compressing the KV cache is a required step to deploy large language models on edge devices. Current quantization methods compress storage but fail to reduce bandwidth as attention calculation requires dequantizing keys from INT4/INT8 to FP16 before use. We observe that attention scoring is mathematically equivalent to the inner product similarity search and we can apply some compression techniques from vector databases to compress KV-cache better. We propose LOOKAT, which applies product quantization and asymmetric distance computation, to transformer architecture by decomposing key vectors into subspaces, learning codebooks and computing attention tables via lookup tables. This transforms attention from memory-bound to compute-bound. LOOKAT achieves 64 $\times$ compression at 95.7\% output fidelity and 32 $\times$ compression at 95.0\% fidelity when tested on GPT-2. LOOKAT requires no architecture changes or training while maintaining rank correlation $\rho>0.95$. Theoretical analysis confirms that rank correlation degrades as $O(d_k/mK)$, with guarantees validated across sequence lengths up to 1024 tokens.