Learned attention gates improve defenses against prompt injection attacks
Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness
Machine LearningComputation and LanguageCryptography and Security
Summary
Prompt injection attacks trick AI models into following harmful instructions hidden in user inputs. The authors introduce Heads of Security (HoS), a way for AI models to separate trusted instructions from untrusted input by adjusting how parts of the model pay attention to each. This method makes AI models much harder to fool with tricky prompts, while still performing well on normal tasks. Compared to other defenses, HoS requires less effort to break and maintains useful abilities in many settings.
What this means in practice
- •For ai product developers: Build more secure AI services that resist prompt injection attacks using learned attention gating techniques like HoS.$Commercial implications: Provides a practical defense method to improve AI product security against advanced prompt injections for developers deploying large language models.
- •For cybersecurity engineers: Integrate learned attention interventions to enhance threat detection against instruction manipulation within AI systems.
Authors
Igor Shilov, Xiaoxue Yang, Alexander Panfilov, Evgenii Kortukov, Matthieu Meeus, Peter Romov, Wojciech Samek, Jonas Geiping, Yves-Alexandre de Montjoye
Abstract
We introduce Heads of Security (HoS), a prompt injection defense that enforces instruction-data separation via a small set of learned per-head attention gates. HoS learns how each head should redistribute attention between trusted and untrusted tokens, enabling models to reason over untrusted content without following injected instructions. We evaluate HoS on Llama-3.3-70B-Instruct and GPT-OSS-20B. We compare it against SecAlign++, a state-of-the-art finetuning-based defense, and ASIDE, a defense that, like ours, builds instruction-data separation into the model architecture. We show that HoS offers greater robustness to adaptive attacks, better performance on standard utility benchmarks, and matches top performance on standard agentic prompt injection benchmarks. More specifically, the black-box RL-Hammer attack reaches 97% ASR against GPT-OSS-20B and 85% against SecAlign, but only 10% against HoS. Claudini, a state-of-the-art white-box attack, strengthened against HoS through autoresearch, requires 3 times more FLOPs to reach the same ASR on HoS as it does on GPT-OSS-20B or SecAlign. HoS preserves base-model utility across standard multiple-choice benchmarks and prompting setups, whereas SecAlign and ASIDE degrade in few-shot settings. On agentic benchmarks with static attacks, HoS and SecAlign perform similarly and substantially improve the security-utility trade-off over the base model. ASIDE struggles with tool use, achieving only 10% utility. These results show that prompt injection robustness can be achieved and improved through targeted attention interventions.