Learned Token Pruning for Transformers
Open accessAbstract
Efficient deployment of transformer models in practice is challenging due to their inference cost including memory footprint, latency, and power consumption, which scales quadratically with input sequence length. To address this, we present a novel token reduction method dubbed Learned Token Pruning (LTP) which adaptively removes unimportant tokens as an input sequence passes through transformer layers. In particular, LTP prunes tokens with an attention score below a threshold, whose value is learned for each layer during training. Our threshold-based method allows the length of the pruned sequence to vary adaptively based on the input sequence, and avoids algorithmically expensive operations such as top-k token selection. We extensively test the performance of LTP on GLUE and SQuAD tasks and show that our method outperforms the prior state-of-the-art token pruning methods by up to ∽2.5% higher accuracy with the same amount of FLOPs. In particular, LTP achieves up to 2.1× FLOPs reduction with less than 1% accuracy drop, which results in up to 1.9× and 2.0× throughput improvement on Intel Haswell CPUs and NVIDIA V100 GPUs. Furthermore, we demonstrate that LTP is more robust than prior methods to variations in input sequence lengths. Our code has been developed in PyTorch and open-sourced
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
Claims other papers make about this one
Each is checked against this paper’s own text: does it say what the citing paper says it does?
Applying the same pruning percentage can potentially under-prune short sequences or over-prune long sequences.
“Therefore, applying the same pruning percentage can potentially under-prune short sequence or over-prune long sequence [103].”
From A Survey of Techniques for Optimizing Transformer Inference (2023) · cites this paper as [103]
Cites
BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
Devlin, J., Chang, M.-W., Lee, K., & Toutanova, K. (2019) · Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) · doi:10.18653/v1/n19-1423 · Cited by 48,293 · see citing works
RoBERTa: A Robustly Optimized BERT Pretraining Approach
Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., & Stoyanov, V. (2019) · arXiv · doi:10.48550/arxiv.1907.11692 · Cited by 16,515 · see citing works
DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
Sanh, V., Debut, L., Chaumond, J., & Wolf, T. (2019) · arXiv · doi:10.48550/arxiv.1910.01108 · Cited by 4,653 · see citing works
A Survey of Quantization Methods for Efficient Neural Network Inference
Gholami, A., Kim, S., Dong, Z., Yao, Z., Mahoney, M. W., & Keutzer, K. (2022) · Low-Power Computer Vision · doi:10.1201/9781003162810-13 · Cited by 1,219 · cited 1× · see citing works
“For quantization, we use the static uniform symmetric integer quantization method [12], which is easy to deploy in commodity hardware with minimal run-time overhead.”
MethodSelf-citationLinformer: Self-Attention with Linear Complexity
Wang, S., Li, B. Z., Khabsa, M., Fang, H., & Ma, H. (2020) · arXiv · doi:10.48550/arxiv.2006.04768 · Cited by 867 · see citing works
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020) · arXiv · doi:10.48550/arxiv.2006.16236 · Cited by 827 · see citing works
MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
Sun, Z., Yu, H., Song, X., Liu, R., Yang, Y., & Zhou, D. (2020) · Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics · doi:10.18653/v1/2020.acl-main.195 · Cited by 757 · see citing works
SQuAD: 100,000+ Questions for Machine Comprehension of Text
Rajpurkar, P., Zhang, J., Lopyrev, K., & Liang, P. (2016) · arXiv · doi:10.48550/arxiv.1606.05250 · Cited by 678 · see citing works
Movement pruning: Adaptive sparsity by fine-tuning
Sanh, V., Wolf, T., & Rush, A. M. (2020) · arXiv · doi:10.48550/arxiv.2005.07683 · Cited by 218 · see citing works
SqueezeBERT: What can computer vision teach NLP about efficient neural networks?
Iandola, F., Shaw, A., Krishna, R., & Keutzer, K. (2020) · Proceedings of SustaiNLP: Workshop on Simple and Efficient Natural Language Processing · doi:10.18653/v1/2020.sustainlp-1.17 · Cited by 114 · see citing works
Self-citation
Flagged first, then by citations
Cited by
Sorted by citations
A Review on Edge Large Language Models: Design, Execution, and Applications
Zheng, Y., Chen, Y., Qian, B., Shi, X., Shu, Y., & Chen, J. (2025) · ACM Computing Surveys · doi:10.1145/3719664 · Cited by 169 · see citing works
A Survey of Techniques for Optimizing Transformer Inference
Chitty-Venkata, K. T., Mittal, S., Emani, M., Vishwanath, V., & Somani, A. K. (2023) · Journal of Systems Architecture · doi:10.1016/j.sysarc.2023.102990 · Cited by 154 · see citing works
LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models
Jiang, H., Wu, Q., Lin, C.-Y., Yang, Y., & Qiu, L. (2023) · Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing · doi:10.18653/v1/2023.emnlp-main.825 · Cited by 168 · see citing works
Token Merging: Your ViT But Faster
Bolya, D., Fu, C.-Y., Dai, X., Zhang, P., Feichtenhofer, C., & Hoffman, J. (2022) · arXiv · doi:10.48550/arxiv.2210.09461 · Cited by 140 · see citing works
“…any number of tokens, several recent works have attempted to prune the tokens from transformers in both NLP (Goyal et al., 2020; Kim & Cho, 2020; Kim et al., 2021; Lassance et al., 2021) and Vision (Meng et al., 2022; Yin et al., 2022; Kong et al., 2022; Song et al., 2022; Rao et al., 2021;…”
BackgroundInfluentialDiet code is healthy: simplifying programs for pre-trained models of code
Zhang, Z., Zhang, H., Shen, B., & Gu, X. (2022) · Proceedings of the 30th ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering · doi:10.1145/3540250.3549094 · Cited by 41 · see citing works
Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
Ge, S., Zhang, Y., Liu, L., Zhang, M., Han, J., & Gao, J. (2023) · arXiv · doi:10.48550/arxiv.2310.01801 · Cited by 13 · see citing works
LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models
Jiang, H., Wu, Q., Lin, C.-Y., Yang, Y., & Qiu, L. (2023) · arXiv · Cited by 168 · see citing works
Kong, Z., Li, Y., Zeng, F., Xin, L., Messica, S., Lin, X., Zhao, P., Kellis, M., Tang, H., & Zitnik, M. (2025) · arXiv · doi:10.48550/arxiv.2505.18227
“Learned token pruning [71] extends this approach by learning attention-based thresholds to adaptively prune uninformative tokens, thereby reducing computational costs while preserving model performance.”
- Authors
- Kim, S., Shen, S., Thorsley, D., Gholami, A., Kwon, W., Hassoun, J., & Keutzer, K.
- Affiliations
- University of California, Berkeley (US)
- Samsung Semiconductor (United States) (US)
- Venue
- Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining
- Publisher
- ACM
- Published
- 2022-08-14
- Pages
- 784-794
- OpenAlex
- W3180037928
- Cited by
- 139
- References
- 10
- Citations by year
- 2026: 22 · 2025: 43 · 2024: 44 · 2023: 22 · 2022: 6 · 2021: 2
- FWCI
- 10.43
- Type
- conference paper
- Language
- en
- Open access
- hybrid
- doi.org
- Licence
- Locations
- ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD) (publishedVersion, cc-by)
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University)
- Topics
- Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Advanced Neural Network Applications
- Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
- Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
- Keywords
- token pruning, transformer models, inference efficiency, FLOPs reduction, throughput improvement, input sequence length, memory usage, power consumption
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
