LOOKAT: Lookup-Optimized Key-Attention for Memory-Efficient Transformers

TL;DR

LOOKAT introduces a novel key-attention method that leverages vector database compression techniques to significantly reduce memory and bandwidth requirements in transformer models without altering architecture or training.

Contribution

It applies product quantization and asymmetric distance computation to transformer attention, enabling high compression ratios while preserving model fidelity.

Findings

Achieves 64× compression with 95.7% fidelity on GPT-2.

Achieves 32× compression with 95.0% fidelity on GPT-2.

Maintains high rank correlation (>0.95) without architecture changes.

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.