On enhancing the robustness of Vision Transformers: Defensive Diffusion

Self-Ensembling Vision Transformer (SEVIT) is a recent development in computer vision that combines the benefits of both Vision Transformers (ViT) and Self-Ensembling models. SEVIT utilizes the ViT architecture for image feature extraction and employs a self-ensembling technique to improve the model’s generalization and performance Figure 1. [4], The Self-Ensembling Vision Transformer (SEViT) proposed by the author makes use of the observation that the feature representations learned by the early blocks of a ViT are not easily impacted by adversarial perturbations. The SEViT approach involves training multiple classifiers based on these intermediate feature representations and combining their predictions with that of the final ViT classifier to improve resistance against adversarial attacks. Additionally, the consistency between the various predictions can be used to identify potential adversarial samples. Through experiments on chest X-ray and fundoscopy modalities, the effectiveness of the SEViT architecture is demonstrated in defending against various adversarial attacks in the gray-box scenario, where the attacker has knowledge of the target model but not the defense mechanism.

Adversarial attacks are a class of techniques that aim to fool machine learning models by adding imperceptible perturbations to the input data. Two popular adversarial attacks are the fast gradient sign method (FGSM) and the projected gradient descent (PGD) attack. The FGSM attack, proposed by [1], generates an adversarial example by adding a small perturbation to the input data in the direction of the sign of the gradient of the loss function with respect to the input. Despite its simplicity, the FGSM attack has been shown to be effective in generating adversarial examples for various image classification tasks. The PGD attack, proposed by [6], is an iterative version of the FGSM attack that generates adversarial examples by taking multiple small steps in the direction of the sign of the gradient of the loss function while constraining the perturbation to be within a certain bound. The PGD attack has been shown to be more effective than FGSM in generating adversarial examples that are harder for the model to classify correctly, even when the model has been trained to defend against FGSM attacks. Studies have explored different aspects of these attacks, such as their transferability across different models and datasets, their effectiveness against models trained with different defenses, and their impact on the interpretability and robustness of machine-learning models [7] [5]

Model diffusion is a technique that has been gaining increasing attention in the field of computer vision. It involves the propagation of information between different images, allowing for the transfer of style, texture, and other features from one image to another. One of the earliest papers on model diffusion was [8] in which the authors propose a technique for fine-grained image editing using diffusion models. Their approach involves propagating the information from a reference image to a target image in a pixel-wise manner, allowing for precise control over the editing process. The authors demonstrate that their technique can be used for tasks such as image inpainting, super-resolution, and style transfer. In [9], the authors propose a method for removing adversarial perturbations from images using diffusion models. Their approach involves first adding the Gaussian noise to adversarially perturb images and then using the denoising diffusion model to reconstruct back the image. The authors demonstrate that their method is effective in improving the robustness of deep neural networks against adversarial attacks. Other papers have also explored Face Recognition [10] and image segmentation. In [11] the authors explored the diffusion model by adding a different variation of noises including blurring, snowing, and masking, and showed that the diffusion model is still robust enough to reconstruct back the original image. Their method called ’cold diffusion’ is based on the heat diffusion equation and the concept of coldness or lack of randomness. The authors demonstrated the effectiveness of their method through experiments on several datasets, including MNIST and CIFAR-10. In [12], the authors proposed a method for adapting pre-trained vision-language models to medical imaging domains. They leveraged the transformer-based model’s ability to encode global contextual information and applied it to medical images by formulating the task as a multi-label classification problem. The authors demonstrated the effectiveness of their approach on several medical imaging datasets, achieving state-of-the-art results on some of them. This work highlights the potential of diffusion models in adapting pre-trained models to new domains, enabling the use of existing models for new applications.

