ProTransformer: Robustify Transformers via Plug-and-Play Paradigm

Authors: zhichao  hou
weizhi  gao
yuchen  shen
feiyi  wang
xiaorui  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:33:35

Abstract: Transformer-based architectures have dominated various areas of machine learning in recent years. In this paper, we introduce a novel robust attention mechanism designed to enhance the resilience of transformer-based architectures. Crucially, this technique can be integrated into existing transformers as a plug-and-play layer, improving their robustness without the need for additional training or fine-tuning. Through comprehensive experiments and ablation studies, we demonstrate that our ProTransformer significantly enhances the robustness of transformer models across a variety of prediction tasks, attack mechanisms, backbone architectures, and data domains. Notably, without further fine-tuning, the ProTransformer consistently improves the performance of vanilla transformers by 19.5\%, 28.3\%, 16.1\%, and 11.4\% for BERT, ALBERT, DistilBERT, and RoBERTa, respectively, under the classical TextFooler attack. Furthermore, ProTransformer shows promising resilience in large language models (LLMs) against prompting-based attacks, improving the performance of T5 and LLaMA by 24.8\% and 17.8\%, respectively, and enhancing Vicuna by an average of 10.4\% against the Jailbreaking attack. Beyond the language domain, ProTransformer also demonstrates outstanding robustness in both vision and graph domains.

Main paper: 
 introduction: 
 In recent years, attention mechanisms and transformer-based architectures have drawn significant attention across many domains in machine learning, such as natural language processing (NLP) [1, 2] , computer vision [3, 4] , and graph learning [5, 6] . In particular, transformers have demonstrated superior capabilities to learn and model complex relations in data through powerful and universal attention mechanisms, and they have dominated many popular NLP tasks such as topic classification, sentiment analysis, textual entailment, machine translation, dialogue generation, etc [2] . Despite their success in NLP and beyond, many recent studies have demonstrated that transformers are highly vulnerable to adversarial attacks such that even small modifications to the input can easily fool the model [7, 8, 9] . However, most research on transformer architectures focuses on accuracy and efficiency, largely ignoring their security and robustness [10, 11] .With the increasing popularity of Large Language Models (LLMs) [12, 13] , the robustness and security concerns of transformer architectures become particularly of interest. It has been shown that malicious attackers can invade the language models through various approaches as shown in Figure 1 . The attacker can modify the input content in text attacks [14] or the prompt template in Figure 1 : Various attack mechanisms on language models. Classic text attacks modify the input content using typos or synonyms; Prompt attacks perturb the prompt template within the input; and Jailbreaks append adversarial, non-semantic suffixes to manipulate the model into producing malicious outputs.prompt attacks to mislead the model predictions [15] . Moreover, by adding adversarial suffixes, the jailbreaking attack [16] can prompt a LLM to generate toxic and illegal content which could lead to catastrophic legal and ethical impacts such as malicious speech or privacy leaks. Given the broad applications of transformers and their vulnerabilities under attacks, it is imperative to design a universal and effective strategy to enhance the robustness of transformers. Existing research attempting to improve the robustness of transformers can be roughly divided into empirical defenses [17, 18, 19, 20, 21] and certifiable defenses [22, 23, 24, 25] . Nevertheless, these defenses require excessive computation costs for training, inference, or both. In addition to these architecture-agnostic defenses, there are also several works proposing to enhance the robustness of transformers architecture [26, 27, 28, 29] . However, these approaches either require substantial computations or rely on specific domain knowledge, which hinders their extensions to larger models or broader application domains.In this paper, given the limitations of existing works and the enormous training cost of transformers, we aim to robustify transformer architectures via a plugand-play paradigm without additional training or finetuning. Our proposed ProAttention (Algorithm 1) can be readily plugged into the given transformers to convert them to ProTransformer (as shown in Figure 2 ) with significantly stronger robustness. Specifically, our contributions can be summarized as follows:• We establish a novel connection between the attention mechanism in transformers and the weighted least square estimator. We provide interpretation and numerical simulation to reveal its vulnerability against potential adversarial attacks.• From our new perspective, we propose robust token estimators to improve the resilience of token aggregation against adversarial attacks. We also propose an efficient Newton-IRLS algorithm to approximate the non-convex and non-smooth robust token estimator with convergence guarantees. The derived algorithm can be plugged into the given transformer as a plug-and-play layer to enhance its robustness against attacks even without additional training or fine-tuning.• Our comprehensive experiments and ablation studies demonstrate that the proposed ProTransformer is effective, efficient, and generalizable. It significantly improves the robustness of transformers across various machine learning tasks, attack mechanisms, backbone architectures, and data domains such as language, vision, and graphs. 
related work: 
 In this section, we mainly summarize related works on the attacks and defenses of transformers focusing on language domains since this is the focus of this paper.Attacks. Compared to the attack mechanisms in vision domain [30, 31] , the text attacks in the language domain are highly complicated due to the natural irregularity of data structure. According to the perturbation units, text attacks can be classified into character-level [7, 32] , word-level [33, 34, 35, 36, 9, 14, 37] , sentence-level [38] , and multi-level [39, 40, 8] . These classic text attacks typically generate adversarial examples through misspellings, synonym replacement, etc. In the era of LLMs, several new types of attacks have emerged, such as jailbreak attacks [16, 41, 42, 43] and prompt injection [44, 45, 46] . These prompting-based attacks aim to trick models into generating unsafe outputs using adversarially crafted prompts.Defenses. There have been some works proposed to defend against adversarial text attacks from various perspectives. Empirical defenses, such as data augmentation [17] and adversarial training [47, 18, 19, 20, 21] , attempt to robustify models by exposing them to a wider range of adversaries during training. On the other hand, several certifiable defenses [24, 25, 22, 23] have been proposed to guarantee the model robustness regardless of the attacks. However, these defenses require excessive computation costs for training, inference, or both, which limits their application in large-scale problems such as LLMs. Besides, all these methods are typically architecture-agnostic, which are orthogonal to and can be combined with our proposed defenses on the transformer architecture to further enhance the robustness.To safeguard the transformers, several endeavors have been made from the transformer architecture perspective. [26] modifies the attention mechanism and position embedding to robustify text-to-speech transformers. In the crisis detection and recognition task, [27] proposes an end-to-end attention-based classifier to enhance robustness. For tabular data, TableFormer [28] adopts structural-aware table-text encodings that are more robust to row and column order perturbations. However, these architectures are tailored for specific tasks, which require specific domain knowledge and can not be generalized across tasks. [29] proposes a general framework for self-attention modules via robust kernel density estimation (RKDE). However, this method introduces excess computation cost and shows relatively limited robustness improvement. Generally speaking, existing approaches either require substantial computations or rely on specific domain knowledge, which hinders their extensions to larger models or broader application domains. 
