ProTransformer: Robustify Transformers via Plug-and-Play Paradigm

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.

Strengths

1. The method is simple yet effective. With only four lines of codes, it's able to provide good resilience performance across diverse tasks and model architectures. 2. The experiments are quite comprehensive. The study of robust transfomers include language modeling, image classification, and graph representative learning. For language modeling, the evaluations consider classic text attacks, prompt attacks, and jailbreak attacks. 3. The results demonstrate promising performance. 4. The plug-and-play nature of ProTransformer makes it practical for real-world applications, as it can be easily integrated into existing model architectures without significant overhead.

Weaknesses

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

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?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Please refer to weakness and questions.

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.

Strengths

1. This paper is well-written with clear descriptions of their method with both motivation/intuition analysis and theoretical analysis. 2. Comprehensive experiments are conducted to prove their claims of the enhanced robustness of Transformers. This includes robustness analysis for both the traditional adversarial attacks within the word or character level and the recent jailbreak attacks for large language models. Detailed ablation studies and transformers for alternative modalities are also considered. 3. The ProTransformer is simple to use (Plug-and-Play) and has significantly improved the robustness. Moreover, ProTransformers can not only work on language modeling but also show its generalization for transformer-based architectures in vision and graph domains.

Weaknesses

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

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.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

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

Summary

This paper proposes an interpretable robust attention layer to robustify transformer architecture via a plug-and-play paradigm.

Strengths

The proposed method is practical and can be plugged into the given transformer as a plug-and-play layer. The experiments are robust and conclusive, signifying the efficacy and satisfactory performance of the proposed method.

Weaknesses

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

1. Could the authors compare the proposed ProAttention with linear attention under the same experimental setting?

Rating

4

Confidence

1

Soundness

2

Presentation

2

Contribution

2

Limitations

The comparison with different attention mechanisms is missing.

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.

Strengths

1. The idea of ProAttention is novel. Most existing work focuses on either architecture-agnostic defenses or robustness improvements through adversarial training, but these approaches often fail to generalize across different tasks or require substantial computational resources. In this paper, the authors propose a plug-and-play paradigm to enhance robustness against adversarial attacks without additional training or fine-tuning. 2. The ProAttention framework is reasonable, and the adaptation of the Newton-IRLS algorithm makes it efficient and effective. The framework leverages the robust weighted least squares estimator to mitigate the effects of adversarial attacks. 3. The experiments are extensive, and sufficient hyper-parameters are provided for reproduction. The authors detail their experimental setup and provide comprehensive ablation studies to demonstrate the robustness and efficiency of their approach under various attack scenarios.

Weaknesses

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

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?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.