DiT: Efficient Vision Transformers with Dynamic Token Routing

TL;DR

DiT introduces a dynamic token routing strategy in vision transformers, enabling adaptive processing of image tokens based on object scale and visual features, leading to improved accuracy and efficiency across multiple vision tasks.

Contribution

The paper proposes a data-dependent token routing mechanism in vision transformers, allowing adaptive multi-path feature propagation and reduced computation with budget constraints.

Findings

Achieves state-of-the-art accuracy on ImageNet classification.

Demonstrates superior performance in object detection and segmentation tasks.

Offers a flexible backbone adaptable to various vision applications.

Abstract

Recently, the tokens of images share the same static data flow in many dense networks. However, challenges arise from the variance among the objects in images, such as large variations in the spatial scale and difficulties of recognition for visual entities. In this paper, we propose a data-dependent token routing strategy to elaborate the routing paths of image tokens for Dynamic Vision Transformer, dubbed DiT. The proposed framework generates a data-dependent path per token, adapting to the object scales and visual discrimination of tokens. In feed-forward, the differentiable routing gates are designed to select the scaling paths and feature transformation paths for image tokens, leading to multi-path feature propagation. In this way, the impact of object scales and visual discrimination of image representation can be carefully tuned. Moreover, the computational cost can be further reduced by giving budget constraints to the routing gate and early-stopping of feature extraction. In experiments, our DiT achieves superior performance and favorable complexity/accuracy trade-offs than many SoTA methods on ImageNet classification, object detection, instance segmentation, and semantic segmentation. Particularly, the DiT-B5 obtains 84.8\% top-1 Acc on ImageNet with 10.3 GFLOPs, which is 1.0\% higher than that of the SoTA method with similar computational complexity. These extensive results demonstrate that DiT can serve as versatile backbones for various vision tasks.