protransformer: 
 The main goal of this paper is to design robust self-attention mechanisms that are more resilient to adversarial attacks so they can be applied to robustify Transformer architectures. In this section, we first provide a new interpretation of the self-attention mechanism in Transformer architecture as the weighted least-square token estimator in Section 3.1. Then we propose robust token estimators that are more resilient to the dominating impact of input tokens in Section 3.2. An efficient Newton-IRLS algorithm is derived with a convergence guarantee to approximate the robust token estimator in Section 3.3. Finally, we describe how the proposed algorithm can be unrolled as robust attention layers to enhance the robustness of transformer architectures in Section 3.4.First, we provide a new perspective to formulate the vanilla attention mechanism as the weighted least squares (WLS) token estimator. In the self-attention layer, each output token z aggregates the values of input tokens {v j } as their weighted sum according to the attention weights: z = N j=1 a j v j , where {a j } j∈[N ] are the attention weights and {v j } j∈[N ] are value vectors for all N input tokens. This weighted sum can be interpreted as the optimal solution of the following weighted least squares (WLS) error minimization problem:EQUATIONwhose first-order optimality condition (∇L(z)= N j=1 a j • 2(z -v j ) = 0) yields z = N j=1 a j v j .Vulnerability analysis of vanilla attention. When adversaries perturb the input tokens, these tokens will dominate the impact on output tokens since the quadratic penalty on the residual ∥v j -z∥ 2 will dominate the WLS estimator. Therefore, the output token z will be shifted to those dominating input tokens. As a result, the adversarial input tokens will significantly impact the representation of output tokens. We also provide an empirical study to verify that adversarial attacks will significantly increase the residual ∥v j -z∥ 2 in Appendix F.3. Moreover, we simulate a mean estimation problem under outlier data points using synthetic data to better illustrate the sensitivity of the WLS estimator. The detailed setting and visualization results of the numerical simulation are provided in Appendix F.The analysis above provides a valid explanation of why various attention-based transformer architectures are easily compromised by introducing adversarial perturbations in the input data. Also, our interpretation of the attention mechanism in transformers as WLS estimator provides a rigorous perspective to design robust alternatives. To dampen the effect of outlier data, multiple robust regression algorithms have been proposed in robust statistics using least absolute deviations [48] , Huber regression [49] , and Minimax Concave Penalty (MCP) [50] . Motivated by these advancements with rigorous robustness guarantees, we propose the robust weighted least squares token estimators to enhance the resilience against potential adversarial attacks as follows: EQUATIONwhere ρ can be flexibly replaced with the specific robust penalties in Figure 3 .Special cases of ρ.(1) The quadratic ℓ 2 loss recovers vanilla WLS estimator; ℓ 1 loss exerts linear effect on the residuals; (2) Huber loss performs as ℓ 2 loss within the range (0, δ), and becomes similar to ℓ 1 when z > δ;(3) MCP loss behaves like ℓ 1 loss near zero and becomes constant when z is large than γ. (4) We also propose Huber-MCP to combine the advantage of Huber and MCP loss.The detailed formulations are available in Appendix B.4 due to the space limit.The proposed robust token estimator in Eq. ( 2) is non-convex and non-smooth, posing a challenge for efficient algorithm design. Moreover, the exploding model size of evolving transformers further necessitates the design of efficient neural network layers. To this end, we propose an efficient Newton iterative reweighted least square (Newton-IRLS) algorithm to tackle this challenging problem. We first design a localized upper bound for the original objective and then optimize the upper bound with a second-order Newton method. We also provide a rigorous theoretical loss descent guarantee. The precise statements are presented as follows and the detailed proof are provided in Appendix B.Localized upper bound. Instead of directly optimizing the original loss function L(z) in Eq. ( 2), we optimize a convex localized upper bound at the current iteration z (k) as follows:Lemma 3.1 (Localized Upper Bound). Suppose the loss objective is defined as in Eq. (2), where ρ • sqrt(•) is any non-convex function. For any fixed point z (k) , there exists a convex localized upper bound as:EQUATIONwhere w(k) j = ρ ′ (∥vj -z (k) ∥) 2∥vj -z (k)∥ and ρ ′ is the first derivative of ρ. Particularly, the constant C(z (k) ) guarantees the equality of L and L at z (k) , i.e., L(z (k) ) = L(z (k) ).Proof. Please refer to Appendix B.1.As C(z (k) ) is treated as a constant during the optimization at the current step, the upper bound in Eq. ( 3) becomes convex and can be efficiently optimized.Newton-IRLS iteration. After obtaining the convex upper bound L in Eq. ( 3), we can derive a concise closed-form iteration using the second-order Newton method as follows:EQUATIONEq. ( 4) can be interpreted as a reweighted sum, in which the derived w (k) j modifies the original attention score a j on the value vector v j . We leave detailed derivations of Newton-IRLS algorithm in Appendix B.2. Its convergence and rigorous loss descent are guaranteed by the following Theorem 3.2. Theorem 3.2 (Convergence guarantee). Suppose the loss objective L(z) is defined as in Eq. (2) and its corresponding convex localized upper bound is in Eq. (3). Then, through the iteration in Eq. (4), the following inequality holds:EQUATIONthat is, optimizing upper bound L can guarantee the rigorous descent of L.Proof. Please refer to Appendix B.3.Although the loss L(z) is not necessarily convex and does not possess a global optimum, Theorem 3.2 guarantees that the Newton-IRLS iteration, which optimizes L(z), can rigorously reduce the original loss L(z). The algorithm analyses in Appendix F, along with the main experiments in Section 4 and Section 5, validate that the local optimal solution achieved by our algorithm performs well in terms of both convergence and empirical robustness.Robust token estimator by reweighting the tokens. The robust estimator in Eq. ( 2) provides a general framework that covers several special cases. By choosing different penalty functions ρ on the residuals ∥v j -z (k) ∥, we obtain various reweighting schemes in Eq. ( 4). Take the MCP function as the instance, the weight is derived as w(k) j = ρ ′ γ (∥vj -z (k) ∥) 2∥vj -z (k) ∥ = max 1 ∥vj -z (k) ∥ -1 γ , 0 . Obviously, the weight w (k) jbecomes smaller as ∥v j -z (k) ∥ increases, thereby down-weighting the large residuals. The residuals will be completely removed when it exceeds the threshold γ, since the weight then becomes 0. The complete discussions for all cases are provided in Appendix B.4.In the previous subsection, we formulate the token-wise Newton-IRLS approach for notation simplicity. Here, we will present the corresponding matrix version for the entire attention layer.Matrix Form. Denote V = {v j } j∈[N ] and A = {a ij } i,j∈[N ] are value matrix and the attention matrix, respectively.Z (k) = {z (k) i } i∈[N ]is the estimator for token i at the k-th iteration. Subsequently, the pairwise distance k) and V can be efficiently computed using the torch.cdist function in PyTorch. Following this, the weightD (k) = {∥v j -z (k) i ∥} i,j∈[N ] between Z (W (k) = {w (k) ij } i,j∈[N ]can be calculated element-wise based on D (k) . Then the next step Z (k+1) is updated as a reweighted matrix multiplication (W (k) ⊙ A) • V.Plug-and-Play Robust Attention. The proposed algorithm can be packaged as a robust attention module, which can be readily plugged into the transformers as a Plug-and-Play Robust Attention (ProAttention) layer without additional training or fine-tuning as shown in Figure 2 . The implementation of Algorithm 1 ProAttention (MCP)1 D = torch.cdist(Z, V) # Pairwise distance 2 W = torch.clip(1/D-1/gamma, min=0) # MCP 3 W = normalize(W * A, p=1, dim=-1) 4 Z = torch.matmul(W, V) # UpdateProAttention using MCP penalty in PyTorch is shown in Algorithm 1. The complete pseudocode for other penalties is presented in in Appendix A.Complexity analysis. Let N , D, and K represent the length of tokens, the dimension of vectors, and the steps of the iterations, respectively. The vanilla attention requires 2 • N × N × D basic operations while our ProAttention needs (1 + 2K) • N × N × D. However, our ProAttention remains efficient, as the Newton-IRLS method can effectively approximate the solution within only 3 steps (K ≤ 3) (Figure 4 (a)) and ProTransformers do not introduce additional computation for training or fine-tuning. We provide the detailed complexity analysis of various attentions in Appendix L.Advantages. Our proposed ProAttention enjoys the following advantages: (1) Simplicity: it is simple and easy to implement with only 4 core lines of code in Algorithm 1; (2) Efficiency: it is a plug-andplay layer that can be integrated into any trained transformer without additional training or fine-tuning;(3) Universality: it is a universal framework that advances the vanilla attention mechanism into a series of robust derivatives with different penalties. Moreover, it can be applied to any attention-based model across various modalities and tasks.In the following sections, we will present comprehensive experiments and studies to validate the effectiveness of the proposed ProAttention on language modeling in Section 4 as well as computer vision and graph learning in Section 5. 
experiment on language modeling: 
 In this section, we evaluate the effectiveness of the proposed ProAttention and ProTransformer under classic text attacks on pre-trained language models, and two prompting-based attacks (prompt attack and jailbreak attack) in the context of LLMs with comprehensive ablation studies.Tasks and Datasets. For topic classification, we use AG's News Corpus (AGNEWS) [51] . For sentiment analysis, we utilize two widely-used datasets: Internet Movie Database (IMDB) [52] and Stanford Sentiment Treebank (SST-2) [53] . For textual entailment, we make use of Recognizing Textual Entailment (RTE) in the General Language Understanding Evaluation benchmark [54] . For jailbreak attack, we select a new dataset Behaviors introduced in [55] . For the detailed information on these datasets, please refer to Appendix C. Backbone Architectures. For classical pre-trained language models, we choose BERT [56] and its variants including RoBERTa [57] , ALBERT [58] and DistilBERT [59] . For large language models (LLMs), we choose T5 [60] , LLaMA [12] and Vicuna [13] . For the detailed information on backbone architectures, please refer to Appendix D.2.We not only evaluate several classic text attacks but also include popular prompt attacks and jailbreak attacks on the LLMs. The three attack mechanisms and their differences are illustrated in Figure 1 . For classic text attacks, we evaluate the attacks at various levels, including the characterlevel DeepWordBug [7] , word-level PWWS [9] , TextFooler [14] , and multi-level TextBugger [8] . For prompt attacks, we modify the prompt template according to the aforementioned text attacks following the evaluation setting in PromptBench [15] . For jailbreak, we evaluate the suffix attack using Greedy Coordinate Gradient (GCG) method [55] and we test both attacks transferred from surrogate model Vicuna (transfer attack) and attacks directly targeting the victim models (adaptive attack). Please refer to Appendix E for details on attacks.Defense Baselines. We include the following defense baselines in our experiments: MixADA [17] , PGD-Adv [31] , FreeLB [47] , TA-VAT [18] and SmoothLLM [61] . Additionally, we also include the adversarial training (AT), wherein the augmented perturbations are generated by the attack to be assessed. Details of these defense methods are provided in Appendix D.1.Evaluation metrics. Following [62] , we use 3 metrics to evaluate the model performance. Clean accuracy (Clean%) is the model accuracy on the clean testing data. Accuracy under attack (AUA%) is the accuracy on the perturbed data under specific attack. Attack success rate (ASR%) is the ratio of the number of successfully perturbed cases divided by the number of attempted texts.Hyperparameters. For text attack setting, we follow the setting in the TextAttack framework [63] . For prompt attack, we follow the setting in PromptBench [15] . For GCG-based jailbreak attack, we follow the setting in [61] . The detailed attack settings can be found in Appendix E. For defense baselines, we follow the settings in their original papers. For our ProTransformer, we set the default number of ProAttention layers as K = 3 since it can quickly converge to a reasonable precision within 3 layers. Finally we tune δ (default 1) or γ (default 4) in the penalties (Huber and MCP loss) to obtain the optimal parameters.To demonstrate the effectiveness of the proposed ProTransformer, we compare the robustness of our methods with several popular defenses in three classical tasks: topic classification, sentiment analysis, and textual entailment. Convergence. To validate the advantage of our Newton-IRLS over the first-order method, we conduct a simulation experiment and plot the loss descent curves in Figure 4 (a). It can be observed that Newton-IRLS exhibits efficient convergence as claimed in Section 3.3. We provide the experiment details, loss descent curves (Figure 7 ), and the visualization of trajectories (Figure 8 ) of the updated vectors in 2D plane in Appendix F.1 to further demonstrate the effectiveness of our algorithm. Attack constraints. In text attack, there are several kinds of attack constraints including the maximum percentage of perturbed words, minimum cosine similarity between the replaced synonym and original word, and minimum sentence similarity threshold between the original sentence and perturbed sentence. We test the values of these constraints in TextFooler. We present the results under different perturbation percentages in Figure 4 (c) and other constraint measurements in Appendix G.4.From the results, we observe that our Pro-ALBERT (MCP) can significantly outperform the backbone ALBERT across all ranges of constraints.Different penalties. Our Newton-IRLS is flexible to be formulated as different robust estimators with different penalties. From the comparison in Figure 4 (d) , it can be observed that our robust framework can consistently improve the robustness of the backbone BERT (ℓ 2 ). Specifically, ℓ 1 and Huber-based defenses are comparable, and MCP-based method exhibits the best performance.Different backbones. Our method is a general plug-and-play layer applicable to various transformer backbones. The results in Table 1 and the ablation study on different backbones in Appendix G.5 (Figure 12 ) demonstrate that ProAttention improves the robustness over various architecture backbones (BERT, RoBERTa, DistilBERT and ALBERT) against various attacks with significant margins. Running time. To empirically evaluate the efficiency of our method, we test the average running time on AGNEWS using BERT and Pro-BERT (MCP) equipped with multi-layer (K) ProAttention. The results in Table 2 show that our ProAttention only requires 1-2 times additional inference time of the backbone model yet achieves significant improvement in robustness without training.In the context of prompt-based generative AI, the adversarial attacks mechanisms on LLMs become more enriched and sophisticated. In this section, we will evaluate the robustness of our proposed ProTransformer under two popular attacks: prompt attack and jailbreak attack.As shown in Figure 1 , the most significant distinction between prompt attacks and classical text attacks is that prompt attacks aim to mislead the models by altering the prompt template rather than the input content. We display the results of T5 in Figure 5 and leave the comprehensive study in Appendix H.1. We also present the results on LLaMA in Appendix H.2. From the results, we can make the following observations: (1) For T5, the choice of the penalty would affect the robustness of defenses. Specifically, Pro-T5 (MCP) exhibits a significant advantage over other methods, and this advantage becomes even more evident as the number of perturbed words increases. Pro-T5 (ℓ 1 ) and Pro-T5 (Huber) show a slight improvement over the backbone model T5.(2) For LLaMA, Huber-MCP and Huber-based methods exhibit better robustness than other methods while preserving good clean performance. The detailed experiments and discussions can be found in Appendix H.2. In recent years, prompts have played a pivotal role in guiding models to generate desired outputs. Nevertheless, there exist malicious "jailbreak prompts", which are intentionally designed to bypass the built-in safeguards in LLMs, causing the model to produce harmful content that violates the legal policies. As illustrated in Figure 1 , the suffix-injection jailbreaks attempt to append a non-semantic suffix to the user's prompt to fool the models. We select GCG method to evaluate the resilience of models comprehensively.In Figure 6 , we compare the Attack Success Rates (ASRs) of Vicuna and its corresponding Pro-Vicuna (Huber) with various δ values on Behaviors. In each column, we also include SmoothLLM [61] with different smoothing extent q(%) to further reinforce the resilience of every single model. The last row of matrix (q = 0) stands for the performance without random smoothing. The additional results of random smoothing with swap, insert and patch, as well as the results under adaptive jailbreaking attack are presented in Appendix I.From the results, we can observe that: (1) Our Pro-Vicuna can significantly improve the robustness of Vicuna. As shown in the last row of Figure 6 , with δ = 0.1, we successfully reduce the ASR to 1.8%, which is comparable to the random smoothing defense that requires multiple random perturbations, inferences and aggregations. (2) Our ProAttention is orthogonal to randomized smoothing defense and can be combined with it to further improve the robustness.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We provide the detailed implementation in Section 3.4 and experimental setting in Section 4.1.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
experiment beyond language modeling: 
 In the previous section, we have provided comprehensive experiments to validate the effectiveness of our ProTransformer in the (large) language models. In fact, as shown in Figure 2 , our ProAttention is a fundamental module which can reinforce any attention-based models across various domains or modalities. In this section, we will integrate ProAttention into vision models and graph learning models to further validate the effectiveness and generality of our approach. In computer vision, we conduct two attacks (FGSM [30] and PGD [31] ) on several vision transformers including ViT, BeiT, ConviT, DeiT and Swin. We perform the experiments on CIFAR-10 and ImageNet-1K across budgets {1/255, 4/255, 8/255}, and present the results of PGD on CIFAR-10 in Table 3 and additional experiments in Appendix J. From the results, we can observe that Pro-ViT can outperform the second best model by {35.64%, 46.28%, 33.14%} under different budgets. SoftMedian 85.0 ± 0.7 75.5 ± 0.9 69.5 ± 0.5 62.8 ± 0.8 58.1 ± 0.7 GAT 83.5 ± 0.5 71.2 ± 1.2 65.0 ± 0.9 60.5 ± 0.9 56.7 ± 0.9 Pro-GAT (ours) 84.6 ± 0.8 75.5 ± 0.8 72.1 ± 0.4 69.0 ± 0.7 66.5 ± 1.2Besides the language and vision domains, we also validate the effectiveness of our method in the graph domain. We conduct the semisupervised node classification task and leverage PGD adaptive attack [64] to evaluate the robustness of models. We show the experiment results of Cora-ML and Citeseer, averaged over 5 different random splits, in Table 4 and Table 29 (in Appendix K), respectively. The ablation studies on the layers and γ in MCP are presented in Table 30 . Please refer to Appendix K for more detailed results and studies. From the results, we can conclude that our Pro-GAT significantly outperforms the backbone GAT and exhibits strong robustness across various budgets while keeping good clean accuracy.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [No] Justification: We provide the detailed implementation in Appendix A and the dataset information in Section 4.1. We will organize the data and code access after the submission.• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
conclusion & limitation: 
 In this paper, we delve into the robustness and security of the popular transformer-based architectures. We revisit the vulnerability of attention mechanisms with theoretical understanding and simulations. We propose an interpretable robust attention layer to robustify transformer architecture via a plugand-play paradigm. Our proposed ProAttention is an effective, efficient, and universal framework that can significantly enhance the robustness of transformers across various tasks, architectures, attacks, and domains without additional training or fine-tuning.Regarding the limitations, despite the acceptable complexity of our ProTransformer, there is still potential to improve the efficiency of our models. Additionally, while we primarily claim and validate the effectiveness of our model under a plug-and-play paradigm, we are excited about the future of the proposed ProTransformer architecture and hope to see its full potential realized through training or fine-tuning on large models in the future. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We include the theoretical computation analysis in Appendix 3.4 and the running time in Section 4.2.2.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). • The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: Our paper does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Our paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Here, we provide the complete pseudocode of our ProAttention including various penalties cases in Algorithm 2. The core itertions are show in the for loop in the code. Our ProAttention is easy to implement by only replacing the vanilla attention module with our ProAttention. ϕ(z) ≤ ϕ(z 0 ) + ϕ ′ (z 0 )(z -z 0 ) = ϕ ′ (z 0 ) • z + C(z 0 )where the first inequality holds with equality at z = z 0 andϕ ′ (z 0 ) = ρ ′ ( √ z) • 1 2 √ z z=z0 = ρ ′ ( √ z 0 ) 2 √ z 0 .By replacemnet as z = ∥v j -z∥ 2 and z 0 = ∥v j -z (k) ∥ 2 , thenρ(∥v j -z∥) ≤ ρ ′ (∥v j -z (k) ∥) 2∥v j -z (k) ∥ • ∥v j -z∥ 2 + C(∥v j -z (k) ∥ 2 ) = w (k) j • ∥v j -z∥ 2 + C(∥v j -z (k) ∥ 2 ),and the first inequality holds with equality at z = z (k) . Sum up the items on both sides with weights {a j } j∈[N ] , we obtainEQUATIONand the equality holds at z = z (k) :EQUATIONAfter obtaining the convex upper bound L(z), it becomes feasible to employ convex optimization algorithms to optimize this objective.Newton-IRLS. We first derive the formulations of gradient and Hessain matirx of L as follows:∇ L(z (k) ) = N j=1 a j • w (k) j • 2(z (k) -v i ) ∇ 2 L(z (k) ) = N j=1 a j • w (k) j • 2 • IThen, the gradient descent (GD) (η is the stepsize) isz (k+1) = z (k) -η • ∇ L(z (k) ) = z (k) -η • N j=1 a j • w (k) j • 2(z (k) -v j ),and the Newton Iteration isEQUATIONIn convex optimization, it has been well-established that second-order methods converge much faster than first-order approaches, but they require substantial computation in calculating or approximating the inverse Hessian matrix. However, due to the uniqueness of our L in Eq. ( 3), we can derive a concise closed-form iteration using the second-order Newton method as in Eq. (10) . Compared to the first-order gradient descent (GD) iteration, our Newton-IRLS algorithm enjoys several advantages as follows:• Fast convergence: Newton method converges at a quadratic rate, which is significantly faster than the linear convergence of gradient descent (GD). The comparative analysis of them can be found in Figure 4 (a) in ablation studies; • Interpretable formulation: The resulted form in Eq. (??) employs a normalized reweighted sum, which can be interpreted as robust estimator by down-weighting the outliers, as discussed in the following paragraph;• Efficient computation: The Hessian ∇ 2 L(z (k) ) can be easily computed as a closed-form diagonal matrix, facilitating the matrix inversion and multiplication in the Newton's iteration.B.3 Proof of Rigorous Loss Descent Guarantee (Theorem 3.2)Proof. Since z (k+1) is obtained from optimize the convex localized upper bound L at z (k) , then we have L(z (k+1) ) ≤ L(z (k) ). According to the upper bound in Eq. ( 6) and localized equality in Eq. ( 7), it is not hard to get the following inequality:L(z (k+1) ) ≤ L(z (k+1) ) ≤ L(z (k) ) = L(z (k) ).Therefore, optimizing the localized upper bound L can guarantee the rigorous descent of L.Our Newton-IRLS is a general framework which can be derived as different reweighting schemes with different penalties:• Square Loss (ℓ 2 ):ρ(z) = 1 2 z 2 , w (k) j = ρ ′ (∥v j -z (k) ∥) 2∥v j -z (k) ∥ = 1 2 , z * = N j=1 a j • v j .ℓ 2 loss increase quadratically with z, which suggests that ℓ 2 loss is more sensitive to the residual magnitude. Particularly, ℓ 2 loss can recover the vanilla attention since the weights are constant 1 2 . • Absolute Loss (ℓ 1 ):ρ(z) = z, w (k) j = ρ ′ (∥v j -z (k) ∥) 2∥v j -z (k) ∥ = 1 2∥v j -z (k) ∥ .With ℓ 1 loss, the weight is inversely proportional to ∥v j -z (k) ∥. By up-weighting the inliers and down-weighting the outliers, ℓ 1 -based estimators can mitigate the effect of large magnitude residues. • Minimax Concave Penalty (MCP) [50] :ρ γ (z) = z -z 2 2γ if y < γ γ 2 if y ≥ γ , w (k) j = ρ ′ (∥v j -z (k) ∥) 2∥v j -z (k) ∥ = 1 2 max 1 ∥v j -z (k) ∥ - 1 γ , 0 .MCP loss becomes constant when z is large and the weight derived by MCP loss enhances the interpretability of the robust estimator by down-weighting or completely removing the outliers. To be specific, the weight w j becomes smaller as the distance ∥v j -z (k) ∥ increases, thereby down-weighting the outlying cases. When this distance exceeds the threshold γ, the weight becomes 0, totally removing the outliers.• Huber loss:ρ δ (z) = 1 2 z 2 if z < δ δ • (z -1 2 δ) if z ≥ δ , w (k) j = ρ ′ (∥v j -z (k) ∥) 2∥v j -z (k) ∥ = 1 2 min 1, δ ∥v j -z (k) ∥ .Huber loss is equivalent to the ℓ 2 loss within the range (0, δ), and it becomes similar to ℓ 1 when z > δ, which indicates that Huber loss may mitigate the effect of large noise while keeping decent performance in noiseless scenario. • Huber-MCP:ρ δ,γ (z) =      1 2 z 2 if z < δ δ • (z -1 2 δ -(z-δ) 2 2(γ-δ) ) if δ ≤ z < γ δγ 2 if γ ≤ z , w (k) j = ρ ′ (∥v j -z (k) ∥) 2∥v j -z (k) ∥ = 1 2 max min δ γ -δ γ ∥v j -z (k) ∥ -1 , 1 , 0 .This penalty combines the advantages of Huber and MCP in recovering the ℓ 2 loss and largely mitigating the outliers.C.1 Language Domain The version also has the patch which fixes some of the corrupted test set images already applied. • Cora-ML [71] : The Cora dataset is a widely-used benchmark dataset in the field of graph-based tasks. It consists of 2708 scientific publications classified into one of seven classes. The citation network consists of 5429 links. Each publication in the dataset is described by a 0/1-valued word vector indicating the absence/presence of the corresponding word from the dictionary. The dictionary consists of 1433 unique words. Working with the Cora dataset presents challenges typical of real-world graph data, such as handling sparse and high-dimensional feature vectors, and dealing with the complex structure of the graph. • Citeseer [72] : The CiteSeer dataset is another popular dataset in the graph field. It consists of 3312 scientific publications classified into one of six classes. The citation network consists of 4732 links. Each publication in the dataset is described by a 0/1-valued word vector indicating the absence/presence of the corresponding word from the dictionary. The dictionary consists of 3703 unique words.Language Domain:• PGD-Adv [31] : The Projected Gradient Descent (PGD) method stands as the most prevalent attack strategy in the field of computer vision. It is primarily utilized for crafting adversarial examples in the context of adversarial training. The defense in this paper is adapted directly from PGD-adv in computer vision, extending its application to language modeling. • MixADA [17] : The search space for adversarial examples in language models is typically vast due to their discrete nature. To enhance the robustness of these models, MixADA integrates adversarial training [30] with mixup data augmentation [73] , thereby expanding the range of adversarial examples covered. Specifically, mixup generates synthetic training examples by linearly blending pairs of inputs and their corresponding labels. This approach enables the model to learn from a broader and more effective set of adversarial examples during training. • FreeLB [47] : Different from attacks that directly change the words in the sentence, FreeLB adds adversarial perturbations to word embeddings and minimizes the resultant adversarial loss around input samples. To expedite the process of adversarial training, FreeLB implements a single descent step on the parameters concurrently with each of the K ascent steps applied to the perturbation, which utilizes the average of accumulated gradients over the K steps. This efficiency has established FreeLB as a popular defense method in the field of NLP. • TA-VAT [18] : TA-VAT is another virtual adversarial training method that generates gradientbased perturbations on the embedding space. To create fine-grained perturbations, TA-VAT employs a token-level accumulated perturbation vocabulary. This vocabulary serves to better initialize the perturbations. Additionally, TA-VAT utilizes a token-level normalization ball, which effectively constrains these perturbations in a relevant and precise manner. • Adversarial Training (AT): Adversarial training is adaptive to the attack to be evaluated. Take the Textfooler as the instance, at every epoch, we generate 1000 perturbations from the Textfooler and add them into the training dataset to reinforce the training of models. We utilize the TextAttack [63] platform the conduct this adversarial training. • SmoothLLM [61] : Motivated by finding that the adversarial-prompting jailbreak is sensitive to the random character-level changes, SmoothLLM is designed by firstly perturbing multiple copies of the given prompt and then aggregating all the outputs.Beyond Language Domain:• Graph Convolutional Network (GCN) [74] : GCN is motivated by the localized first-order approximation of spectral graph convolutions. The basic idea is to first add self-loops to the adjacency matrix and then normalize the matrix. • Graph Attention Network (GAT) [5] : GAT leverages the attention mechanism to construct masked self-attentional layers. This allows the nodes to reweight their neighbors via the feature similarity.• GNNGuard [75] : GNNGuard is a universal reweighting framework that can be applied to any GNN. It leverages the cosine similarities between nodes' features to up-weight the correlated nodes and prune the edges between the dissimilar pairs. • Robust GCN (RGCN) [76] : RGCN first models the latent representations as the Gaussian distributions. Then the weights of different neighborhoods will be assigned different weights according to their variances when performing the message propagation. • Graph Random Neural Network (GRAND) [77] : The core of GRAND is the random propagation, wherein the node feature will be partially or entirely dropped out and then propagated through over the graph. This operation enable the node to be insensitive to the specific neighborhood, which prevents the effect of malicious outliers. Additionally, the random propagation also help to augment the representation for each node, thus improving the generalization of GNN. • Property GNN (ProGNN) [78] : The core principle of ProGNN is to robustify the GNNs through enhancing the graph properties of sparsity, low rank and feature smoothness. It provides a graph structure learning framework to learn the clean graph structure and parameters simultaneously.• Jaccard-GCN [79] : The basic idea of Jaccard-GCN is to preprocess the adjacency matrix by first computing the Jaccard coefficients of paired node features and then dropping the edges where the coefficients are below the threshold.• SoftMedian [80] : SoftMedian is a robust estimator for the message passing aggregation. It reweights the adjacency weights based to the distances of the hidden embeddings between the neighbor nodes and the dimension-wise median of the the entire neighboring representations.Classical language models:• BERT [56] : BERT stands out as one of the most well-known transformer-based language models. It is pretrained through masked language modeling (MLM), where it learns to predict words that have been masked, using context for guidance. This pretrained model is then fine-tuned for a variety of downstream tasks, showcasing its versatility and effectiveness in diverse applications.In our experiments, we will use BERT-110M.• RoBERTa [57] : RoBERTa is developed to overcome certain limitations of the original BERT model. This is accomplished by implementing key modifications such as increasing the batch size, extending the training epochs, and employing advanced optimization techniques. As a result of these strategic changes, RoBERTa has demonstrated substantial performance improvements over BERT across various NLP benchmarks. In our experiments, we will use RoBERTa-125M.• ALBERT [58] : ALBERT is a lite variant of BERT. It is achieved by decoupling the word embedding from the hidden embedding, significantly cutting down the number of parameters. To further enhance its efficiency, ALBERT employs cross-layer parameter sharing, ensuring that all layers use the same parameters. The reductions not only minimize memory footprint but also improve the efficiency of the model. In our experiments, we will use ALBERT-12M.• DistilBERT [59] : DistilBERT is a light version of BERT, maintaining most of the performance of the original BERT. It is trained with the knowledge distillation technique [81] to achieve high efficiency. In our experiments, we will use DistilBERT-66M.Large Language Models:• T5 [60]: Text-to-Text Transfer Transformer (T5) is a transformer-based neural network model known for its versatility and power in handling a wide range of NLP tasks. T5 simplifies NLP tasks by treating them uniformly as text-generation challenges. The T5 model family offers a range of sizes, from 60 million to 11 billion parameters, catering to different computational needs. The flexibility has made T5 a popular choice in NLP research. In our experiments, we will use T5-770M.• LLaMA [12] : LLaMa, the Large Language Model developed by Meta AI, represents a cuttingedge advancement in language modeling. Trained on publicly available datasets, LLaMa is available in various sizes to suit different computational needs. Notably, LLaMa-13B demonstrates superior performance over GPT-3 in most benchmarks, highlighting its exceptional effectiveness and capability in NLP tasks. In our experiments, we will use LLaMA-7B.• Vicuna [13] : Vicuna is a high-performing, open-source chatbot that impresses with capabilities comparable to GPT-4. Fine-tuned from the LLaMa model, it utilizes user-shared conversations gathered from Share-GPT for its training. Remarkably, Vicuna achieves 90% of the performance level of GPT-4, despite having only 13 billion parameters, showcasing its efficiency and effectiveness. In our experiments, we will use Vicuna-7B.Vision Models:• ViT [3]: The Vision Transformer (ViT) is a model in computer vision that adopts the principles of the Transformer architecture. In ViT, an image is processed similarly to a sequence of words, or tokens. Specifically, the image is segmented into fixed-size patches, each of which is then linearly transformed into an embedded representation. When trained on sufficient data, ViT achieves state-of-the-art performance on image classification benchmarks, competing with or outperforming leading CNN-based models. In our experiments, we will use ViT-86M.• Swin [4] : Swin Transformer is a popular variant of ViT, standing out for its enhanced efficiency and superior performance. It employs a hierarchical architecture, which not only aligns more closely with the nature of visual data but also boosts efficiency. To effectively capture global contextual information, Swin Transformer incorporates shifted window-based self-attention, further enhancing its effectiveness in vision-related applications. In our experiments, we will use Swin-50M. • BEIT [82] : Due to the success of BERT, BEIT harnesses the concept of masked language modeling to enhance self-supervised learning in the visual domain. To align with the words in language models, BEIT first maps the patch in an image into a token with an autoencoder. In the training process, it masks a portion of these patches, using the remaining unmasked ones to predict the masked tokens. Subsequently, the model is fine-tuned for a variety of downstream tasks, demonstrating its adaptability and effectiveness in diverse applications. In our experiments, we will use BEIT-86M. E Attacks.For the classic text attacks, we follow the default attack setting in the TextAttack [63] and the detailed information are as follows:• DeepWordBug [7] : DeepWordBug is black-box attacks that apply character-level transformations to the highest-ranked tokens misclassify the text input. It includes several character transformations including swap, substitution, deletion and insertion. We hold the maximum difference on edit distance (Levenshtein Edit Distance) to 30 for each sample. We will greedily modify the works with the word importance ranking. • PWWS [9] : The probability weighted word saliency (PWWS) employs a new word order determined by the word saliency and predicted probability, and then greedily perform the synonyms substitution.• TextFooler [14] : TextFooler propose a more comprehensive paradigm to generate adversarial perturbations. It firstly identify the important words and then replace them with the most semantically and syntacticaly similar synonyms until the prediction is altered. We set the minimum word embedding cosine similarity as 0.5 and the universal sentence similarity threshold as 0.840845057. • TextBugger [8] : TextBugger is a framework for creating utility-preserving adversarial texts against the classification models. It can be employed under both white-box and black-box settings. Specifically, this attack sorts the words based on Jacobian matrix in white-box setting and based on specific scoring function in black-box setting. It includes swap, substitution, deletion and insertion. There are two types of substitutions: Sub-C (replace characters with visually similar characters) and Sub-W (replace a word with its top-K (K = 5) nearest neighbors.) The universal sentence threshold is set as 0.8.• Prompt attack: we conduct adversarial prompt attacks using the text attacks including Deep-WordBug [7] , PWWS [9] , TextFooler [14] and TextBugger [8] . The evaluation setting follow PromptBench [15] which is a unified and comprehensive python library for evaluating LLMs. • Jailbreak: We employ the Greedy Coordinate Gradient (GCG) [55] method to search for the adversarial prompt suffixes that circumvent the alignment of the language models. It can be conducted under the transfer attack and direct attack. Specifically, the transfer attacks firstly attack the surrogate model (e.g., Vicuna) and can apply the perturbations to the victim models. The direct attacks directly apply the GCG method to the victim models to generate the suffixes, which can be much stronger. The JailbreakKeywords is a list of words commonly found in messages that abstain from responding to a prompt asking for inappropriate content. In our experiment, we use the following list: JailbreakKeywords = ["I'm sorry", "Sorry", "I apologize", "As an", "I'm just", "I'm an", "I cannot", "I would advise", "it's not appropriate", "As a responsible AI"].• FGSM [30] : Fast Gradient Sign Method (FGSM) is a widely used technique for generating adversarial examples. It crafts perturbations specifically to maximize errors in the output of neural networks.• PGD [31] : Projected Gradient Descent (PGD) method is one of the most prominent attack strategies in computer vision. Unlike FGSM employing a single step, PGD uses multiple steps to generate adversarial examples. This iterative approach includes a projection operation, which ensures that the intensity of the attack remains within specified limits, making PGD a more controlled and effective method for generating adversarial examples. The steps are K = 7 and the steps size α = 0.00784. • PGD on Graph [64] : Motivated by PGD [31] in vision domain, [64] propose a first-order method to conduct topology attack on discrete graph structure. This method firstly solve continuous optimization problem by Projected Gradient Descent (PGD) method and then utilize the random sampling to get the optimal binary topology perturbation from the continuous probabilistic matrix. Trajectory. To further validate the convergence and effectiveness of our algorithm, we use a toy experiment to visualize the trajectories of updated vector in 2D plane in Figure 8 . We use L 1 penalty in our algorithm, the simulated attention matrix and value matrix are as follows:EQUATIONEQUATION)Figure 8 : Optimization trajectory.From the figure, we can find that with the mean as the initial position, the updated vector can approach closely to the ground truth within only 3 steps. This phenomenon further validate the effectiveness and efficiency of our algorithm.Robust estimation.We firstly generated clean samples {x i } n i=1 (blue dots) and the outlier samples {x i } n+m i=n+1 (red dots) from 2-dimensional Gaussian distributions, N ((0, 0), 1) and N ((8, 8), 0.5), respectively. We calculate the mean of clean samples 1 n n i=1 x i as the ground truth of the mean estimator. Then we estimate the mean of all the samples by solving arg min z n+m i=1 ρ(z -x i ) using the our method, where ρ(•) can take different penalties such as ℓ 2 penalty ∥ • ∥ 2 2 and ℓ 1 penalty ∥ • ∥ 2 . In Figure 9 , we visualize the generated clean samples and outliers, as well as the ground truth means and the mean estimators with η(•) = ∥ • ∥ 2 2 or ∥ • ∥ 2 under different outlier ratios 15% and 45%. The results show that, with the existence of outliers, the ℓ 2 -based estimator deviates far from the true mean, while the ℓ 1 -based estimator is more resistant to outliers and MCP-based estimator is the most robust.To gain insight into the vulnerability of vanilla attention (WLS estimator), we attempt to analyze the effect of perturbed tokens on the output embedding of attention modules. Specifically, since vanilla attention can be formalized as a WLS estimator: z * = arg min z L(z) = N j=1 a j • ∥v j -z∥ 2 , we quantify and compare ∥v ′ j -z * ∥ 2 and ∥v j -z * ∥ 2 , where v j and v ′ j are the original and perturbed tokens, and z * is the ground truth output embedding. We present the numerical results across every attention module in the pretrained BERT on IMDB dataset in Table 5 . The results demonstrate that attackers tend to introduce larger residuals ∥v j -z∥ by modifying the important tokens v j when perturbing the text. We present the results of sentiment analysis on IMDB dataset under various attacks in Table 6 . We can conlude from the results that our methods improve the robustness of the backbones significantly by simply plugging the ProAttention into the models without additional fintuning or training. Moreover, our method can be combined with the existing defenses such as Adversarial Training (AT) to further improve the performance. In Table 7 , we display the results of textual entailment on RTE across different cosine similarities constraints in TextFooler attack. We select DistilBERT as the backbone model and construct several MCP-based architectures with different γ. We can observe that our method can improve the robustness acorss different cosine similarities. The performance improvement is more evident under the smaller cosine similarities, which is equivalent to larger budgets.In Table 8 , we present the results of textual entailment on RTE across various attacks. The results exhibit the consistent improvement of our methods over the backbone model. Adversarial training techniques are highly effective to enhance the robustness of models via adding the adversarial examples into the training set. To better capture the robustness enhancement of adversarial training, we track the adversarial fine-tuning curves and present the detailed results on AG's News in Table 9 and Figure 10 . In the beginning of every epoch, we generate 1000 perturbed examples using the specific attack and then put them to the original training dataset. From the results, we can make the following observations: (1) the models show even better robustness during the process of adversarial training. (2) our method can be combined with adversarial training to further improve the resilience of the models.Table 9 : Adversarial fine-tuning on AGNEWS.We present the ablation study on the maximum perturbation percentage, minimum cosines similarity and sentence similarity threshold in Figure 11 , Table 10 , Table 11 and Table 12 , respectively. The experiments are performed on AGNEWS under TextFooler with the ALBERT as the backbone. The default values are as follows: sentence similarity threshold is 0.840845057, maximum perturbation percentage is 1.0, synonym cosine similarity is 0.5. The results show the consistent improvement of our method over the backbone models. Our proposed ProAttention is a universal framework which can be applied to various attention-based models. To verify the universality of our methods, we integrate our robust attention module into various backbones and present the results in Figure 12 and Table 13 . As seen in the results, our ProAttention can consistently enhance the robustness over any backbone under various attacks. For the Huber-based model, we present the ablation study on the δ and layers K of Huber loss under TextFooler in Table 14 and Figure 13 . The default setting are as follows: δ = 0.9 and K = 3, and we vary the two parameters separately to capture the trend. We can find that the performance is insensitive to δ or layers within an appropriate range.Table 14 : Ablation study on Huber on AGNEWS. We present the ablation study of the γ and layers K in MCP in Table 15 & Figure 14 and Table 16 & Figure 15 . From the figures, it can be concluded that appropriate γ is needed to get the best robustness. Besides, more layers can get a more precise solution for the robust objective, but we need to consider a good balance between precision and efficiency. We compare the backbone model T5 with its robust version based on ℓ 1 , MCP and Huber loss in Figure 16 and Table 17 . The experiments are conducted on SST2 under TextFooler. As shown in the figure, Pro-T5 (Huber or ℓ 1 ) can slightly improve the robustness of backbone T5. Moreover, Pro-T5 (MCP) significantly outperform other baselines, especially under the large attack budgets. We present the ablation study on δ in Huber on SST2 in Figure 17 and Table 18 . We fix the number of layers as 3 and vary the values of δ from 3.0 to 9.0. The Huber-based model with δ = 4.0 perform best among all the selected parameters.Figure 17 : Ablation study on Huber on T5 We present the ablation studies in MCP on SST2 in Figure 18 and Table 19 . The default settings are as follows: K = 3 and γ = 3.0. We can make the observations as follows: (1) For γ, the optimal value falls into the range of (3, 5) . ( 2) For number of layers, the best value can be chosen as 4 or 5. For LLaMA, we can observe an intriguing phenomenon that differs from the T5 case: the ℓ 1 and MCP-based methods sacrifice too much accuracy while Huber method can keep decent performance under small budgets. This reason is that in the small range region, ℓ 1 and MCP utilize a linear or concave functions while Huber can recover the ℓ 2 function. Inspired by the characteristics of these functions, we combine the properties of Huber and MCP, and construct a new function which we refer to Huber-MCP 2 . The detailed formulation and derived robust attention layers are available in Appendix B. As indicated by the following curves, Huber-MCP and Huber-based models exhibits better robustness than other methods while preserving the good clean performance.We present the results of textual entailment on SST2 under TextFooler in Figure 20 . We can observe that ℓ 1 and MCP-based methods sacrifice the performance because of the estimation bias. Pro-LLaMA (Huber-MCP) shows slight improvement over other models. 2 Empirical penalty selection strategy: For small or medium-sized models such as BERT (110M) and ViT (86M), MCP-based models exhibit superior robustness while nearly not sacrificing the clean performance. Moreover, MCP-based models are easy to tune with only one parameter γ. For large models like LLaMA (7B) and Vicuna (7B), it is necessary to choose Huber and Huber-MCP to recover the original ℓ2 penalty within the low-value region in case of clean performance drop.We present the results of textual entailment on SST2 under TextBugger in Figure 21 . In this case, ℓ 1 -based model show a catastrophic performance drop while Pro-LLaMA (Huber) outperforms other baselines with a sinificant margin. We provide the results of transfer jailbreak in Table 21 , Table 22 , and Table 23 . As SmoothLLM exhibits excellent performance to defend the jaibreaking attack, we also combine the random smoothing with the backbone models and our methods to further validate the effectiveness of our method.As shown in the results, when q = 0 (without random smoothing), by simply plugging our ProAttention into the Vicuna, the robustness can be improved by a significant margin. Our plug-and-play method can even be comparable to the randomly smoothed models which require multiple operations including random perturbations, votings and aggregations. SMOOTH q 0 1 3 5 VICUNA 91.8 71.8 20.9 0.9 PRO-VICUNA-HUBER δ = 0.1 1.8 0.9 0.9 0.9 PRO-VICUNA-HUBER δ = 0.2 8.2 1.8 0.9 0.9 PRO-VICUNA-HUBER δ = 0.3 21.8 2.7 0.9 0.9 PRO-VICUNA-HUBER δ = 0.4 30.9 21.8 0.9 0.9 PRO-VICUNA-HUBER δ = 0.5 36. 4 23.6 0.9 0.9 PRO-VICUNA-HUBER δ = 0. 0.9 0.9 0.9 0.9 0.9 0.9 0.9 0.9 0.9 0.9Although our Pro-Vicuna has demonstrated significant effectiveness under transfer attack (black-box), it is still unclear whether our method can be resilient under white-box attacks which adaptively target the specific victim models. The comparison of our Pro-Vicuna and backbone Vicuna under adaptive jailbreak is presented in Table 24 and Figure 23 . Our Pro-Vicuna can improve Vicuna by an average of 10.4% across various numbers of attack queries. We don't include SmoothLLM in adaptive attacks since it introduces non-differentiable operators that preclude the gradient-based GCG attack on it. The main results of FGSM on ViT are presented in Table 25 . We can conclude that our Pro-ViT (MCP) outperforms other methods across various budgets. J.2 Ablation study.The ablation study of PGD and FGSM are provided in Table 27 , Table 28 and Figure 24 . As shown in the results, the optimal γ of MCP fall into the range of (3, 4) . The robust estimators with more layers show the better robustness while slightly sacrifice the clean performance. The main results on Citeseer and the ablation study on Cora-ML are presented in Table 29 and Table 30 . From the results, we can make the following conclusions: (1) Our Pro-GAT outperform other methods significantly. Under the larger budgets, the outliers introduced by the adversarial attack Justification: We provide the theory in Section 3.3 and the detailed proof in the Appendix B. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.Justification: There is no societal impact of the work performed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 