Knowledge distillation is a process of training a smaller, faster, and more efficient model, known as the student model, to mimic the behavior of a larger and more accurate model, known as the teacher model. This technique has gained significant attention in recent years due to its ability to compress complex models, reduce their computational requirements, and accelerate their inference time. One of the earliest works on knowledge distillation was proposed by [13], where they introduced the concept of transferring the knowledge of a deep neural network to a smaller network by minimizing the difference between the outputs of the teacher and student models. In their work, they demonstrated that knowledge distillation could significantly improve the performance of a smaller network on various image classification tasks. Following this work, several studies have explored different aspects of knowledge distillation, such as the choice of the loss function, the impact of the teacher and student model architectures, and the effect of data augmentation techniques. One such work is the study by [14], where they proposed a novel loss function, called the attention transfer loss, which leverages the attention maps of the teacher model to guide the training of the student model. Their results showed that attention transfer loss could significantly improve the performance of the student model on various image classification tasks. Moreover, recent research has extended knowledge distillation to other domains, such as natural language processing, speech recognition, and object detection [15].

We hypothesis that (1) reverse diffusion process is effective in removing adversarial noise, (2) adversarial images generated specifically for the vision transformer would not be transferable to the student model (CNN). Faris et al. [4] conducted a study revealing that Vision Transformers (ViTs) lack robustness against white box attacks. In our research, we address this issue from a different perspective by leveraging the capabilities of a diffusion model to mitigate attack-agnostic noise. By exploiting the denoising properties of the diffusion model, we employ a reverse diffusion process to remove adversarial noise and obtain a clean image.

In a gray box scenario, where the attacker possesses comprehensive knowledge of the base model, including model weights, prediction probabilities, and training data, but lacks information about the defensive model, [4] propose the SEViT model. SEViT incorporates trained Multi layer Perceptions (MLPs) attached to each block of the vision transformer to enhance robustness. However, deploying this approach becomes infeasible in real-world due to computational and memory complexities, as each MLP block contains a substantial number of parameters (625M). To overcome this challenge, our methodology incorporates performing knowledge distillation to adopt a lightweight student CNN model trained on soft probabilities distilled from the teacher model (ViT). Our hypothesis aligns with the observations made by [17], who demonstrated that adversarial images exhibit higher transferability within the same family of model variants.

We implemented three variation of diffusion phases in diffusion model (where this diffusion phase consists of forward diffusion and reconstruction phases) on ViT. The three variation of diffusion phases are blurring, reconstruction and direct reconstruction. In the blurring variation, we take the image processed solely by the forward diffusion phase. In the reconstruction variation, we reconstruct the image from its blurred state back to a fine-grained image (See Figure 4). Moreover, the direct reconstruction involves performing the reconstruction directly without any iterations. The diffusion model acts as a defensive purifier added before the ViT that aims to remove the adversarial noise in the input image. Our proposed diffusion framework, with its diffusion phase, aims to purify the input image before feeding it into the original ViT module for classification. We performed the same implementation of combining defensive diffusion phase with SEViT instead of ViT with the aim to compare with the baseline performance of robustness in both the models. Our experiments yield several noteworthy observations, as outlined in Table 1. On the other hand, SEViT, which is originally built to be robust against attack samples, achieved significantly higher robust accuracy on attack samples

In order to improve robust accuracy while preserving clean accuracy of ViT, along with improved computational performance, we employed knowledge distillation to extract a student CNN model from the original teacher ViT. During the distillation process, we utilized the training set to train the student ViT, which generated soft targets for each input sample. These soft targets were then used to create a new dataset (X, Soft target) for training the distilled student CNN. A temperature parameter (set to 1 in our case) was employed to determine the degree of softness in the student model. We conducted the same defensive diffusion experiments on the student CNN as we did on ViT and SEViT. This involved implementing three diffusion variations as purifiers for the student CNN and subsequently testing it on both actual attack samples and diffused attack samples as shown in Table 2.

