🛡️ ProTransformer

Robustify Transformers via Plug-and-Play Paradigm

Zhichao Hou¹ · Weizhi Gao¹ · Yuchen Shen² · Feiyi Wang³ · Xiaorui Liu^1,✉

¹North Carolina State University    ²Carnegie Mellon University    ³Oak Ridge National Laboratory

No retraining. No fine-tuning. Just swap the attention.
A 4-line robust attention layer that plugs into any pretrained transformer — language, vision, or graph.

📰 News

• [2024] 🎉 ProTransformer is accepted to NeurIPS 2024!
• [2024.10] 📄 Paper is released on arXiv and the code is open-sourced.

✨ Highlights

📈 Robustness gains at a glance

Without any further fine-tuning, ProTransformer improves the robustness of vanilla models by:

…and it also shows strong robustness in the vision and graph domains.

🧠 Method

1. Attention is a weighted least-squares (WLS) estimator

Each output token aggregates the value vectors with attention weights, which is exactly the solution of a weighted least-squares problem:

The quadratic penalty lets a few adversarially perturbed tokens dominate the output — the root of the vulnerability.

2. Robust token estimator

We replace the quadratic loss with a robust penalty $\rho$ (ℓ₁, Huber, MCP, Huber-MCP):

3. Newton-IRLS → ProAttention

Optimizing a convex localized upper bound with a Newton step yields a closed-form re-weighted attention:

with guaranteed descent $\mathcal{L}(z^{(k+1)}) \le \mathcal{L}(z^{(k)})$. Outlier tokens with large residuals are automatically down-weighted (and fully removed beyond $\gamma$ for MCP).

🔧 ProAttention in 4 lines (MCP)

D = torch.cdist(Z, V)                         # pairwise distance
W = torch.clip(1 / D - 1 / gamma, min=0)      # MCP weights
W = normalize(W * A, p=1, dim=-1)             # re-weight attention
Z = torch.matmul(W, V)                        # update

class ProAttention(nn.Module):
    def __init__(self, L=3, norm="L2", epsilon=1e-2, gamma=4.0, t=1.0, delta=4.0):
        super().__init__()
        self.L, self.norm, self.epsilon = L, norm, epsilon
        self.gamma, self.t, self.delta = gamma, t, delta

    def forward(self, A, V):
        M = torch.matmul(A, V)                      # vanilla attention output
        if self.norm == "L2":
            return M

        for _ in range(self.L):                     # Newton-IRLS iterations
            dist = torch.cdist(M, V)

            if self.norm == "L1":
                w = 1 / (dist + self.epsilon)
            elif self.norm == "MCP":
                w = 1 / (dist + self.epsilon) - 1 / self.gamma
                w[w < self.epsilon] = self.epsilon
            elif self.norm == "Huber":
                w = self.delta / (dist + self.epsilon)
                w[w > 1.0] = 1.0
            elif self.norm == "HuberMCP":
                w = self.delta / (self.gamma - self.delta) * (self.gamma / (dist + self.epsilon) - 1)
                w = w.clamp(min=self.epsilon, max=1.0)

            ww_norm = nn.functional.normalize(w * A, p=1, dim=-1)
            M = (1.0 - self.t) * M + self.t * torch.matmul(ww_norm, V)

        return M

🧩 Model Zoo

Every model below is a drop-in replacement for its 🤗 Transformers counterpart and loads the original pretrained weights. ProAttention is off by default (norm="L2" = vanilla attention) and switched on with one call.

🚀 Quick Start

0️⃣ Installation

As a library — just the ProTransformer models:

pip install protransformers

To run the quick-start scripts — clone the repo and install with the extra dependencies (TextAttack, torchvision, …):

git clone https://github.com/chris-hzc/ProTransformer.git
cd ProTransformer

conda create -n protransformer python=3.10 -y
conda activate protransformer
pip install -e ".[examples]"

1️⃣ Plug ProAttention into any model in 3 lines

from protransformers import AutoModelForSequenceClassification, set_pro_attention

model = AutoModelForSequenceClassification.from_pretrained("textattack/bert-base-uncased-ag-news")
set_pro_attention(model, norm="MCP", gamma=4.0, L=3)   # 🛡️ now it's a ProTransformer — no training needed
# set_pro_attention(model, norm="L2")                   # ↩️ back to vanilla attention

2️⃣ Pro-BERT family under classic text attacks

# Pro-BERT (MCP) under TextFooler on AG News
python text_attack.py --backbone bert --norm MCP --gamma 4.0 --L 3 --data ag-news --attack tf

# Vanilla BERT baseline
python text_attack.py --backbone bert --norm L2 --data ag-news --attack tf

# Other backbones: roberta | albert | distilbert
python text_attack.py --backbone albert --norm MCP --gamma 4.0 --attack tf

3️⃣ Pro-ViT under FGSM / PGD

