ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT
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 and inter-layer dependencies that complicate optimization, 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), a training-free method that filters redundant noise channels at inference time. 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\%p) with 39.4\% FLOPs reduction. ToaSt also transfers effectively to diverse downstream tasks (COCO detection, ADE20K segmentation, CIFAR-100 classification), achieving 52.2 versus 51.9 mAP on COCO. Code: \href{https://github.com/SHANNonLab-HUFS/ToaSt}{github.com/SHANNonLab-HUFS/ToaSt}.
Lay Summary
Modern AI systems that understand images — from medical scans to self-driving cars to photo search — have become remarkably powerful, but also remarkably expensive to run. They demand massive computation, making them slow and costly to deploy on everyday devices like phones, drones, and cameras. We developed ToaSt, a method that makes these AI models smaller and faster while keeping their accuracy intact. Like trimming a tree by removing dead branches, ToaSt identifies parts of the model that contribute little to its decisions and removes them. Unlike most existing techniques, ToaSt requires no additional training, so it can be applied instantly to any pretrained model. Tested across nine different image-recognition models, ToaSt reduced computation by up to 40 percent without losing accuracy — and in several cases, the slimmer models became slightly more accurate, because removing noisy components helped the model focus on what matters. The method also transferred smoothly to other tasks, including object detection and scene segmentation. By cutting the cost of running AI vision systems, our work helps bring these technologies to more devices and more users, while reducing the energy footprint of increasingly large AI models.