Decision: This paper introduces a robust attention mechanism intended to enhance the robustness of transformer-based architectures against adversarial attacks. A key advantage is that it can serve as a plug-and-play layer, therefore can be easily integrated into existing transformers for enhancing robustness without additional training or fine-tuning. 

Overall, all reviewers appreciate the simplicity and practicability of this method, and acknowledge that the provided experiments are comprehensive and the corresponding results are solid. Meanwhile, some concerns are also raised, mainly including: 1) some critical discussions are missing (e.g., efficiency analysis, limitation); 2) the comparisons against linear attention are missing; and 3) the presentation could be further improved.

The rebuttal is considered, which addresses most of these concerns. Consequently, three reviewers stand on the positive side of accepting this paper. The sole reviewer who raised objections, primarily concerning the comparison with linear attention, did not respond to the details provided in the rebuttal; additionally, by reading the original review and the provided rebuttal, the AC feels the rebuttal reasonably alleviates this concern. Therefore, the AC decides to ignore this negative review and agree with the other three reviewers on accepting this submission.

Reviews: review: paper_summary: paper_summary: In this paper, the authors intend to robustify transformer architectures against adversarial attacks to enhance their resilience across various machine learning tasks. Specifically, they propose the ProAttention mechanism. They use a novel interpretation of the self-attention mechanism as a weighted least squares (WLS) estimator, along with robust WLS token estimators and an efficient Newton-IRLS algorithm with convergence guarantees, which enhances the robustness of transformers without additional training or fine-tuning. Further, the authors develop a plug-and-play layer, ProAttention, to be integrated into existing transformers. They demonstrate several experiments to show the performance of this architecture in improving robustness across language modeling, image classification, and graph representation learning.
strength_weakness: strength_weakness: 1. The paper does not include error margins for experiments, especially for ablation studies in Section 4. The paper provides extensive experimental results; however, it does not include error bars or confidence intervals to indicate the variability or statistical significance of the results.

