# Token Pruning using a Lightweight Background Aware Vision Transformer

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

## Key facts

- Cited by 3 papers in the map (data through 14 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.

## Placed in

- Efficient LLM inference and serving / [Token pruning and merging for ViTs](https://mlsanity.com/inference/token-pruning-and-merging-for-vits/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Deformable DETR: Deformable Transformers for End-to-End Object Detection](https://mlsanity.com/paper/2010.04159/) (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)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Generalized Intersection over Union: A Metric and A Loss for Bounding Box Regression](https://mlsanity.com/paper/1902.09630/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [DETRs Beat YOLOs on Real-time Object Detection](https://mlsanity.com/paper/2304.08069/) (2023)
- [ViTPose: Simple Vision Transformer Baselines for Human Pose Estimation](https://mlsanity.com/paper/2204.12484/) (2022)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT](https://mlsanity.com/paper/1909.05840/) (2019)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detection](https://mlsanity.com/paper/2106.00666/) (2021)
- [Training with Quantization Noise for Extreme Model Compression](https://mlsanity.com/paper/2004.07320/) (2020)
- [All Tokens Matter: Token Labeling for Training Better Vision Transformers](https://mlsanity.com/paper/2104.10858/) (2021)
- [Sparse DETR: Efficient End-to-End Object Detection with Learnable Sparsity](https://mlsanity.com/paper/2111.14330/) (2021)
- [Semantic Segmentation using Vision Transformers: A survey](https://mlsanity.com/paper/2305.03273/) (2023)
- [Efficient 8-Bit Quantization of Transformer Neural Machine Language Translation Model](https://mlsanity.com/paper/1906.00532/) (2019)
- [Revisiting Token Pruning for Object Detection and Instance Segmentation](https://mlsanity.com/paper/2306.07050/) (2023)
- [Less is More: Focus Attention for Efficient DETR](https://mlsanity.com/paper/2307.12612/) (2023)
- [Adaptive Sparse ViT: Towards Learnable Adaptive Token Pruning by Fully Exploiting Self-Attention](https://mlsanity.com/paper/2209.13802/) (2022)
- [Dynamic Spatial Sparsification for Efficient Vision Transformers and Convolutional Neural Networks](https://mlsanity.com/paper/2207.01580/) (2022)

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

- [Vision Transformers on the Edge: A Comprehensive Survey of Model Compression and Acceleration Strategies](https://mlsanity.com/paper/2503.02891/) (2025)
- [Focus: A Streaming Concentration Architecture for Efficient Vision-Language Models](https://mlsanity.com/paper/2512.14661/) (2025)
- [Increasing the Efficiency of DETR for Maritime High-Resolution Images](https://mlsanity.com/paper/2605.10269/) (2026)

## 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.