SAC-ViT: Semantic-Aware Clustering Vision Transformer with Early Exit

TL;DR

SAC-ViT tackles the ViT efficiency bottleneck by introducing a two-stage framework with an Early Exit on downsampled inputs and a non-iterative Semantic-Aware Clustering stage that partitions tokens into target and non-target groups based on semantic importance. Target tokens are upscaled and processed with localized self-attention, while non-target tokens are reused, reducing the effective token set and computation from $O(N^2D)$ toward a clustered, lower-cost regime. The model is trained end-to-end with a CE loss on SAC outputs and KL supervision from the EE stage, enabling adaptive compute via a threshold $oldsymbol{f}$. Empirical results on ImageNet show SAC-ViT significantly reduces FLOPs (e.g., ~62% at certain thresholds) and improves throughput (up to ~2x) with minimal or improved accuracy compared to strong ViT baselines, indicating practical impact for resource-constrained deployment.

Abstract

The Vision Transformer (ViT) excels in global modeling but faces deployment challenges on resource-constrained devices due to the quadratic computational complexity of its attention mechanism. To address this, we propose the Semantic-Aware Clustering Vision Transformer (SAC-ViT), a non-iterative approach to enhance ViT's computational efficiency. SAC-ViT operates in two stages: Early Exit (EE) and Semantic-Aware Clustering (SAC). In the EE stage, downsampled input images are processed to extract global semantic information and generate initial inference results. If these results do not meet the EE termination criteria, the information is clustered into target and non-target tokens. In the SAC stage, target tokens are mapped back to the original image, cropped, and embedded. These target tokens are then combined with reused non-target tokens from the EE stage, and the attention mechanism is applied within each cluster. This two-stage design, with end-to-end optimization, reduces spatial redundancy and enhances computational efficiency, significantly boosting overall ViT performance. Extensive experiments demonstrate the efficacy of SAC-ViT, reducing 62% of the FLOPs of DeiT and achieving 1.98 times throughput without compromising performance.

Figures (5)

• Figure 1: Comparing our SAC-ViT with state-of-the-art adaptive ViT optimization methods, SAC-ViT achieves better efficiency/Top-1 accuracy trade-off. SAC-ViT, DVT wang2021not, CF-ViT chen2023cf, and LF-ViT hu2024lf are all build up DeiT touvron2021training. $*$ indicates that the input resolution is $288\times 288$.
• Figure 2: Overview of the SAC-ViT framework. SAC-ViT consists of an Early Exit (EE) stage and a non-iterative Semantic-Aware Clustering (SAC) stage. In the EE stage, downsampled images are processed to extract global semantic information and generate initial results. If these results don't meet the EE terminate criteria, the information is clustered into target and non-target tokens. In the SAC stage, target tokens are mapped back to the original image, cropped, embedded, and then combined with reused non-target tokens from the EE stage. Multi-Head Self-Attention (MHSA) is applied within each cluster. Notably, SAC-ViT uses the same network parameters in both stages and performs end-to-end optimization.
• Figure 3: Performance analysis of removing each of designs.
• Figure 4: Quantitative analysis of SAC-ViT's early exit stage and semantic-aware clustering stage.
• Figure 5: Qualitative analysis of SAC-ViT's early exit and semantic-aware clustering stages reveals that images recognized during the early exit stage generally have simple backgrounds, with the objects of interest clearly visible throughout the image. In contrast, during the SAC stage, SAC-ViT correctly identifies objects with complex backgrounds by clustering the target and background based on semantic information, followed by performing local self-attention to enhance recognition efficiency. (The shaded and non-shaded regions of the SAC stage indicate non-target and target clustering, respectively.)