ToFe: Lagged Token Freezing and Reusing for Efficient Vision Transformer Inference

TL;DR

This paper introduces ToFe, a novel framework for vision transformers that freezes and reuses tokens to significantly reduce computation while maintaining high accuracy, suitable for resource-limited devices.

Contribution

ToFe is the first method to freeze and reuse tokens dynamically in vision transformers, balancing efficiency and performance through end-to-end training.

Findings

Reduces LV-ViT computational cost by 50%.

Achieves less than 2% accuracy drop.

Outperforms existing token reduction methods.

Abstract

Although vision transformers (ViT) have shown remarkable success in various vision tasks, their computationally expensive self-attention hinder their deployment on resource-constrained devices. Token reduction, which discards less important tokens during forward propagation, has been proposed to enhance the efficiency of transformer models. However, existing methods handle unimportant tokens irreversibly, preventing their reuse in subsequent blocks. Considering that transformers focus on different information among blocks, tokens reduced in early blocks might be useful later. Furthermore, to adapt transformer models for resource-constrained devices, it is crucial to strike a balance between model performance and computational overhead. To address these challenges, in this paper, we introduce a novel Token Freezing and Reusing (ToFe) framework, where we identify important tokens at each stage and temporarily freeze the unimportant ones, allowing their lagged reusing at a later stage. Specifically, we design a prediction module for token identification and an approximate module for recovery of the frozen tokens. By jointly optimizing with the backbone through computation budget-aware end-to-end training, ToFe can adaptively process the necessary tokens at each block, thereby reducing computational cost while maintaining performance. Extensive experiments demonstrate that ToFe reduces the computational cost of LV-ViT model by 50% with less than 2% drop in Top-1 accuracy, achieving a better trade-off between performance and complexity compared to state-of-the-art methods.