# Token-Picker: Accelerating Attention in Text Generation with Minimized Memory Transfer via Probability Estimation

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 16 references that are in the map.
- First posted on arXiv 21 Jul 2024.

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

- [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)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning](https://mlsanity.com/paper/2012.09852/) (2020)

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

- [BitStopper: An Efficient Transformer Attention Accelerator via Stage-fusion and Early Termination](https://mlsanity.com/paper/2512.06457/) (2025)
- [LAPA: Log-Domain Prediction-Driven Dynamic Sparsity Accelerator for Transformer Model](https://mlsanity.com/paper/2512.07855/) (2025)
- [HiKV: Hierarchical Importance-Aware KV Cache with Hardware Acceleration for LLM Decoding](https://mlsanity.com/paper/2607.22389/) (2026)

## Abstract (arXiv)

The attention mechanism in text generation is memory-bounded due to its sequential characteristics. Therefore, off-chip memory accesses should be minimized for faster execution. Although previous methods addressed this by pruning unimportant tokens, they fall short in selectively removing tokens with near-zero attention probabilities in each instance. Our method estimates the probability before the softmax function, effectively removing low probability tokens and achieving an 12.1x pruning ratio without fine-tuning. Additionally, we present a hardware design supporting seamless on-demand off-chip access. Our approach shows 2.6x reduced memory accesses, leading to an average 2.3x speedup and a 2.4x energy efficiency.