Token Pruning using a Lightweight Background Aware Vision Transformer
On the map: Efficient LLM inference and serving → Token pruning and merging for ViTs
Abstract (arXiv)
High runtime memory and high latency puts significant constraint on Vision Transformer training and inference, especially on edge devices. Token pruning reduces the number of input tokens to the ViT based on importance criteria of each token. We present a Background Aware Vision Transformer (BAViT) model, a pre-processing block to object detection models like DETR/YOLOS aimed to reduce runtime memory and increase throughput by using a novel approach to identify background tokens in the image. The background tokens can be pruned completely or partially before feeding to a ViT based object detector. We use the semantic information provided by segmentation map and/or bounding box annotation to train a few layers of ViT to classify tokens to either foreground or background. Using 2 layers and 10 layers of BAViT, background and foreground tokens can be separated with 75% and 88% accuracy on VOC dataset and 71% and 80% accuracy on COCO dataset respectively. We show a 2 layer BAViT-small model as pre-processor to YOLOS can increase the throughput by 30% - 40% with a mAP drop of 3% without any sparse fine-tuning and 2% with sparse fine-tuning. Our approach is specifically targeted for Edge AI use cases.
Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Placed in Token pruning and merging for ViTs, one of 12 directions in Efficient LLM inference and serving. Builds on 29 of its 35 references that are in the map. First posted on arXiv 12 Oct 2024.
Builds on
29 of 35 references.
Closest ancestors · 2
Same direction, newest first
- 2023Revisiting Token Pruning for Object Detection and Instance Segmentation12 Jun
- 2022Adaptive Sparse ViT: Towards Learnable Adaptive Token Pruning by Fully Exploiting Self-Attention28 Sep
Other references · 27
- 2017Attention Is All You Needcited by 59,932
- 2015Deep Residual Learning for Image Recognitioncited by 51,809
- 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
- 2014Microsoft COCO: Common Objects in Contextcited by 17,498
- 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,266
- 2021Swin Transformer: Hierarchical Vision Transformer using Shifted Windowscited by 12,082
- 2020End-to-End Object Detection with Transformerscited by 7,384
- 2020Deformable DETR: Deformable Transformers for End-to-End Object Detectioncited by 3,410
- 2020Big Bird: Transformers for Longer Sequencescited by 1,673
- 2020Reformer: The Efficient Transformercited by 1,658
- 2019Generating Long Sequences with Sparse Transformerscited by 1,523
15 more
- 2019Generalized Intersection over Union: A Metric and A Loss for Bounding Box Regressioncited by 1,361
- 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 844
- 2023DETRs Beat YOLOs on Real-time Object Detectioncited by 526
- 2022ViTPose: Simple Vision Transformer Baselines for Human Pose Estimationcited by 450
- 2019Reducing Transformer Depth on Demand with Structured Dropoutcited by 434
- 2019Q-BERT: Hessian Based Ultra Low Precision Quantization of BERTcited by 355
- 2020Sparse Sinkhorn Attentioncited by 250
- 2021You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detectioncited by 170
- 2020Training with Quantization Noise for Extreme Model Compressioncited by 145
- 2021All Tokens Matter: Token Labeling for Training Better Vision Transformerscited by 144
- 2021Sparse DETR: Efficient End-to-End Object Detection with Learnable Sparsitycited by 98
- 2023Semantic Segmentation using Vision Transformers: A surveycited by 73
- 2019Efficient 8-Bit Quantization of Transformer Neural Machine Language Translation Modelcited by 58
- 2023Less is More: Focus Attention for Efficient DETRcited by 38
- 2022Dynamic Spatial Sparsification for Efficient Vision Transformers and Convolutional Neural Networkscited by 23
Built on by
Citing papers per year: 2025: 2; 2026: 1.
3 papers.
- 2026Increasing the Efficiency of DETR for Maritime High-Resolution Images11 May 2026
- 2025Focus: A Streaming Concentration Architecture for Efficient Vision-Language Models16 Dec 2025
- 2025Vision Transformers on the Edge: A Comprehensive Survey of Model Compression and Acceleration Strategies26 Feb 2025
Similar, not cited
Close in wording, no citation either way
- 2024Token Cropr: Faster ViTs for Quite a Few Tasks
- 2023No Token Left Behind: Efficient Vision Transformer via Dynamic Token Idling
- 2021SPViT: Enabling Faster Vision Transformers via Soft Token Pruning
- 2023PPT: Token Pruning and Pooling for Efficient Vision Transformers
- 2023MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers
- 2025Similarity-Aware Token Pruning: Your VLM but Faster
- 2025Where Do Tokens Go? Understanding Pruning Behaviors in STEP at High Resolutions
- 2025Token Transforming: A Unified and Training-Free Token Compression Framework for Vision Transformer Acceleration