# Memory-efficient Transformers via Top-k Attention

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

## Key facts

- Cited by 60 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 55 references that are in the map.
- First posted on arXiv 13 Jun 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Deep contextualized word representations](https://mlsanity.com/paper/1802.05365/) (2018)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge](https://mlsanity.com/paper/1811.00937/) (2018)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [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)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [RACE: Large-scale ReAding Comprehension Dataset From Examinations](https://mlsanity.com/paper/1704.04683/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [How Much Knowledge Can You Pack Into the Parameters of a Language Model?](https://mlsanity.com/paper/2002.08910/) (2020)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [The NarrativeQA Reading Comprehension Challenge](https://mlsanity.com/paper/1712.07040/) (2017)

## Built on by (most cited first, 30 of 60)

- [Diagonal State Spaces are as Effective as Structured State Spaces](https://mlsanity.com/paper/2203.14343/) (2022)
- [Memorizing Transformers](https://mlsanity.com/paper/2203.08913/) (2022)
- [BiFormer: Vision Transformer with Bi-Level Routing Attention](https://mlsanity.com/paper/2303.08810/) (2023)
- [The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers](https://mlsanity.com/paper/2210.06313/) (2022)
- [Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems](https://mlsanity.com/paper/2312.15234/) (2023)
- [MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding](https://mlsanity.com/paper/2408.11049/) (2024)
- [Loki: Low-rank Keys for Efficient Sparse Attention](https://mlsanity.com/paper/2406.02542/) (2024)
- [Universal Hopfield Networks: A General Framework for Single-Shot Associative Memory Models](https://mlsanity.com/paper/2202.04557/) (2022)
- [Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers](https://mlsanity.com/paper/2406.16747/) (2024)
- [Masked Hard-Attention Transformers Recognize Exactly the Star-Free Languages](https://mlsanity.com/paper/2310.13897/) (2023)
- [Attention Entropy is a Key Factor: An Analysis of Parallel Context Encoding with Full-attention-based Pre-trained Language Models](https://mlsanity.com/paper/2412.16545/) (2024)
- [Long-Context Generalization with Sparse Attention](https://mlsanity.com/paper/2506.16640/) (2025)
- [Pick of the Bunch: Detecting Infrared Small Targets Beyond Hit-Miss Trade-Offs via Selective Rank-Aware Attention](https://mlsanity.com/paper/2408.03717/) (2024)
- [Nonparametric Modern Hopfield Models](https://mlsanity.com/paper/2404.03900/) (2024)
- [On Learning the Transformer Kernel](https://mlsanity.com/paper/2110.08323/) (2021)
- [AdaSplash: Adaptive Sparse Flash Attention](https://mlsanity.com/paper/2502.12082/) (2025)
- [Treeformer: Dense Gradient Trees for Efficient Attention Computation](https://mlsanity.com/paper/2208.09015/) (2022)
- [ZETA: Leveraging Z-order Curves for Efficient Top-k Attention](https://mlsanity.com/paper/2501.14577/) (2025)
- [IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs](https://mlsanity.com/paper/2405.02842/) (2024)
- [HATA: Trainable and Hardware-Efficient Hash-Aware Top-k Attention for Scalable Large Model Inference](https://mlsanity.com/paper/2506.02572/) (2025)
- [LoMA: Lossless Compressed Memory Attention](https://mlsanity.com/paper/2401.09486/) (2024)
- [HashEvict: A Pre-Attention KV Cache Eviction Strategy using Locality-Sensitive Hashing](https://mlsanity.com/paper/2412.16187/) (2024)
- [Pretraining the Vision Transformer using self-supervised methods for vision based Deep Reinforcement Learning](https://mlsanity.com/paper/2209.10901/) (2022)
- [What makes Models Compositional? A Theoretical View: With Supplement](https://mlsanity.com/paper/2405.02350/) (2024)
- [MultiMax: Sparse and Multi-Modal Attention Learning](https://mlsanity.com/paper/2406.01189/) (2024)
- [Exploiting Sparsity for Long Context Inference: Million Token Contexts on Commodity GPUs](https://mlsanity.com/paper/2502.06766/) (2025)
- [Mamba Adaptive Anomaly Transformer with association discrepancy for time series](https://mlsanity.com/paper/2502.07858/) (2025)
- [Spark Transformer: Reactivating Sparsity in FFN and Attention](https://mlsanity.com/paper/2506.06644/) (2025)
- [IceCache: Memory-efficient KV-cache Management for Long-Sequence LLMs](https://mlsanity.com/paper/2604.10539/) (2026)
- [Consciousness-Inspired Spatio-Temporal Abstractions for Better Generalization in Reinforcement Learning](https://mlsanity.com/paper/2310.00229/) (2023)

## Abstract (arXiv)

Following the success of dot-product attention in Transformers, numerous approximations have been recently proposed to address its quadratic complexity with respect to the input length. While these variants are memory and compute efficient, it is not possible to directly use them with popular pre-trained language models trained using vanilla attention, without an expensive corrective pre-training stage. In this work, we propose a simple yet highly accurate approximation for vanilla attention. We process the queries in chunks, and for each query, compute the top-k scores with respect to the keys. Our approach offers several advantages: (a) its memory usage is linear in the input size, similar to linear attention variants, such as Performer and RFA (b) it is a drop-in replacement for vanilla attention that does not require any corrective pre-training, and (c) it can also lead to significant memory savings in the feed-forward layers after casting them into the familiar query-key-value framework. We evaluate the quality of top-k approximation for multi-head attention layers on the Long Range Arena Benchmark, and for feed-forward layers of T5 and UnifiedQA on multiple QA datasets. We show our approach leads to accuracy that is nearly-identical to vanilla attention in multiple setups including training from scratch, fine-tuning, and zero-shot inference.