Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Aryan Karmore

TL;DR

LOOKAT recasts attention as an approximate retrieval problem to break the memory-bound KV-cache bottleneck for edge deployment. By applying product quantization across subspaces and computing attention scores via asymmetric distance computations with precomputed lookup tables, it avoids key dequantization and dramatically reduces bandwidth while preserving rank order. The approach achieves up to 64× KV-cache compression with over 95% output fidelity on GPT-2, and maintains attention structure (Spearman ρ > 0.95) across sequence lengths up to 1024 tokens, with theoretical guarantees on rank preservation. It requires no architectural changes or training and relies on small codebooks per layer, though it currently targets key compression and depends on calibration data for codebook learning. This work demonstrates that bridging attention with vector retrieval techniques yields practical memory-speed improvements for edge AI devices, while outlining pathways for extending to values and further hardware optimizations.

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 $ρ> 0.95$. Theoretical analysis confirms that rank correlation degrades as $O(d_k/mK)$, with guarantees validated across sequence lengths up to 1024 tokens.

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

TL;DR

LOOKAT recasts attention as an approximate retrieval problem to break the memory-bound KV-cache bottleneck for edge deployment. By applying product quantization across subspaces and computing attention scores via asymmetric distance computations with precomputed lookup tables, it avoids key dequantization and dramatically reduces bandwidth while preserving rank order. The approach achieves up to 64× KV-cache compression with over 95% output fidelity on GPT-2, and maintains attention structure (Spearman ρ > 0.95) across sequence lengths up to 1024 tokens, with theoretical guarantees on rank preservation. It requires no architectural changes or training and relies on small codebooks per layer, though it currently targets key compression and depends on calibration data for codebook learning. This work demonstrates that bridging attention with vector retrieval techniques yields practical memory-speed improvements for edge AI devices, while outlining pathways for extending to values and further hardware optimizations.

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 compression at 95.7\% output fidelity and 32 compression at 95.0\% fidelity when tested on GPT-2. LOOKAT requires no architecture changes or training while maintaining rank correlation . Theoretical analysis confirms that rank correlation degrades as , with guarantees validated across sequence lengths up to 1024 tokens.
Paper Structure (40 sections, 1 theorem, 10 equations, 4 figures, 4 tables, 1 algorithm)

This paper contains 40 sections, 1 theorem, 10 equations, 4 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Literature Review
  3. Scalar Quantization for KV-Cache Compression
  4. Post-Training Quantization (PTQ).
  5. Quantization Aware Training
  6. Structural Approaches
  7. Attention Pruning
  8. Low Rank Factorisation
  9. Approximate Retrieval and Outlier Quantization
  10. Product Quantization
  11. Outlier Aware Quantization
  12. LOOKAT: Bridging Attention and Approximate Retrieval
  13. Method
  14. Problem Setup
  15. Scalar Quantization:INT4 and INT8
  16. ...and 25 more sections

Key Result

Proposition 1

Let $s$ be the accurate attention score and $\hat{s}$ the LOOKAT-approximated score. For PQ with $m$ subspaces and $K$ centroids per subspace, the expected rank correlation satisfies:

Figures (4)

  • Figure 1: Standard Attention Mechanism compared against LOOKAT attention. LOOKAT does not dequantize keys before computation and eliminates the bandwidth bottleneck.
  • Figure 2: Query is split into subspaces to precompute lookup tables (blue). Attention scores are computed via table lookups (green) using compressed key indices, replacing $O(L \cdot d_k)$ with $O(L \cdot m)$ lookups.
  • Figure 3: Comprehensive Evaluation of LOOKAT. (Top Left) Cosine Similarity measures output fidelity . (Top Right) Values of KL divergence on the log scale. (Bottom Left) Spearman's Rank correlation remains above 0.95. (Bottom Right) Pareto frontier showing that LOOKAT methods (red) dominate scalar quantization (blue) in high-compression regime. Error bars show standard deviation over 3 text samples.
  • Figure 4: Attention Pattern Reconstruction. Each row compares FP16 reference (left) with LOOKAT-4 (right) across natural language, code, and technical text. Attention peaks align spatially despite 32 $\times$ compression.

Theorems & Definitions (1)

  • Proposition 1