Token Cropr: Faster ViTs for Quite a Few Tasks

Benjamin Bergner, Christoph Lippert, Aravindh Mahendran
12/1/2024
cs.CVcs.LG

No official code found. Best implementation: benbergner/cropr (31 stars, MIT, last commit 2025-07-21).

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.

View on arXivView PDF

Code Implementations(1)

benbergner/croprCommunity100%
MIT

A token pruning method that accelerates ViTs for various tasks while maintaining high performance.

313PythonNov 22, 20241 years agoMIT
attention-mechanismcomputer-visiondeep-learningneural-networkstoken-pruning+1 more

Similar Papers

Cite this paper

@article{bergner2024token,
  title  = {Token Cropr: Faster ViTs for Quite a Few Tasks},
  author = {Benjamin Bergner and Christoph Lippert and Aravindh Mahendran},
  year   = {2024},
  eprint = {2412.00965},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2412.00965v1},
  journal = {CVPR 2024 2024}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2412.00965.svg)](https://codeofpaper.com/paper/2412.00965)

Discussion