Robust attention.
Plug and play.
ProTransformer swaps the attention layer of any pretrained Transformer for a robust estimator that adversarial tokens can't hijack. No retraining or fine-tuning. Works on text, vision, graphs and LLMs.
$ pip install protransformers
Attention is a weighted average, and averages break easily.
Each attention output is a weighted mean of value tokens. In statistics terms it's a weighted least-squares estimator, and least squares has a breakdown point of zero: one bad token can drag the output anywhere.
Adversarial attacks exploit exactly this. A few swapped words, an invisible pixel pattern, a gibberish jailbreak suffix or a handful of fake graph edges each inject a token that pulls the representation off course, layer after layer.
With ρ = ℓ₁, Huber or MCP, tokens far from the consensus get down-weighted automatically, so outliers stop steering the output.
See it in action
Each dot is a value token in 2D with equal attention weight. Click the canvas to inject adversarial tokens and drag any dot to move it. The grey ✕ is vanilla attention, and the coloured ◆ is ProAttention, with its IRLS path. Each dot's size shows the weight ProAttention gives it.
Robust across language, LLMs, vision and graphs
The same ProAttention module drops into BERT, LLaMA, ViT and GAT, using the original pretrained weights.
Schematic of the GCG suffix-injection attack (Zou et al., 2023). The suffix is a nonsense token string optimized to flip the model into compliance. ProAttention treats those tokens as outliers. The numbers below are attack success rates from the paper (Table 21, Vicuna-7B, transfer GCG).
ViT splits an image into 16×16 patches and attends over them. A PGD perturbation pushes a few patch embeddings far away from the rest, and vanilla attention averages them in. ProAttention sees those patches as outliers relative to the other patches and down-weights them in every layer.
Toy view: the red patches are perturbed. Hover over the grid to watch ProAttention down-weight them; vanilla attention keeps every weight uniform.
Drop-in replacements for 🤗 Transformers
Load the original checkpoints, then switch ProAttention on with one call. With norm="L2" you get the vanilla model back, bit for bit.
BertForSequenceClassificationRobertaForSequenceClassificationAlbertForSequenceClassificationDistilBertForSequenceClassificationLlamaForCausalLMT5ForConditionalGenerationViTForImageClassificationgraph_attack.GATThree lines to a robust model
- Install from PyPI with
pip install protransformers. - Load any supported checkpoint from
protransformersinstead oftransformers. - Switch on ProAttention with
set_pro_attention. That's all.
BERT inference on AG News goes from 6.1 ms to 11.7 ms with K = 2. That's about 1–2× the backbone's time, and nothing extra at training time.
from protransformers import (
AutoModelForSequenceClassification, set_pro_attention
)
model = AutoModelForSequenceClassification.from_pretrained(
"textattack/bert-base-uncased-ag-news"
)
# Plug in ProAttention: no retraining, same weights
set_pro_attention(model, norm="MCP", gamma=4.0, L=3)
# LLMs: load with attn_implementation="eager"
# Vanilla again: set_pro_attention(model, norm="L2")ProTransformer: Robustify Transformers via Plug-and-Play Paradigm
1North Carolina State University · 2Carnegie Mellon University · 3Oak Ridge National Laboratory
Advances in Neural Information Processing Systems (NeurIPS) 2024
@inproceedings{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},
booktitle = {Advances in Neural Information
Processing Systems},
volume = {37},
pages = {137557--137609},
year = {2024}
}