Research map / Architectures and efficiency / Attention and transformers
ProTransformer: Robustify Transformers via Plug-and-Play Paradigm
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ProAttention, a plug-and-play robust attention layer that reframes attention as a weighted least-squares estimator and replaces it with robust penalties (MCP, Huber, l1) solved by a Newton-IRLS iteration, improving transformer robustness without training.
Problem
Transformers, including LLMs and vision/graph transformers, are highly vulnerable to adversarial input perturbations (text attacks, prompt attacks, jailbreaks), and existing defenses need costly training or are architecture-specific.
Approach
Attention output is the optimal solution of a weighted least-squares token estimator; replacing the quadratic penalty with robust penalties (MCP, Huber, l1) and optimizing via a Newton-IRLS scheme yields a reweighted attention that down-weights or removes outlier tokens, usable as a drop-in layer at inference time.
Key results
- Pro-BERT (MCP) raises AGNEWS accuracy under TextFooler from 19.7% to 39.2% without fine-tuning; ProAttention improves ALBERT, DistilBERT, RoBERTa, BERT under TextFooler by 28.3%, 16.1%, 11.4%, 19.5% respectively.
- Pro-Vicuna (Huber, delta=0.1) reduces GCG jailbreak ASR to 1.8%, comparable to SmoothLLM.
- Pro-ViT attains 77.39% under PGD at 1/255 budget on CIFAR-10 vs 41.75% for the second-best model; Pro-GAT reaches 66.5% at 40% graph perturbation on Cora-ML vs 56.7% for GAT.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| plug in ProAttention with MCP penalty | BERT · AGNEWS, TextFooler | AUA%: 19.7 -> 39.2 | Helped |
| add adversarial training to Pro-BERT (MCP) | Pro-BERT (MCP) · AGNEWS, TextFooler | AUA%: 39.2 -> 56.8 | Helped |
| add ProAttention (MCP) to BERT + AT | BERT + AT · AGNEWS, TextFooler | AUA%: 42.1 -> 56.8 | Helped |
| plug in ProAttention with MCP penalty | ALBERT · AGNEWS, TextFooler | AUA%: 20.6 -> 48.9 | Helped |
| use MCP rather than the vanilla attention penalty (ablation) | BERT · AGNEWS, TextFooler | AUA% | Helped |
| use MCP rather than the vanilla attention penalty (ablation) | BERT · AGNEWS, TextBugger | AUA% | Helped |
| use MCP rather than the vanilla attention penalty (ablation) | BERT · AGNEWS, DeepWordBug | AUA% | Helped |
| use MCP rather than the vanilla attention penalty (ablation) | BERT · AGNEWS, PWWS | AUA% | Helped |
| use Huber or ℓ1 penalty instead of MCP (ablation) | BERT · AGNEWS, adversarial text attacks | AUA% | Mixed |
| plug in ProAttention | ViT · CIFAR-10, PGD 1/255 | accuracy: 34.61 -> 77.39 | Helped |
| plug in ProAttention | ViT · CIFAR-10, PGD 4/255 | accuracy: 1.83 -> 48.11 | Helped |
| plug in ProAttention | ViT · CIFAR-10, PGD 8/255 | accuracy: 0.26 -> 33.40 | Helped |
Surprising findings
- Inference-time-only robustification without any fine-tuning yields large gains, rivaling some adversarial training methods.
- Newton-IRLS converges within about 3 iterations, adding only 1-2x inference time over the backbone.
Limitations
- Efficiency still has room for improvement; additional inference cost grows with the number of iterations.
- Results are only for plug-and-play setting; training or fine-tuning with ProAttention on large models is left for future work.
- Hyperparameters (delta, gamma) are tuned per setting.
Experimental setup
Datasets: AG's News IMDB SST-2 RTE Behaviors CIFAR-10 ImageNet-1K Cora-ML Citeseer
Benchmarks: AG's News IMDB SST-2 RTE Behaviors CIFAR-10 ImageNet-1K Cora-ML Citeseer
Metrics: Clean accuracy Accuracy under attack (AUA%) Attack success rate (ASR%)
Baselines: FreeLB PGD-Adv MixADA TA-VAT Adversarial training SmoothLLM GCN GAT GNNGuard RGCN GRAND ProGNN
Builds on: Newton-IRLS SmoothLLM GCG TextAttack PromptBench
Closest papers by meaning
- Adventures of Trustworthy Vision-Language Models: A Survey (AAAI 2024 · similarity 0.86)
- Transformers Learn Low Sensitivity Functions: Investigations and Implications (ICLR 2025 · similarity 0.85)
- Adapters Strike Back (CVPR 2024 · similarity 0.85)
- Viewing Transformers Through the Lens of Long Convolutions Layers (ICML 2024 · similarity 0.85)
- Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models (ICML 2024 · similarity 0.84)
- SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.