ProTransformer: Robustify Transformers via Plug-and-Play Paradigm

Thanks for your recognition of the efficiency and effectiveness of our method. We are glad to solve134 your concerns and answer your questions with the following illustrations.

W1: 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 Table 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?

Answer. I include the average running time across different architectures. The results show the consistent conclusion with the BERT backbone: Our ProTransformer will only include 1-2 times of inference time while achieving great robustness. We want to point out that robustifying models is usually very time-consuming and labor-intensive. For example, adversarial training methods usually require substantial training time (e.g., 7-100 times of the normal training time [1]), while our ProTransformer can directly plug in the robust attention into a pre-trained backbone model without any training. Moreover, randomized smoothing methods usually require a lot of inference time (e.g., 1,000 times for the language domain [2] and 100,000 times for the vision domain [3]). Additionally, most robust architectures are quite complicated, require training the models from scratch, and exhibit limited robustness. Therefore, our ProTransformer, which only requires 1-2 times more inference time, is quite efficient and economical.

[1] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. "Towards deep learning models resistant to adversarial attacks." In International Conference on Learning Representations, 2018

[2] Lou, Qian, et al. "CR-UTP: Certified Robustness against Universal Text Perturbations." arXiv preprint arXiv:2406.01873 (2024).

[3] Cohen, Jeremy, Elan Rosenfeld, and Zico Kolter. "Certified adversarial robustness via randomized smoothing." International conference on machine learning. PMLR, 2019.

[4] Han, Xing, et al. "Designing robust transformers using robust kernel density estimation." Advances in Neural Information Processing Systems 36 (2024).

Table 5: Average running time (ms) on AG-NEWS.

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

[*] Robustifying Token Attention for Vision Transformers

Answer. We validate the advantage of our ProTransformer over RVT [1] and TAP&ADL [2] in the following table.

Table 6: Adversarial robustness under PGD.

[1] Towards Robust Vision Transformer

[2] Robustifying Token Attention for Vision Transformers

Q1: For the main results in Tab.1, how many layers (K) in the transformer models are equipped with ProAttention?

Answer. We equipped ProAttention with 3 layers (K).

Q2: For the results in Tab.2, what is the robustness performance by employing different numbers of ProAttention layers?

Answer: We already include the ablation study on the number of layers in the Appendix G.6.2. The results validate that our algorithm can converge well within 3 steps.

Thanks for your recognition of the novelty and effectiveness of our method. We are glad to solve33 your concerns and answer your questions with the following illustrations.

W1: Lack of analysis of clean performance in jailbreak attacks. I want to know if ProTransformer would hurt the generation quality of LLMs.

Answer. In ProTransformer, using Huber loss will almost not hurt the generation quality while using L1 or MCP will slightly sacrifice the generation performance. We construct 100 synthetic texts (including prompts and reference texts) evaluate the generation with the BLEU score and a series of ROUGE scores with the vicuna-7b as the backbone. The following results verify that Huber model keep similar score with L2 (vanilla), while L1 and MCP may sacrifice a little bit quality.

Table 2: BLEU and ROUGE score on generation quality.

W2: 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.

[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).

Answer. For AdvGlue Benchmark, we already included the evaluation under typos-based Textbugger and embedding similarity-based Textfooler. We add the experiment of the context-aware BERT-Attack in the following table, the results show our Pro-BERT outperforms all of the other baselines.

Table 3: ASR (%) under BERT-Attack.

Additionally, we also evaluate the advantage of our ProTransformer under semantic-level jailbreak attack PAIR in the following table.

Table 4: ASR (%) under PAIR Jailbreaks.

W3: Lack of comparison with baseline defense methods.

Answer. Actually we have included the comparison with baseline defense methods for each domains:

• Language domain: backbone baselines ALBERT, DistilBERT, RoBERTa and BERT, and popular defenses including MixADA, PGD-Adv, FreeLB, TA-VAT and SmoothLLM.
• Vision domain: We include Deit, Convit, BeiT, Swin, ViT.
• Graph domain: GCN, GAT, GNNGuard, RGCN, GRAND, ProGNN, Jaccard-GCN and Soft-Median.

Q1: When using white-box adaptive attacks like PGD and GCG, do you compute the gradient through your ProAttention blocks? Answer. Yes. We compute the gradient through the ProAttention blocks.

Q2: 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.

Answer. In Section 4.2.1, the task is classification, so both Huber and MCP losses would not sacrifice too much clean performance. But in jailbreak, we evaluate the model with the generated response. In this scenario, Huber loss has $L_2$ penalty in the low-value region, which enables the model to keep better generation quality. But MCP will hurt the generation quality, result in some non-semantic answer.

L1: The limitations of the paper are not adequately discussed. The cost of the evaluation with ProTransformer should be included in the limitations.

Answer. Please refer to the global response for the concern of limitations of ProTransformer.

Thanks for your recognition of the efficiency and effectiveness of our method. We are glad to solve134 your concerns and answer your questions with the following illustrations.

W1: 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, Liangxiu, et al. "DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention." arXiv preprint arXiv:2405.18428 (2024).

