Learned attention gates make prompt injection costlier
Reviewing Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness · Igor Shilov, Xiaoxue Yang, Alexander Panfilov et al.
Researchers add small learned gates to each attention head so a model weighs trusted and untrusted text differently, sharply cutting prompt injection success while keeping the base model frozen. The defence holds up better than most against attackers who know it is there, but it only enforces a trust boundary someone else has to draw.
A black-box reinforcement learning attacker, trying to make a model obey instructions hidden in content it reads, succeeds 97% of the time against an undefended open model. Against the same model fitted with a small set of learned attention gates, it succeeds 10% of the time. The interesting part is where that defence lives.
Prompt injection works because a model reads everything as one stream. The user's request and a sentence in a retrieved web page are both just tokens. Attention lets any token influence any other, so the model has no built-in sense of whose words it should obey.
The researchers' system, Heads of Security, labels every token as trusted, untrusted or template (the chat formatting around them). Each attention head then gets a learned gate for each label. The gate scales how much attention that head pays to each kind of token. The base model stays frozen. Only the gates are trained, on paired examples where the right answer follows the user rather than the planted instruction. That is tens of thousands of numbers, not millions.
What I like is what the gates learned. Most turn down attention from trusted text to untrusted text. A minority turn it up, and appear to help the model notice that an injection is present. Suppressing untrusted content evenly across every head did worse. So the model is not blinded to the data. It reads it with a different weighting in each head.
It was tested against an attacker who knew
This is what I look for in a defence paper and rarely find. The researchers ran adaptive attacks, including a white-box optimisation attack strengthened specifically against the gates. It still got through. It needed about three times the compute to match its success against the base model or a fine-tuning defence. That is a higher cost, not immunity, and the researchers say so plainly.
It also costs little in usefulness. Few-shot benchmark scores barely move, while a fine-tuning defence and a rotation-based architectural defence both degrade there. The rotation-based one struggles even to call tools.
The gates only enforce the line you draw
Every token must be labelled before the model sees it. The gates enforce that boundary; they do not decide it. If a pipeline marks a scraped page as trusted, the defence has nothing to work with. The same applies when a workflow legitimately needs the model to follow instructions found in external content. An attacker facing this will go after the labelling, not the gates.
There are narrower limits too. Training used single-turn examples from a fixed corpus. The method assumes standard pairwise attention, so it does not carry straight over to linear-attention architectures.
This does not end prompt injection. It moves the fight from the model's attention to whoever assigns trust. That is a better place to have it.
Further analysis
Problem
The paper studies indirect prompt injection, where attacker-planted instructions in external content are mistaken for authorised commands. As models gain access to external data and tools, such injections enable attackers to hijack agentic behaviour. The authors focus on enforcing an instruction-data trust boundary so models may read and reason over untrusted content without following instructions embedded within it.
Approach
Heads of Security (HoS) is an architectural defence that inserts a small set of learned per-head attention gates into existing decoder transformer layers. Tokens are assigned to three channels: trusted, untrusted and template. Each head learns per-channel gate logits that bias pre-softmax attention scores, scaling attention to keys from each channel via a sigmoid. Gates are trained while base model weights remain frozen, using a preference-based optimisation recipe and paired examples where the preferred response follows the trusted user instruction rather than an injected instruction in the untrusted span. HoS thus redistributes attention between trusted and untrusted tokens on a head-by-head basis. The method is evaluated on Llama-3.3-70B-Instruct and GPT-OSS-20B and compared to two baselines: SecAlign++, a finetuning/adapter defence, and ASIDE, an architectural rotation-based defence.
Key Findings
- Large reductions in attack success: a state-of-the-art black-box RL attack reaches 97% attack success rate (ASR) on undefended GPT-OSS-20B and 85% on SecAlign but only 10% on HoS; for Llama-3.3-70B the strongest attacker yields 99%/82%/1% ASR for undefended/SecAlign/HoS respectively.
- Increased white-box effort: a white-box optimisation attack specialised against HoS requires roughly three times more FLOPs to reach comparable ASR than against the base model or SecAlign, showing HoS raises the computational cost of successful adaptive attacks.
- Preserved utility: HoS largely preserves base-model performance on multiple-choice and prompting benchmarks, including few-shot MMLU-Pro where HoS scores ~67% versus 68% for the base model while SecAlign and ASIDE fall substantially; HoS also maintains multi-turn instruction-following robustness where SecAlign degrades.
- Agentic trade-offs: on agentic benchmarks that measure tool use and injection resistance, HoS and SecAlign both reduce ASR and improve the security-utility frontier, while ASIDE frequently fails to call tools and yields very low utility.
- Mechanistic insight: most gates learn to reduce attention to untrusted keys for trusted queries, while a minority increase attention to untrusted content and appear to support injection detection; uniform suppression across heads performs worse than learned per-head gates for the security-utility trade-off.
- Parameter efficiency: HoS trains only tens of thousands of scalar gate parameters (for example ~46k for Llama-3.3-70B, ~18k for GPT-OSS-20B) versus millions for SecAlign adapters or full fine-tuning.
Limitations
HoS enforces a supplied trust boundary but does not decide which sources are trusted; it is unsuitable when external content legitimately contains authorised instructions. The evaluation used a fixed training corpus with single-turn preference pairs, so broader training data could change results. The method as implemented targets dense pairwise attention and is not directly compatible with linear-attention architectures. Finally, attackers can adapt given knowledge of the defence and unlimited resources, so HoS raises attack cost rather than proving absolute immunity.
Implications
Offensively, the results indicate that an attacker aiming to succeed against HoS must invest substantially more computational resources and adopt stronger adaptive strategies, for example advanced white-box optimisation and iterative autoresearch to craft suffixes that bypass per-head gates. Attackers might instead target trust-assignment mechanisms or workflows that legitimately require following external instructions in order to bypass the instruction-data separation. The paper demonstrates that attention-level interventions materially change the attacker effort and strategy required to mount successful prompt-injection attacks.
The paper title, authors and abstract
Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness
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.
