通过识别图像背景冗余信息，用轻量模型提前剪枝，提升ViT在边缘设备的推理速度。
Token Pruning using a Lightweight Background Aware Vision Transformer
- 用分割图或框标注训练轻量ViT区分前景与背景令牌。
- 在VOC和COCO上背景识别准确率达75%~88%，可显著减少输入令牌数。
- 无需稀疏微调即可提速30%-40%，适合边缘AI部署场景。
高运行内存和高延迟对视觉Transformer（ViT）的训练与推理构成重大限制，尤其在边缘设备上。令牌剪枝通过重要性标准减少输入到ViT的令牌数量。本文提出一种背景感知视觉Transformer（BAViT），作为目标检测模型如DETR/YOLOS的预处理模块，通过新方法识别图像中的背景令牌，从而在输入到基于ViT的目标检测器前完全或部分剪枝。利用分割图和/或边界框标注训练少量层的ViT以分类令牌为前景或背景。在VOC数据集上使用2层和10层的BAViT，背景与前景分离准确率分别为75%和88%；在COCO数据集上分别为71%和80%。将2层BAViT-small模型作为YOLOS的预处理器，在不进行稀疏微调的情况下可提升吞吐量30%-40%，仅损失3% mAP；进行稀疏微调后损失降为2%。该方法专为边缘AI应用场景设计。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。