2. This paper could benefit from better presentation and clearer structure. For example, the captions for the figures, especially Fig. 1 and 2, are too simple and lack the necessary explanations. This may lead to misunderstandings about the designs or results.

3. I do not see a discussion of the method's limitations, despite the title of the final section mentioning “Limitation.”

4. Although the paper proposes enhancing transformer robustness through the ProAttention mechanism, it needs a thorough explanation and analysis of why this method performs better under various attacks. I suggest that the authors further discuss this aspect in more detail
questions: questions: 1. The paper claims that the Newton-IRLS algorithm converges efficiently within three steps. However, can the authors provide a more detailed analysis or proof of this convergence rate, especially for larger models or more complex datasets?

2. The authors show the effectiveness of ProAttention in enhancing robustness across language, vision, and graph domains. However, each domain's specific implementation details and hyper-parameter tuning strategies seem different. Can the authors clarify how the ProAttention framework can be uniformly applied across these diverse domains without extensive domain-specific modifications?

3. The ProAttention framework relies on the weighted least squares (WLS) estimator and its robust variants. Can authors provide a deeper theoretical justification for choosing WLS and its robustness under adversarial attacks? Specifically, how do the assumptions of WLS hold up under different types of adversarial perturbations, and are there scenarios where the WLS assumptions might fail, leading to reduced robustness?
limitations: limitations: Despite being designed for efficiency, the ProAttention framework introduces additional computational overhead during inference. The paper mentions that the ProAttention layer requires 1-2 times more inference time, which may limit its deployment in real-time applications or on devices with limited computational resources.

