用查表法压缩注意力缓存，让小设备也能跑大模型。
LOOKAT: Lookup-Optimized Key-Attention for Memory-Efficient Transformers
- 将键向量分块后查表压缩，变内存瓶颈为计算瓶颈。
- 在GPT-2上实现64倍压缩，输出准确率仍达95.7%。
- 无需训练或改架构，适合边缘部署的轻量推理场景。
压缩KV缓存是将大语言模型部署到边缘设备的必要步骤。现有量化方法虽减少存储，但注意力计算需将INT4/INT8键解量化至FP16，无法降低带宽。我们发现注意力评分在数学上等价于内积相似性搜索，可借鉴向量数据库的压缩技术优化KV缓存。提出LOOKAT，通过分解键向量为子空间、学习码本并利用查表计算注意力表，实现对变压器架构的改进。该方法使注意力从内存密集型转为计算密集型。在GPT-2上测试，LOOKAT实现64×压缩时输出保真度达95.7%，32×压缩时保真度为95.0%。无需架构改动或训练，保持秩相关性ρ>0.95。理论分析表明，秩相关性退化随O(d_k/mK)变化，且在最长1024个标记的序列上验证有效。
原文摘要 · Abstract (English)
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 $ρ> 0.95$. Theoretical analysis confirms that rank correlation degrades as $O(d_k/mK)$, with guarantees validated across sequence lengths up to 1024 tokens.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。