Answer: Thanks for your comments. We would like to emphasize that while we aim to develop algorithm with acceptable efficiency, efficiency is not the focus of this work. Therefore, linear attention is irrelevant to this context. Instead, our ProAttention targets on improve the robustness of transformers. Linear attention aims to improve the efficiency of transformers but is still vulnerable against adversarial attacks as verified in next answer. Their contributions are totally orthogonal.

Q1: Could the authors compare the proposed ProAttention with linear attention under the same experimental setting?

L1: The comparison with different attention mechanisms is missing.

Answer: Thanks for your comments. We compare different attention mechanisms (vanilla attention, Linear attention [1], Longformer attention [2], and our ProAttention) on IMDB across various attacks in the following table. Our method shows better robustness compared to other attention mechanisms including the linear one.

[1] Katharopoulos, Angelos, et al. "Transformers are rnns: Fast autoregressive transformers with linear attention." International conference on machine learning. PMLR, 2020.

[2] Beltagy, Iz, Matthew E. Peters, and Arman Cohan. "Longformer: The long-document transformer." arXiv preprint arXiv:2004.05150 (2020).

Table 1: The results of sentiment analysis on IMDB.

Thanks for your recognition of the efficiency and effectiveness of our method. We are glad to solve134 your concerns and answer your questions with the following illustrations.

W1: We want to clarify that the proposed ProAttention is plugged into fixed pre-trained models without finetuning, hence there is no randomness of optimization of training. Also, the attack strategies strictly follow the procedures without randomness. Besides, our evaluation method is common in robustness of language models, and consistent with the existing benchmarks [1] and text attacks works.

[1] Searching for an effective defender: Benchmarking defense against adversarial word substitution, 2021.

W2: Thanks for pointing it out. We have referred to the figures and included the explanations in the main text. We will include necessary captions for the figures in the revised version.

• For Figure 1, we present three attack mechanisms to manipulate and harm the language models:

  1. Classic text attacks modify the input content.
  2. Prompt attacks perturb the prompt template.
  3. Jailbreaks add adversarial non-semantic suffixes.

• For Figure 2, we show the schematic overview of the ProTransformer: We design our ProTransformer by plugging ProAttention into transformers, replacing the vanilla attention mechanism without additional training. Furthermore, our ProTransformer is versatile and can be applied across various domains, including language, image, and graph processing.

W3: Please refer to the global response for the concern of limitations of ProTransformer.

W4: We have provided the detailed motivations and explanations for the robustness in Section 3. We are glad to summarize them again here:

1. In Section 3.1, we highlight the vulnerability of vanilla attention due to the quadratic impact of L2 penalty in the WLS (L2) estimator. In particular, the attacker can easily manipulate the input tokens to dominate the output because of the quadratical penalty on the difference between input and output tokens.
2. In Section 3.2 and 3.3, we demonstrate that robust penalties, such as L1 or MCP, can mitigate the quadratic impact to a linear or constant impact. Specifically, when some tokens are attacked and introduce large residues, our ProAttention can adaptively downweight their attention weights to mitigate their impact.

To simplify the understanding, it is helpful to think about a special case of it when all the attention weight are equal: median estimation (L1 case) is notably more robust than mean estimation (quadratic case) because median estimation can be formulated as the minimizer of the sum of absolute deviations (L1 norm) which reduces the impact of outliers.

Q1: We would like to clarify that in general, there is no way to prove that an optimization algorithm converges in three steps, especially for the challenging non-convex and non-smooth problems proposed in this paper. In addition, there is no clear connection between the convergence rate of the neural network layers and the size of models or datasets. What we provide in our paper is to theoretically guarantee each iteration step will decrease the function values so running more steps will at least not hurt. We then provide empirical evaluations to show it indeed converges fast in a few steps, which makes the algorithm practical. The exact convergence rate is neither the focus nor the primary contribution of this work, but this can be an exciting future research.

Q2: It is common in adversarial defenses to control the balance between clean and robust performance through key hyperparameters. For instance, the famous paper TRADES [1] needs to carefully tune the combination between clean loss and robust loss to achieve best performance even for different attack budgets for the same model architectures and datasets, not to mention totally different domains evaluated in this paper. Since we plug-in our robust layers into the given models without training, it is expected some reasonable hyperparameter tuning is necessary to obtain optimal performance.

We will provide some general guidance in the revision: (1) Huber loss better maintains clean performance so it typically works better when the attack is weak; (2) MCP loss works better when the attack is strong; (3) L1 loss is a special case of Huber loss; (4) the proposed Huber-MCP penalty can flexibly reduce to different options such as Huber, MCP, and L1 based on the values of hyperparameters δ and γ as shown in Figure 3 in the paper.

As a future work, we can make those hyperparameters as learnable parameters learned through fine-tuning or adversarial training, which can avoid hyperparameter search but also need more training costs.

[1] Zhang, Hongyang, et al. "Theoretically principled trade-off between robustness and accuracy." (ICML, 2019)

Q3: We want to clarify that WLS is not an assumption; the vanilla attention is exactly equivalent to the WLS estimator (L2) under our new interpretation. Our robust ProAttention is an improved variant.

L1: Please refer to the global response for concern of efficiency of ProTransformer

L2: Please refer to Q3.