We implemented our proposed diffusion method on both ViT and SEViT models, and the results are presented in Figure 5. As a defender, our objective is to defend against attack samples while maintaining a robust accuracy close to the clean accuracy. This means that we want our model to accurately classify both clean and attack samples, ensuring correct predictions. In the case of original ViT, we observe that the inference using attack samples (Original) significantly decreases the clean accuracy to approximately 55%. However, when utilizing our proposed diffusion purifier with ViT, we achieve higher robust accuracy across all three diffusion variations. By leveraging the denoising properties of the diffusion model, we apply a reverse diffusion process that effectively removes the adversarial noise from the attack sample, resulting in a cleaner image. Among the three diffusion variations, the blurred variation performs the best, achieving robust accuracy of approximately 85% in the case of PGD attack. On the other hand, the reconstructed and direct reconstructed variations show lower robust accuracy, with around 61% robust accuracy in the PGD attack scenario. Therefore, we imply that blurred variation of diffusion phase of our approach is able to produce enhanced robust accuracy when tested with attack samples compared to raw ViT (i.e, ViT without diffusion).

Furthermore, SEViT demonstrates higher robust accuracy even without the use of diffusion phases, with clean accuracy of approximately 94% and attack accuracy of around 89% in the case of FGSM attack. This higher robustness is consistent across other types of attacks as well. However, when the diffusion phase is applied to SEViT, there is a decrease in robust accuracy across all three diffusion variations. This reduction in robust accuracy can be attributed to the introduction of over-smoothing caused by the diffusion phases, which results in the loss of important details and features necessary for robust classification. This over-smoothing effect ultimately leads to reduced discriminative power and slightly lower robust accuracy when SEViT is combined with the diffusion phase. When considering the best performing variation, which is the blurred variation among the three diffusion phases, we observe that when the blurred diffusion variation is applied to ViT, it achieves a robust accuracy that is comparable to SEViT (See Figure 6). SEViT utilizes an ensemble of MLPs, with each MLP containing 625 million parameters. This introduces computational complexities and makes practical deployment infeasible. In contrast, our blurred variation of the diffusion phase demonstrates almost similar robust performance to SEViT, with only a 1-2% lower difference in robust accuracy. This suggests that our approach of incorporating the blurred diffusion variation into ViT can be an efficient and robust solution for developing a ViT model that is both effective and practical. By using the blurred diffusion variation, we can achieve similar levels of robustness to SEViT without the computational complexity associated with the ensemble of MLPs.

We conducted a thorough comparison between the top-performing student CNN model and the teacher ViT model, specifically evaluating their performance on adversarial images and the output images obtained through defensive diffusion. The results, presented in the accompanying table, revealed a substantial decrease in accuracy for the ViT model when exposed to white box attacks. On the other hand, the CNN model consistently maintained a performance level comparable to that achieved with clean images. The accuracy of the ViT model witnessed a significant decline, plummeting from 88% to approximately 20% when confronted with white box attacks Figure 7 and 8. In contrast, the CNN model exhibited remarkable resilience, maintaining an accuracy rate of over 80% even under adversarial attack. These outcomes highlight the robustness of the CNN model and its ability to withstand adversarial perturbations while preserving its high accuracy. The comparison results demonstrate the student CNN model’s superior performance in the face of adversarial attacks, emphasizing its robustness and resilience compared to the ViT model. These findings contribute to our understanding of adversarial vulnerabilities across different model families and have implications for developing more robust defenses against adversarial attacks.

Furthermore, we proceeded to compare the accuracy of both models using the images generated from our method, specifically the blurred images. Surprisingly, the accuracy of the ViT model was effectively restored when evaluated on these blurred images. Notably, the CNN model exhibited robust performance, maintaining accuracy levels comparable to the ViT model. For instance, the ViT model’s accuracy was restored from 2% to over 80% Figure 8 when subjected to the AutoPGD attack, which is known to be particularly challenging. This outcome serves as compelling evidence of the efficacy of our approach. The restoration of accuracy in the ViT model when evaluated on blurred images indicates the success of our defensive diffusion technique. Moreover, the fact that the CNN model’s accuracy remains consistently high, even after undergoing the defensive diffusion process, suggests that the features learned by the CNN model are inherently more robust. This observation underlines the resilience of the CNN model in the face of adversarial perturbations and further validates the effectiveness of our proposed method.