PAT-ViT: Token Pruning-based Adversarial Tuning for Robust Vision Transformers

Recent studies have demonstrated that Vision Transformers (ViTs) are vulnerable to adversarial attacks. While adversarial training is a recognized strategy for enhancing model robustness, it demands substantial computational resources during the training process. Furthermore, the high computational complexity of ViTs during inference presents further challenges. To address these challenges, we propose token pruning-based adversarial tuning for robust ViTs (PAT-ViT). PAT-ViT enhances the robustness of ViTs by reducing their predictability to attackers and tuning them using adversarial data. Unlike conventional methods, PAT-ViT reduces training overhead by tuning pre-trained ViTs rather than training ViTs from scratch. Moreover, PAT-ViT improves inference efficiency by pruning partial input image. Compared to the state-of-the-art (SOTA) robust method, PAT-ViT boosts robust accuracy by 7.5% while also achieving a 10.3% increase in clean accuracy. During inference, PAT-ViT reduces FLOPs by 0.7× relative to SOTA. Additionally, it reduces the training time by 7.5× to 13.1× compared to prior works.