EN
Learned attention gates reduce prompt injection in two open models
Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness · arXiv:2610.37657
In plain words
Models sometimes need to read tool outputs or other outside content, but that content can contain instructions the model should not follow.
Heads of Security (HoS) is a defense that learns which parts of the model should pay less attention to those instructions while still letting other parts process the content.
In one test, an adaptive attacker succeeded on 10% of held-out goals against GPT-OSS-20B with HoS, compared with 97.2% against the undefended model. This result applies to that attacker and test setup.
The research
When a language model reads tool output, it may need the information but should not obey instructions planted there by an attacker. It remains unclear whether a defense can enforce that distinction while leaving the model able to process the content.
HoS learns per-head attention gates to limit obedience to injected instructions while allowing models to process external content.
The challenge is to let a model use information in external content without treating planted instructions as commands. HoS addresses this by learning attention gates that let some model heads limit attention to untrusted content while others can still process it.
Against a SecAlign-trained adaptive attacker, the attack succeeded on 10.0% of 100 held-out InjecAgent goals against GPT-OSS-20B with HoS, compared with 85.2% with SecAlign and 97.2% without a defense.
Success means the attacker's tool was named in the model's first tool call; the rates average five runs at medium reasoning effort.
These results cover two open-weight models and the tested attacks and tasks. HoS depends on a correctly assigned trust boundary, and finite evaluations cannot rule out a successful attack by a capable adversary.
What the evidence shows
- Separate data from instructions
Heads of Security (HoS) is a prompt-injection defense that helps language models treat external content as data rather than authorized instructions. This matters when a model reads tool output: it may need the information there while ignoring instructions planted to redirect its actions.
The paper asks whether a model can separate those roles without losing its ability to use the content.
- Learn gates for each head
HoS assigns tokens to trusted, untrusted, and chat-template groups, then learns attention gates for each attention head. Some heads can reduce attention to untrusted tokens, while other heads can still process them.
The trust boundary is supplied to the system; HoS does not determine which sources or instructions should be trusted.
- Train on competing responses
The authors train the gates on preference pairs: one response follows the user's instruction, while another follows an injected instruction. They compare HoS with SecAlign, which uses the same preference dataset but has different trainable components and places the untrusted text differently. This comparison frames the results without implying that every training condition was identical.
- Measure resistance to attack
In an adaptive attack test, an attacker trained against SecAlign succeeded on 10% of 100 held-out InjecAgent goals against HoS on GPT-OSS-20B. The same attacker succeeded on 85.2% against SecAlign and 97.2% against the undefended model.
Success rates were averaged over five runs at medium reasoning effort, so this is evidence about that attacker and evaluation setup.
- Check utility and fixed suppression
The authors checked ordinary question answering as well: on five-example MMLU-Pro questions, GPT-OSS-20B with HoS scored 67.2% accuracy, close to the base model's 68.3%; SecAlign scored 59.4%. A separate comparison on AgentDojo and AgentDyn found that learned gates could achieve a different balance between legitimate task completion and injected-task success than fixed uniform suppression.
In that GPT-OSS-20B comparison, AgentDojo excluded Workspace user tasks 35 and 38 because they can exceed the base model's context limit; AgentDyn excluded no tasks. Together, these results support a bounded case for HoS, not a guarantee across tasks or attacks.
Deep
- Original paper authors
- Igor Shilov, Xiaoxue Yang, Alexander Panfilov, Evgenii Kortukov, Matthieu Meeus
- Submitted to arXiv
- License
- CC BY 4.0