# Vanilla ViT vs. Pro-ViT on CIFAR-10, PGD (7 steps, α = 2/255) with budgets 1/255, 4/255, 8/255
python image_attack.py --norm L2
python image_attack.py --norm MCP --gamma 4.0

# FGSM, custom budgets
python image_attack.py --norm MCP --attack fgsm --budgets 1 2 4 8

Uses nateraw/vit-base-patch16-224-cifar10 by default (--model to change); CIFAR-10 is downloaded to --data_dir.

4️⃣ Pro-GAT under adaptive PGD topology attack

A GAT is trained once per split on the clean graph; ProAttention is then plugged into its neighbor aggregation without retraining, and both models are attacked adaptively with the PGD topology attack of Xu et al. (2019) (global evasion, budget = fraction of edges flipped).

# GAT vs. Pro-GAT on Cora, 5 splits, budgets 5% / 10% / 20%
python graph_attack.py --data cora --norm MCP --gamma 4.0 --L 3 --budgets 0.05 0.1 0.2

# Citeseer, larger budgets
python graph_attack.py --data citeseer --norm MCP --gamma 4.0 --budgets 0.1 0.2 0.3 0.4

Graphs and 10/10/80 splits follow Mujkanovic et al. (2022); the .npz files are downloaded to --data_dir automatically.

5️⃣ Chat with Pro-Vicuna / Pro-LLaMA / Pro-T5

Prints the vanilla and the ProTransformer answer side by side:

# Pro-Vicuna
python llm_chat.py --model lmsys/vicuna-7b-v1.5 --norm Huber --delta 0.1 \
                   --prompt "Give me two tips for staying focused while studying."

# Pro-LLaMA (interactive session)
python llm_chat.py --model meta-llama/Llama-2-7b-chat-hf --norm MCP --gamma 4.0

# Pro-T5
python llm_chat.py --model google/flan-t5-large --norm MCP --gamma 4.0

6️⃣ Pro-Vicuna / Pro-LLaMA under jailbreak attacks

Evaluates the attack success rate (ASR ↓) on AdvBench harmful behaviors with refusal-prefix matching. Bring your own adversarial suffixes (e.g. GCG suffixes transferred from a surrogate model), one per line:

# Vanilla Vicuna
python llm_jailbreak.py --model lmsys/vicuna-7b-v1.5 --norm L2 --suffix_file suffixes.txt

# Pro-Vicuna (Huber, δ = 0.1)
python llm_jailbreak.py --model lmsys/vicuna-7b-v1.5 --norm Huber --delta 0.1 \
                        --suffix_file suffixes.txt --output pro_vicuna.csv

📊 Experimental Results

Pro-BERT (MCP) is competitive with adversarial-training defenses (FreeLB, PGD, MixADA, TA-VAT) without any training, and combining it with adversarial training (Pro-BERT (MCP) + AT) sets the best results across all four attacks. See the paper for results on LLMs (T5, LLaMA, Vicuna), jailbreaks, ViT, and GAT.

📁 Repository Structure

ProTransformer/
├── text_attack.py          # 2️⃣ Classic text attacks on Pro-BERT / RoBERTa / ALBERT / DistilBERT
├── image_attack.py         # 3️⃣ FGSM / PGD on Pro-ViT
├── graph_attack.py         # 4️⃣ Adaptive PGD topology attack on Pro-GAT
├── llm_chat.py             # 5️⃣ Chat with Pro-Vicuna / Pro-LLaMA / Pro-T5
├── llm_jailbreak.py        # 6️⃣ Jailbreak ASR of Pro-Vicuna / Pro-LLaMA
├── pyproject.toml          # pip install -e ".[examples]"
├── requirements.txt
├── protransformers/        # Modified 🤗 Transformers (v4.40) with ProAttention
│   ├── pro_attention.py    #   ← ProAttention + set_pro_attention
│   └── models/{bert,roberta,albert,distilbert,vit,llama,t5}/   # patched attention layers
└── figures/

📖 Citation

If you find ProTransformer useful in your research, please consider citing our paper and giving this repo a ⭐:

@article{hou2024protransformer,
  title={Protransformer: Robustify transformers via plug-and-play paradigm},
  author={Hou, Zhichao and Gao, Weizhi and Shen, Yuchen and Wang, Feiyi and Liu, Xiaorui},
  journal={Advances in Neural Information Processing Systems},
  volume={37},
  pages={137557--137609},
  year={2024}
}

🙏 Acknowledgements

This codebase builds upon 🤗 Transformers and TextAttack. We thank the authors for their great work.

📄 License

This project is released under the Apache 2.0 License. protransformers/ is a modified copy of 🤗 Transformers, which is also licensed under Apache 2.0.

📬 Contact

For questions, please open an issue or contact Zhichao Hou (zhou4@ncsu.edu).