Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Token Pruning using a Lightweight Background Aware Vision Transformer

Sudhakar Sah, Ravish Kumar, Honnesh Rohmetra, Ehsan Saboori

TL;DR

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.

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.

Token Pruning using a Lightweight Background Aware Vision Transformer

TL;DR

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.

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.

Paper Structure

This paper contains 15 sections, 4 equations, 6 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Vision Transformers
  4. Runtime Memory Improvement
  5. Token Pruning
  6. Methodology
  7. Auxiliary Annotations
  8. BAViT Architecture
  9. Accumulative Cross Entropy Loss
  10. Model Training
  11. BAViT Integration with ViT based Detection
  12. Results
  13. BG/FG Classification Model
  14. Token reduction using BAViT
  15. Conclusion

Figures (6)

  • Figure 1: Comparison of background token identification results between Sparse DETR, Focus DETR and BAViT models
  • Figure 2: Three VOC images (left to right) a) original image b) foreground object area in transparency c)$16\times16$ grid with red grids bring foreground and gray being background
  • Figure 3: Background aware ViT architecture. (Left) 2 layers for BG and FG patch classification. (Right) BAViT attached as a pre-processing step to YOLOS (DETR type object detector) object detector
  • Figure 4: FG/BG token classification (16x16) on COCO images. Top- original image, bottom - sparse image generated from BAViT with sparsity percentage.
  • Figure 5: Converting BG tokens to FG tokens using a post processing convolution operation
  • ...and 1 more figures