Token Cropr: Faster ViTs for Quite a Few Tasks

TL;DR

This paper introduces a fast token pruning method for Vision Transformers that uses auxiliary prediction heads to select relevant tokens, achieving significant speedups across various vision tasks with minimal performance loss.

Contribution

A novel token pruning approach using auxiliary heads for end-to-end token relevance prediction, applicable to multiple vision tasks, and achieving high efficiency with minimal accuracy impact.

Findings

Achieves 1.5 to 4x speedup across tasks

On ADE20k, 2x speedup with only 0.1 mIoU drop

Method maintains high performance with fast inference

Abstract

The adoption of Vision Transformers (ViTs) in resource-constrained applications necessitates improvements in inference throughput. To this end several token pruning and merging approaches have been proposed that improve efficiency by successively reducing the number of tokens. However, it remains an open problem to design a token reduction method that is fast, maintains high performance, and is applicable to various vision tasks. In this work, we present a token pruner that uses auxiliary prediction heads that learn to select tokens end-to-end based on task relevance. These auxiliary heads can be removed after training, leading to throughput close to that of a random pruner. We evaluate our method on image classification, semantic segmentation, object detection, and instance segmentation, and show speedups of 1.5 to 4x with small drops in performance. As a best case, on the ADE20k semantic segmentation benchmark, we observe a 2x speedup relative to the no-pruning baseline, with a negligible performance penalty of 0.1 median mIoU across 5 seeds.