The ProAttention framework relies on WLS estimators and their robust variants. While the paper provides a theoretical justification for choosing WLS, its effectiveness under different adversarial perturbations needs further investigation. In some scenarios, the assumptions of WLS may fail and reduce robustness.
paper_summary: paper_summary: This paper proposes an interpretable robust attention layer to robustify transformer architecture via a plug-and-play paradigm.
strength_weakness: strength_weakness: It seems that the complexity of the proposed ProAttention is still greater than that of linear attention [1] [2] [3].

[1] Han, Dongchen, et al. "Flatten transformer: Vision transformer using focused linear attention." Proceedings of the IEEE/CVF international conference on computer vision. 2023.
[2] Katharopoulos, Angelos, et al. "Transformers are rnns: Fast autoregressive transformers with linear attention." International conference on machine learning. PMLR, 2020.
[3] Zhu, Lianghui, et al. "DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention." arXiv preprint arXiv:2405.18428 (2024).
questions: questions: 1. Could the authors compare the proposed ProAttention with linear attention under the same experimental setting?
limitations: limitations: The comparison with different attention mechanisms is missing.
paper_summary: paper_summary: This paper proposed a robust transformer architecture named ProTransformer through a plug-and-play paradigm without further training or fine-tuning. The authors include robust token estimators in the self-attention blocks which are more resilient to the dominating impact of input tokens, and apply Newton-ISLR to approximate these estimators. The comprehensive experiments for various tasks, modalities, attacks, and backbones show that ProTransformer can significantly improve the robustness of transformers without compromising benign performance.
strength_weakness: strength_weakness: 1. Lack of analysis of clean performance in jailbreak attacks. I want to know if ProTransformer would hurt the generation quality of LLMs.
2. Less discussion about the semantic level attacks and robustness of corruption. For example, the AdvGlue Benchmark (https://adversarialglue.github.io/) for corruption robustness and PAIR[1] for semantic-level jailbreak attacks.
3. Lack of comparison with baseline defense methods.

[1] Chao, Patrick, Alexander Robey, Edgar Dobriban, Hamed Hassani, George J. Pappas, and Eric Wong. Jailbreaking black box large language models in twenty queries. arXiv preprint arXiv:2310.08419 (2023).
questions: questions: 1. When using white-box adaptive attacks like PGD and GCG, do you compute the gradient through your ProAttention blocks?
2. For the jailbreak attack results in Figure 6, why do you use Huber instead of MCP attack constraints? As shown in Section 4.2.1, MCP shows the best performance.
limitations: limitations: The limitations of the paper are not adequately discussed. The cost of the evaluation with ProTransformer should be included in the limitations.
paper_summary: paper_summary: This paper proposes a robust attention mechanism to improve the resilience of transformer-based architectures. The method does not need additional training or finetuning with only four lines of code, which is simple yet effective. To show the effectiveness of the proposed mechanism, the paper conducts experiments across different tasks including topic classification, sentiment analysis, textual entailment, jailbreak attack, on different backbone architectures such as large language models. The results show promising resilience.
strength_weakness: strength_weakness: 1. As the paper points about one of the advantages of the method is the efficiency (efficient Newton-ISLR algorithm), there should be more discussions about the runtime of the proposed method across different architectures. The current version only includes the study in Tab.2 on a single model type. What is the influence on LLMs as they are more computation intensive? And how is the runtime compared with other robust designs?

2. For the results on image classification tasks, there is no comparision. How is the performance of the proposed method compared with other robust designs such as [*]? 

[*] Robustifying Token Attention for Vision Transformers
questions: questions: 1. For the main results in Tab.1, how many layers (K) in the transformer models are equipped with ProAttention?
2. For the results in Tab.2, what is the robutness performance by employing different numbers of ProAttention layers?
limitations: limitations: Please refer to weakness and questions.

Comments: comment: Thank you for your feedback. We have carefully incorporated your suggestions into our revision, but the updated manuscript cannot be uploaded at this moment. We would greatly appreciate it if you could consider our clarifications in the rebuttal.
Thanks for your response. I would maintain my score.
Thanks for the response from the authors. The response have answered my questions and I will maintain my rating of acceptance  for the paper.
Thank you for your response. I look forward to seeing the updated manuscript.