ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT

TL;DR

ToaSt introduces a decoupled framework for efficient Vision Transformer pruning, combining head-wise structured pruning and token channel selection to significantly reduce computational costs while maintaining high accuracy.

Contribution

It proposes a novel decoupled approach applying specialized strategies to different ViT components, improving efficiency and robustness over existing methods.

Findings

Achieves 39.4% FLOPs reduction on ViT-MAE-Huge with 88.52% accuracy.

Outperforms baselines in accuracy-efficiency trade-offs across nine models.

Transfers effectively to downstream tasks like object detection.

Abstract

Vision Transformers (ViTs) have achieved remarkable success across various vision tasks, yet their deployment is often hindered by prohibitive computational costs. While structured weight pruning and token compression have emerged as promising solutions, they suffer from prolonged retraining times and global propagation that creates optimization challenges, respectively. We propose ToaSt, a decoupled framework applying specialized strategies to distinct ViT components. We apply coupled head-wise structured pruning to Multi-Head Self-Attention modules, leveraging attention operation characteristics to enhance robustness. For Feed-Forward Networks (over 60\% of FLOPs), we introduce Token Channel Selection (TCS) that enhances compression ratios while avoiding global propagation issues. Our analysis reveals TCS effectively filters redundant noise during selection. Extensive evaluations across nine diverse models, including DeiT, ViT-MAE, and Swin Transformer, demonstrate that ToaSt achieves superior trade-offs between accuracy and efficiency, consistently outperforming existing baselines. On ViT-MAE-Huge, ToaSt achieves 88.52\% accuracy (+1.64 \%) with 39.4\% FLOPs reduction. ToaSt transfers effectively to downstream tasks, achieving 52.2 versus 51.9 mAP on COCO object detection. Code and models will be released upon acceptance.