Heads of Security
Learned attention gates improve prompt injection robustness
Paper
Code
Models
TutorialA learned gate on every attention head allows the model to read tool outputs without following instructions injected in them.
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 3x 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.
<|start|> userCountR’s<|start|> toolSendmoney<|start|> assistanttokens
tokens
scores
gate
scores
weights
(For illustration) A single head with hand-picked scores, shown to illustrate the mechanism. Each token is assigned a channel (trusted, untrusted, template). Before the softmax, each head scales its attention scores by a learned, channel-dependent gate, so that an instruction-following head can be prevented from reading untrusted content while other heads remain free to process that content, preserving the model's ability to reason over it. HoS trains only gate parameters: 18,432 scalars on GPT-OSS-20B, 46,080 on Llama-3.3-70B, trained with the SecAlign++ DPO recipe. The base model weights never change.
HoS achieves strong security and utility
Adaptive attacks
HoS holds up against attacks trained or optimized to break it: the black-box RL-Hammer attack on InjecAgent, with attackers trained against each target model, and the white-box Claudini attack on AgentDojo and InjecAgent, optimized specifically against HoS.
RL-Hammer on InjecAgent · black-box
Following the original RL-Hammer setup, we use GRPO to train a Llama-3.1-8B-Instruct attacker to rewrite an attacker goal into an injection optimized against a particular target model, and evaluate on held-out InjecAgent goals. Rows identify the attacker’s training target and columns its evaluation target. Shaded cells show attacks against the model they were trained to defeat; the remaining cells measure how well an attack transfers.
Eval target (ASR %) | ||||
|---|---|---|---|---|
| GPT-OSS-20B | w/ SecAlign | w/ HoS | GPT-5.6 Luna | |
| InjecAgent (static) | 0% | 0% | 0% | 0% |
| RL train target | ||||
| GPT-OSS-20B | 95% | 27% | 5% | 1% |
| w/ SecAlign | 97% | 85% | 10% | 6% |
| w/ HoS | 89% | 46% | 3% | 6% |
| GPT-5.6 Luna | 87% | 31% | 2% | 36% |
| Highest ASR | 97% | 85% | 10% | 36% |
Eval target (ASR %) | |||
|---|---|---|---|
| Llama-3.3-70B | w/ SecAlign | w/ HoS | |
| InjecAgent (static) | 16% | 5% | 0% |
| RL train target | |||
| Llama-3.3-70B | 99% | 4% | 0% |
| w/ SecAlign | 99% | 82% | 1% |
| w/ HoS | 72% | 6% | 0% |
| GPT-5.6 Luna | 76% | 0% | 0% |
| Highest ASR | 99% | 82% | 1% |
Among the attackers shown, the highest ASR against GPT-OSS-20B is 97% without a defense and 85% with SecAlign, compared with only 10% with HoS; for Llama-3.3-70B the corresponding rates are 99%, 82% and 1%. HoS also withstands these attacks better than GPT-5.6-Luna: the Luna-trained attacker reaches 36% ASR on Luna, against a highest ASR of 10% on the HoS model.
Claudini · white-box
Starting from the best-performing attack reported by Claudini, we use GPT-5.6-Sol to develop attack variants over 50 autoresearch iterations against HoS-defended GPT-OSS-20B, and evaluate the selected attack against GPT-OSS-20B and its SecAlign and HoS variants on 100 AgentDojo user-task–injection-task pairs and the 100 held-out InjecAgent goals used for RL-Hammer. For each example and target model, the attack optimizes a 120-token suffix under a budget of 1017 FLOPs; an example counts as successful if any suffix found within the budget succeeds.
Cumulative ASR against attack compute. HoS requires roughly three times as much attack compute to reach 50% ASR as either the base model or SecAlign, despite the attack being developed specifically against HoS.
Benign utility
On standard multiple-choice capability benchmarks, HoS stays close to the base model. SecAlign and ASIDE perform comparably zero-shot, but degrade with few-shot examples and across conversation turns.
Multiple-choice benchmarks
ARC-C | GSM8K | MMLU | MMLU-Pro | |||||
|---|---|---|---|---|---|---|---|---|
| 0-shot | 0-shot | 1-shot | 8-shot | 0-shot | 5-shot | 0-shot | 5-shot | |
| GPT-OSS-20B | 96.1 | 91.4 | 93.6 | 95.6 | 83.0 | 82.3 | 72.0 | 68.3 |
| w/ SecAlign | 94.8 | 93.3 | 70.9 | 91.8 | 82.8 | 54.4 | 71.6 | 59.4 |
| w/ ASIDE | 94.7 | 91.3 | 93.2 | 93.6 | 83.2 | 74.1 | 71.7 | 52.1 |
| w/ HoS (ours) | 95.4 | 90.5 | 92.3 | 94.6 | 83.2 | 80.0 | 72.3 | 67.2 |
| Llama-3.3-70B | 94.5 | 93.0 | 93.7 | 95.7 | 84.5 | 82.8 | 67.8 | 56.3 |
| w/ SecAlign | 93.8 | 92.2 | 2.0 | 68.0 | 84.5 | 28.4 | 68.6 | 16.1 |
| w/ HoS (ours) | 93.7 | 92.9 | 93.2 | 95.7 | 85.0 | 82.4 | 67.9 | 54.8 |
Accuracy (%). Bold marks large drops from the undefended model. Every model is within 2 points of the undefended model in zero-shot. We use the standard few-shot in-context learning setup for these benchmarks, but SecAlign often answers a few-shot demonstration rather than the final question.
Multi-turn AlpacaFarm
100 three-turn AlpacaFarm conversations, each with untrusted input in one turn (shaded). Undefended Llama and HoS remain near 100% across turns and placements of untrusted input. SecAlign drops most sharply on the input turn and remains degraded on later fully trusted turns. Instruction-completion rate per turn:
SecAlign shows a similar pattern to its few-shot failures: it continues answering an earlier instruction, treats a later instruction as data for the earlier one, or combines instructions across turns. For example, after a question about Thai food, it ignores an input field specifying chocolate chip cookies and writes a Pad Thai recipe. When next asked why it is a bad idea to give a mouse a cookie, it replies, “Now, let’s get back to cooking” and continues the recipe.
Agentic prompt injection benchmarks
Compared to SecAlign and ASIDE, HoS matches top performance on agentic prompt injection benchmarks (AgentDojo, AgentDyn, WASP and LLMail).
Security-utility trade-off for GPT-OSS-20B. Each point is a separately trained model at one learning rate. Both defenses Pareto-dominate the undefended model. ASIDE loses the ability to call tools and stays near-zero utility on all four benchmarks. WASP values are means over five runs due to its small sample size.
Llama-3.3-70B · utility and ASR on agentic benchmarks
| Benchmark | Metric | Llama-3.3-70B | w/ SecAlign | w/ HoS (ours) |
|---|---|---|---|---|
| AgentDojo | Utility ↑ | 53.7% | 67.0%(+13.3%) | 58.5%(+4.8%) |
| Utility attacked ↑ | 36.2% | 63.9%(+27.7%) | 52.4%(+16.2%) | |
| ASR ↓ | 23.1% | 2.1%(-21.0%) | 0.9%(-22.2%) | |
| AgentDyn | Utility ↑ | 33.7% | 41.7%(+8.0%) | 18.3%(-15.4%) |
| Utility attacked ↑ | 30.9% | 38.3%(+7.4%) | 18.4%(-12.5%) | |
| ASR ↓ | 18.5% | 5.4%(-13.1%) | 1.4%(-17.1%) | |
| WASP | Utility ↑ | 66.5% | 62.7%(-3.8%) | 62.2%(-4.3%) |
| Utility attacked ↑ | 62.4% | 79.8%(+17.4%) | 81.4%(+19.0%) | |
| ASR ↓ | 7.6% | 2.4%(-5.2%) | 1.4%(-6.2%) | |
| LLMail | Utility attacked ↑ | 80.2% | 84.6%(+4.4%) | 82.4%(+2.2%) |
| ASR ↓ | 2.1% | 12.7%(+10.6%) | 0.0%(-2.1%) |
Changes relative to the undefended model in parentheses. HoS brings ASR to 1.4% or less everywhere and matches the undefended model's utility except on AgentDyn. SecAlign raises LLMail ASR: it calls tools in 29.95% of generations even though the user only asks for an email summary, these calls often use fabricated recipients and email bodies and sometimes adopt the attacker-specified values from untrusted emails, resulting in the high ASR.
How to use HoS
Live demo on Hugging Face
Pick an AgentDojo scenario, try injecting different instructions in the tool output, and compare the responses from GPT-OSS-20B with and without HoS gates.
Load the gated model
To use the gates in your own code, install the package and load them on top of the base model. Mark which tokens are untrusted, then generate as usual.
pip install heads-of-security
from heads_of_security import load_gated_model
model, tokenizer = load_gated_model("imperial-cpg/hos-gpt-oss-20b")
enc = tokenizer.apply_chat_template_with_channels(messages, tools=tools)
out = model.generate(
input_ids=enc["input_ids"],
attention_mask=enc["attention_mask"],
token_channels=enc["token_channels"],
)What a gates file contains, and how to load a gated model and generate with it.
Analyzing the learned attention gates
On the Analysis page, we look inside the model: how the learned gates redistribute attention, an interactive before-and-after view of the first generated token’s attention on injections, and what HoS does to role confusion.
BibTeX
@article{shilov2026heads,
title = {Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness},
author = {Shilov, Igor and Yang, Xiaoxue and Panfilov, Alexander and Kortukov, Evgenii and
Meeus, Matthieu and Romov, Peter and Samek, Wojciech and Geiping, Jonas and
de Montjoye, Yves-Alexandre},
journal = {arXiv preprint arXiv:2610.37657},
year = {2026},
url = {https://arxiv.org/abs/2610.37657}
}