Research map / Architectures and efficiency / Attention and transformers
kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Gives a theoretical framework for kNN attention via lazy Gumbel sampling and k-MIPS, plus sub-quadratic Markov-chain-based estimators for attention gradients, with small experiments on random inputs, toy transformers and GPT-2 XL fine-tuning.
Problem
kNN attention is used empirically to make Transformers efficient on long sequences, but its approximation guarantees are not theoretically analyzed, and approximate backward passes lack provable algorithms.
Approach
Reformulates self-attention outputs as expectations under softmax distributions, samples them sublinearly via lazy Gumbel sampling over kNN/MIPS indices, and estimates attention gradients (notably D^V) with single-step random-walk (Markov chain) simulations.
Key results
- kNN attention handles self-attention for n=10^6 on a MacBook Air M3 with 8GB RAM, while naive attention runs out of memory beyond n=20000.
- Approximation error on random inputs is minimal for k >= n^{1/8}, below the theoretically predicted sqrt(n) threshold.
- GPT-2 XL (1.61B) was fine-tuned with kNN attention on a single NVIDIA L40 GPU on a Shakespeare dataset.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace exact attention with kNN Attention | self-attention on random inputs · Random inputs | memory scalability | Helped |
| replace exact attention with kNN Attention | self-attention on random inputs · Random inputs | runtime scaling | Helped |
| increase k (ablation) | kNN Attention · Random inputs | approximation error | Helped |
| use smaller k (ablation) | kNN Attention · Random inputs | approximation error | Hurt |
| vary the input scale B while k is small (ablation) | kNN Attention · Random inputs | approximation error | Mixed |
| approximate attention gradients | gradient descent with MSE loss · Attention gradients, MSE | gradient approximation quality | Helped |
| approximate attention gradients | gradient descent with cross-entropy loss · Attention gradients, cross-entropy | gradient-descent convergence | Hurt |
| replace exact attention with kNN Attention | nanoGPT character-level Transformer · Shakespeare character-level language modeling | perplexity | Mixed |
| use kNN Attention for fine-tuning | GPT-2 XL · Shakespeare fine-tuning | fine-tuning feasibility | Helped |
Surprising findings
- Empirical error threshold for k (about n^{1/8}) is well below the theoretically predicted sqrt(n).
- Gradient descent with approximate gradients converges less well in the non-convex cross-entropy case than in the convex MSE case.
Limitations
- Gradient approximation experiments are small (N=100, d=3) and LLM gradient-approximation impact is left for future work.
- Theoretical guarantees rely on bounded-norm assumptions (e.g., ||V||_inf = O(log n), d = O(log n)).
- Optimal k in practice is significantly smaller than the predicted sqrt(n).
Experimental setup
Datasets: Shakespeare dataset random uniform matrices Q,K,V
Benchmarks: random input approximation error Shakespeare character-level perplexity
Metrics: approximation error running time memory perplexity gradient descent convergence
Baselines: naive O(n^2) attention Alman and Song attention gradient approximation
Builds on: Lazy Gumbel sampling Routing Transformer Memorizing Transformers Unlimiformer nanoGPT
Closest papers by meaning
- Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.90)
- Fast attention mechanisms: a tale of parallelism (NeurIPS 2025 · similarity 0.88)
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.88)
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 2024 · similarity 0.88)
- On the Optimal Memorization Capacity of Transformers (ICLR 2025 · similarity 0.88)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.