Token Pruning using a Lightweight Background Aware Vision Transformer

TL;DR

This paper introduces BAViT, a background-aware token pruning method for Vision Transformers that reduces runtime memory and increases throughput on edge devices by effectively identifying and pruning background tokens.

Contribution

The paper presents a novel pre-processing block, BAViT, that classifies and prunes background tokens in ViT-based object detectors, improving efficiency without significant accuracy loss.

Findings

BAViT achieves 75-88% accuracy in background/foreground classification.

Using BAViT as pre-processor increases YOLOS throughput by 30-40%.

The approach maintains competitive mAP with minimal fine-tuning.

Abstract

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.