Rect-ViT: Rectified attention via feature attribution can improve the adversarial robustness of Vision Transformers

Deep neural networks (DNNs) have suffered from input perturbations and adversarial examples (AEs) for a long time, mainly caused by the distribution difference between robust and non-robust features. Recent research shows that Vision Transformers (ViTs) are more robust than traditional convolutional neural networks (CNNs). We studied the relationship between the activation distribution and robust features in the attention mechanism in ViTs, coming up with a discrepancy in the token distribution between natural and adversarial examples during adversarial training (AT). When predicting AEs, some tokens irrelevant to the targets are still activated, giving rise to the extraction of non-robust features, which reduces the robustness of ViTs. Therefore, we propose Rect-ViT, which can rectify robust features based on class-relevant gradients. Performing the relevance back-propagation of auxiliary tokens during forward prediction can achieve rectification and alignment of token activation distributions, thereby improving the robustness of ViTs during AT. The proposed rectified attention mechanism can be adapted to a variety of mainstream ViT architectures. Along with traditional AT, Rect-ViT can also be effective in other AT modes like TRADES and MART, even for state-of-the-art AT approaches. Experimental results reveal that Rect-ViT improves average robust accuracy by 0.64% and 1.72% on CIFAR10 and Imagenette against four classic attack methods. These modest gains have significant practical implications in safety-critical applications and suggest potential effectiveness for complex visual tasks.