# kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers

arXiv: https://arxiv.org/abs/2411.04013 · HTML: https://mlsanity.com/paper/2411.04013/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 47 references that are in the map.
- First posted on arXiv 6 Nov 2024.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Enhancing the Locality and Breaking the Memory Bottleneck of Transformer on Time Series Forecasting](https://mlsanity.com/paper/1907.00235/) (2019)
- [Transformers in Time Series: A Survey](https://mlsanity.com/paper/2202.07125/) (2022)
- [Memorizing Transformers](https://mlsanity.com/paper/2203.08913/) (2022)
- [BookSum: A Collection of Datasets for Long-form Narrative Summarization](https://mlsanity.com/paper/2105.08209/) (2021)
- [Blockwise Self-Attention for Long Document Understanding](https://mlsanity.com/paper/1911.02972/) (2019)
- [Adaptively Sparse Transformers](https://mlsanity.com/paper/1909.00015/) (2019)
- [On The Computational Complexity of Self-Attention](https://mlsanity.com/paper/2209.04881/) (2022)
- [Unlimiformer: Long-Range Transformers with Unlimited Length Input](https://mlsanity.com/paper/2305.01625/) (2023)
- [Asymmetric LSH (ALSH) for Sublinear Time Maximum Inner Product Search (MIPS)](https://mlsanity.com/paper/1405.5869/) (2014)
- [Fast Attention Requires Bounded Entries](https://mlsanity.com/paper/2302.13214/) (2023)
- [Star-Transformer](https://mlsanity.com/paper/1902.09113/) (2019)
- [HyperAttention: Long-context Attention in Near-Linear Time](https://mlsanity.com/paper/2310.05869/) (2023)
- [The De-democratization of AI: Deep Learning and the Compute Divide in Artificial Intelligence Research](https://mlsanity.com/paper/2010.15581/) (2020)
- [Memory-efficient Transformers via Top-k Attention](https://mlsanity.com/paper/2106.06899/) (2021)
- [KDEformer: Accelerating Transformers via Kernel Density Estimation](https://mlsanity.com/paper/2302.02451/) (2023)
- [Is Sparse Attention more Interpretable?](https://mlsanity.com/paper/2106.01087/) (2021)
- [The Fine-Grained Complexity of Gradient Computation for Training Large Language Models](https://mlsanity.com/paper/2402.04497/) (2024)
- [The Cascade Transformer: an Application for Efficient Answer Sentence Selection](https://mlsanity.com/paper/2005.02534/) (2020)
- [Fast Amortized Inference and Learning in Log-linear Models with Randomly Perturbed Nearest Neighbor Search](https://mlsanity.com/paper/1707.03372/) (2017)
- [The I/O Complexity of Attention, or How Optimal is Flash Attention?](https://mlsanity.com/paper/2402.07443/) (2024)

## Built on by (most cited first, 1 of 1)

- [More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding](https://mlsanity.com/paper/2610.04753/) (2026)

## Abstract (arXiv)

Despite their power, Transformers face challenges with long sequences due to the quadratic complexity of self-attention. To address this limitation, methods like k-Nearest-Neighbor (kNN) attention have been introduced [Roy, Saffar, Vaswani, Grangier, 2021] enabling each token to attend to only its k closest tokens. While kNN attention has shown empirical success in making Transformers more efficient, its exact approximation guarantees have not been theoretically analyzed. In this work, we establish a theoretical framework for kNN attention, reformulating self-attention as expectations over softmax distributions and leveraging lazy Gumbel sampling [Mussmann, Levy, Ermon, 2017] with kNN indices for efficient approximation. Building on this framework, we also propose novel sub-quadratic algorithms that approximate self-attention gradients by leveraging efficient sampling techniques, such as Markov Chain-based estimation. Finally, we demonstrate the practical effectiveness of these algorithms through empirical experiments, showcasing their benefits in both training and inference.