Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness
Prompt injection succeeds because models follow instructions hidden in untrusted content. Learned attention gates cut the RL-Hammer attack success rate from 97% to 10%.
The claim
Learned attention gates cut a black-box injection attack success rate from 97% to 10%.
The problem
Useful models must read untrusted content that arrives from the web.
Injected instructions hide inside that content and get obeyed by the model.
Existing defenses trade away utility or break when tools are used.
How it works
HoS learns a small set of gates on attention heads.
Each head redistributes attention across trusted and untrusted tokens.
The model can reason over data without obeying injected instructions.
The defense is evaluated on Llama-3.3-70B-Instruct and GPT-OSS-20B models.
Instruction-data separation is enforced inside the model architecture.
Figures
What changes
The RL-Hammer black-box attack success drops from 97% to 10%.
A white-box attacker needs three times more FLOPs to match it.
Base-model utility stays intact on standard benchmarks.
Limits
Two open instruction-tuned backbones carry the whole evaluation.
Adaptive attacks were rebuilt specifically against this defense.
On agentic benchmarks with static attacks, results stay close to the best defense.
Source · arXiv