Make A Long Image Short: Adaptive Token Length for Vision Transformers

TL;DR

This paper introduces an adaptive token length method for Vision Transformers that shortens long images during inference, significantly reducing computational costs while maintaining accuracy.

Contribution

We propose ReViT and TLA, enabling adaptive token length assignment at test time to accelerate ViT inference without sacrificing performance.

Findings

Significant reduction in inference time across multiple ViT models.

Maintained accuracy with fewer tokens during image processing.

Applicable to various vision transformer architectures.

Abstract

The vision transformer is a model that breaks down each image into a sequence of tokens with a fixed length and processes them similarly to words in natural language processing. Although increasing the number of tokens typically results in better performance, it also leads to a considerable increase in computational cost. Motivated by the saying "A picture is worth a thousand words," we propose an innovative approach to accelerate the ViT model by shortening long images. Specifically, we introduce a method for adaptively assigning token length for each image at test time to accelerate inference speed. First, we train a Resizable-ViT (ReViT) model capable of processing input with diverse token lengths. Next, we extract token-length labels from ReViT that indicate the minimum number of tokens required to achieve accurate predictions. We then use these labels to train a lightweight Token-Length Assigner (TLA) that allocates the optimal token length for each image during inference. The TLA enables ReViT to process images with the minimum sufficient number of tokens, reducing token numbers in the ViT model and improving inference speed. Our approach is general and compatible with modern vision transformer architectures, significantly reducing computational costs. We verified the effectiveness of our methods on multiple representative ViT models on image classification and action recognition.