Vision Transformer with Progressive Sampling

TL;DR

This paper introduces PS-ViT, a vision transformer with a progressive sampling strategy that adaptively locates discriminative regions, significantly improving accuracy and efficiency over vanilla ViT on ImageNet.

Contribution

The paper proposes a novel differentiable progressive sampling method integrated with ViT to better locate important regions, enhancing performance and reducing computational costs.

Findings

PS-ViT achieves 3.8% higher top-1 accuracy than vanilla ViT on ImageNet.

PS-ViT uses approximately 4 times fewer parameters.

PS-ViT requires about 10 times fewer FLOPs.

Abstract

Transformers with powerful global relation modeling abilities have been introduced to fundamental computer vision tasks recently. As a typical example, the Vision Transformer (ViT) directly applies a pure transformer architecture on image classification, by simply splitting images into tokens with a fixed length, and employing transformers to learn relations between these tokens. However, such naive tokenization could destruct object structures, assign grids to uninterested regions such as background, and introduce interference signals. To mitigate the above issues, in this paper, we propose an iterative and progressive sampling strategy to locate discriminative regions. At each iteration, embeddings of the current sampling step are fed into a transformer encoder layer, and a group of sampling offsets is predicted to update the sampling locations for the next step. The progressive sampling is differentiable. When combined with the Vision Transformer, the obtained PS-ViT network can adaptively learn where to look. The proposed PS-ViT is both effective and efficient. When trained from scratch on ImageNet, PS-ViT performs 3.8% higher than the vanilla ViT in terms of top-1 accuracy with about $4\times$ fewer parameters and $10\times$ fewer FLOPs. Code is available at https://github.com/yuexy/PS-ViT.