Learned Token Pruning for Transformers

2021/07/02 by Sehoon Kim, Kim, Sehoon, Sheng Shen +11 · 27 citations
Computer Science · #Advanced Data Storage Technologies #Advanced Neural Network Applications #Algorithm #Artificial intelligence #Computation and Language (cs.CL) #Computer network #Computer science #FLOPS #FOS: Computer and information sciences #Inference #Parallel Computing and Optimization Techniques #Parallel computing #Pruning #Security token #Sequence (biology) #Transformer #cs.CL

paper · pdf · doi:10.48550/arxiv.2107.00910

published in arXiv (Cornell University) (Cornell University) · KDD 2022 (Research Track)

openalex publication_date 2021/07/02 · arxiv created 2022/06/02 · arxiv updated 2022/06/06 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Deploying transformer models in practice is challenging due to their inference cost, which scales quadratically with input sequence length. To address this, we present a novel Learned Token Pruning (LTP) method 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 value which 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 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.1x FLOPs reduction with less than 1% accuracy drop, which results in up to 1.9x and 2.0x throughput improvement on Intel Haswell CPUs and NVIDIA V100 GPUs, respectively. Furthermore, we demonstrate that LTP is more robust than prior methods to variations on input sentence lengths. Our code has been developed in PyTorch and has been open-sourced.

Cited by

• Token Expand-Merge: Training-Free Token Compression for Vision-Language-Action Models
• Rethinking Chain-of-Thought Reasoning for Videos
• Leveraging KV Similarity for Online Structured Pruning in LLMs
• AdmTree: Compressing Lengthy Context with Adaptive Semantic Trees
• Extreme Model Compression for Edge Vision-Language Models: Sparse Temporal Token Fusion and Adaptive Neural Compression
• DeepCoT: Deep Continual Transformers for Real-Time Inference on Data Streams
• The Few Govern the Many:Unveiling Few-Layer Dominance for Time Series Models
• SparseVILA: Decoupling Visual Sparsity for Efficient VLM Inference
• Efficient Adaptive Transformer: An Empirical Study and Reproducible Framework
• Patch Rebirth: Toward Fast and Transferable Model Inversion of Vision Transformers
• Where Do Tokens Go? Understanding Pruning Behaviors in STEP at High Resolutions
• The Better You Learn, The Smarter You Prune: Towards Efficient Vision-language-action Models via Differentiable Token Pruning
• Purge-Gate: Backpropagation-Free Test-Time Adaptation for Point Clouds Classification via Token Purging
• Prune2Drive: A Plug-and-Play Framework for Accelerating Vision-Language Models in Autonomous Driving
• Representation Shift: Unifying Token Compression with FlashAttention
• Polymorph: Energy-Efficient Multi-Label Classification for Video Streams on Embedded Devices
• Local Representative Token Guided Merging for Text-to-Image Generation
• Input Conditioned Layer Dropping in Speech Foundation Models
• Put Teacher in Student's Shoes: Cross-Distillation for Ultra-compact Model Compression Framework
• KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction
• S2AFormer: Strip Self-Attention for Efficient Vision Transformer
• Sparsified State-Space Models are Efficient Highway Networks
• Shifting AI Efficiency From Model-Centric to Data-Centric Compression
• Token Reduction Should Go Beyond Efficiency in Generative Models -- From Vision, Language to Multimodality
• An Effective Training Framework for Light-Weight Automatic Speech Recognition Models
• Balanced and Elastic End-to-end Training of Dynamic LLMs
• Back to Fundamentals: Low-Level Visual Features Guided Progressive Token Pruning
• TC-MGC: Text-Conditioned Multi-Grained Contrastive Learning for Text-Video Retrieval
• Saliency-driven Dynamic Token Pruning for Large Language Models

Related

• A Fast Post-Training Pruning Framework for Transformers
• Zero-TPrune: Zero-Shot Token Pruning through Leveraging of the Attention Graph in Pre-Trained Transformers
• Random-LTD: Random and Layerwise Token Dropping Brings Efficient Training for Large-scale Transformers
• SEVEN: Pruning Transformer Model by Reserving Sentinels
• Memory-Efficient Fine-Tuning of Transformers via Token Selection