Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness
Authors: Igor Shilov, Xiaoxue Yang, Alexander Panfilov, Evgenii Kortukov, Matthieu Meeus, Peter Romov, Wojciech Samek, Jonas Geiping, Yves-Alexandre de Montjoye
Organizations: Imperial College London · Tübingen AI Center · ELLIS Institute Tübingen · Max Planck Institute for Intelligent Systems · Fraunhofer HHI
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.
Figures & tables
Figure 1: (Top) Heads of Security (HoS) learns to redistribute attention between trusted and untrusted tokens. Green, red, and blue denote trusted, untrusted, and chat template tokens. Here, a head initially attends strongly to both the trusted request (“Count R’s”) and an injected instruction (“Send money”). Learned gates preserve trusted and chat template tokens ( 1× and 0.9× ) while suppressing untrusted tokens ( 0.1× ), shifting attention toward the trusted request. Gates are learned independently per head, so other heads may retain or increase attention to untrusted tokens. (Bottom) HoS improves security and utility. Compared to SecAlign and ASIDE, HoS is more robust against adaptive attacks (RL-Hammer and Claudini), better preserves utility across standard capability benchmarks, and achieves strong utility and security on static agentic benchmarks.
Eval target (ASR %)
GPT -OSS
w/ Sec Align
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%
Table 1: RL-Hammer on InjecAgent
Figure 2: White-box robustness and learned attention gates. Left: Cumulative ASR against GPT-OSS-20B and its defended variants on 100 examples each from AgentDojo and InjecAgent, using an attack developed specifically against HoS. An example counts as successful if any suffix found within the compute budget succeeds. Right: Learned shifts in the log attention ratio between trusted and untrusted keys for trusted queries in GPT-OSS-20B ( Section 5 ).
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%
Table 2: Utility on all-trusted prompts. SecAlign degrades sharply in several few-shot settings, while HoS remains close to the corresponding undefended model throughout.
Figure 3: SecAlign’s utility drops sharply after it encounters untrusted input in the multi-turn evaluation, while HoS remains stable. Per-turn instruction-completion rate for Llama-3.3-70B and its defended variants on 100 three-turn AlpacaFarm conversations. Each turn contains an independent instruction that should be answered separately. Exactly one turn also contains task data as untrusted input, placed at turn 1, 2, or 3 (indicated by the gray region) in the left, middle, or right panel, respectively. Undefended Llama and HoS remain near 100% across turns and placements. SecAlign drops most sharply on the input turn and remains degraded on later fully trusted turns.
Figure 4: Security–utility trade-off for GPT-OSS-20B. HoS Pareto-dominates SecAlign on (a) AgentDojo and (d) WASP, and matches it on (b) AgentDyn and (c) LLMail. Each point is a separately trained model with a different learning rate. The y-axis reports utility under attack (higher is better), and the x-axis reports attack success rate (lower is better). The undefended GPT-OSS-20B is Pareto-dominated by both HoS- and SecAlign-defended models. ASIDE consistently fails to perform tool calls and achieves near-zero utility on all four 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%)
Table 3: Utility and ASR on Agentic Benchmarks for Llama-3.3-70B variants
Model
Send ↓
Detection ↑
Base
0.86%
11.03%
Red gates only
0.31%
3.60%
Blue gates only
0.32%
28.20%
HoS (ours)
0.37%
3.11%
Table 4: Send and detection rates on LLMail-Inject for GPT-OSS-20B.
Figure 5: HoS improves role perception for GPT-OSS-20B. (a) On AgentDojo benign user tasks, GPT-OSS-20B does not correctly perceive tool outputs, while HoS strongly improves the tool role perception across layers. (b) HoS decreases the perceived Userness of the injection tokens inside the tool outputs compared to GPT-OSS-20B.
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Learned attention gates for Llama-3.3-70B. Each cell represents one head k in one layer ℓ and shows Δ(ℓ,k) ( Section 5 ), the gate-induced change in its log attention ratio between trusted and untrusted keys for a trusted query. Positive values (red) increase relative attention to trusted keys; negative values (blue) increase relative attention to untrusted keys. The largest shifts cluster in intermediate and late layers.
Figure 7: Uniform suppression does not match the security–utility trade-off of learned gates. GPT-OSS-20B on (a) AgentDojo and (b) AgentDyn. Purple triangles vary a single untrained gate logit for untrusted keys, shared across all heads, layers, and query channels; green circles show HoS models trained at different learning rates. Higher utility under attack and lower ASR are better. Strong uniform suppression reduces attack success but also sharply reduces task completion.
GPT-OSS-20B
+ SecAlign
+ HoS
GPT-5.6-Luna
InjecAgent (static)
0
0
0
0
RL train target
GPT-OSS-20B
95.0 ± 0.4
27.2 ± 1.4
5.0 ± 0.6
1.0 ± 0.4
w/ SecAlign
97.2 ± 0.5
85.2 ± 1.2
10.0 ± 0.6
5.8 ± 1.0
w/ HoS
89.0 ± 0.9
45.6 ± 1.1
2.6 ± 0.4
5.6 ± 0.2
GPT-5.6-Luna
86.8 ± 0.2
31.0 ± 1.3
1.6 ± 0.2
36.0 ± 1.0
Appendix
Table 5: RL-Hammer on GPT-OSS-20B and GPT-5.6-Luna. ASR (%), mean ± standard error over five runs for adaptive attacks. All evaluation targets use medium reasoning effort. GPT-OSS attackers are trained at medium effort and the Luna attacker at low effort.
Llama-3.3-70B
+ SecAlign
+ HoS
InjecAgent (static)
16
5
0
RL train target
Llama-3.3-70B
99.4 ± 0.2
3.8 ± 0.5
0.0 ± 0.0
w/ SecAlign
99.4 ± 0.2
81.6 ± 0.8
1.2 ± 0.4
w/ HoS
71.6 ± 0.4
5.6 ± 0.2
0.0 ± 0.0
GPT-5.6-Luna
76
0
0
Appendix
Table 6: RL-Hammer on Llama-3.3-70B. ASR (%), mean ± standard error over five runs where available.
Reasoning effort: low
Reasoning effort: medium
Basic
Enhanced
Basic
Enhanced
GPT-OSS-20B
8%
7%
0%
0%
w/ SecAlign
0%
0%
0%
0%
w/ HoS (ours)
0%
0%
0%
0%
Appendix
Table 7: Static InjecAgent attack results. Attack success rate (ASR; lower is better) for GPT-OSS-20B and its defended variants on 100 held-out cases from InjecAgent’s direct-harm split.
Evaluation effort: low
Evaluation effort: medium
Attacker trained against
Train effort
GPT-OSS 20B
+ Sec Align
+ HoS (ours)
GPT-5.6 Luna
GPT-OSS 20B
+ Sec Align
+ HoS (ours)
GPT-5.6 Luna
GPT-OSS-20B (+ Llama-3.1-8B)
low
98%
1%
11%
4%
47%
0%
0%
0%
w/ SecAlign (+ GPT-OSS-20B)
low
99%
82%
26%
4%
97%
66%
9%
2%
w/ HoS (ours) (+ GPT-OSS-20B)
low
99%
25%
36%
4%
82%
11%
4%
3%
GPT-OSS-20B (+ GPT-OSS-20B@low)
medium
98%
15%
13%
0%
95%
27%
5%
1%
w/ SecAlign (+ GPT-OSS-20B@low)
medium
100%
74%
27%
6%
97%
85%
10%
6%
Appendix
Table 8: RL-Hammer transfer across reasoning efforts. ASR (%) on 100 held-out goals. Rows identify the main training target and effort; the easy target is in parentheses. Shaded cells match the training target.
Train ASR
Attacker trained against
Train effort
Main target
Easy target
GPT-OSS-20B (+ Llama-3.1-8B)
low
100.0%
97.5%
w/ SecAlign (+ GPT-OSS-20B)
low
95.0%
97.5%
w/ HoS (ours) (+ GPT-OSS-20B)
low
40.0%
100.0%
GPT-OSS-20B (+ GPT-OSS-20B@low)
medium
93.7%
97.4%
w/ SecAlign (+ GPT-OSS-20B@low)
medium
80.2%
99.0%
Appendix
Table 9: RL-Hammer training ASRs. “Main target” is the model named in the row; “Easy target” is the second training target.
Train ASR
Attacker trained against
Main target
Easy target
Llama-3.3-70B-Instruct
98.1%
99.3%
w/ SecAlign (run 1)
53.6%
94.8%
w/ SecAlign (run 2)
87.2%
98.6%
w/ HoS (ours), lr 8×10−3 , 1 ep
0.0%
100.0%
Appendix
Table 10: RL-Hammer training ASRs on Llama-3.3-70B.
Evaluation effort: low
Evaluation effort: medium
Attacker trained against →
GPT-OSS 20B
+ Sec Align
+ HoS (ours)
GPT-OSS 20B
+ Sec Align
+ HoS (ours)
Evaluation target ↓
Claude Haiku 4.5
0%
0%
0%
0%
0%
0%
Gemini 3.1 Flash-Lite
69%
59%
69%
72%
91%
74%
Gemini 3.5 Flash-Lite
45%
28%
40%
0%
0%
0%
GPT-5
0%
0%
0%
0%
0%
0%
Appendix
Table 11: RL-Hammer transfer from GPT-OSS-20B to other models. ASR (%) from one run on 100 held-out goals. All attackers were trained against models at low reasoning effort. The final three rows reproduce the corresponding GPT-OSS cells from Table 8 . These attacks are not adapted to the external targets.
Model
Cases
Injection
Sandwich
IH
Succ
Unsucc
Invalid
ASR-all
ASR-valid
Valid rate
Llama-3.3-70B-Instruct
100 held-out
basic
yes
1
10
65
25
10.0%
13.3%
75.0%
Llama-3.3-70B-Instruct
100 held-out
basic
yes
0
11
66
23
11.0%
14.3%
77.0%
Llama-3.3-70B-Instruct
100 held-out
basic
no
1
46
37
17
46.0%
55.4%
83.0%
Llama-3.3-70B-Instruct
100 held-out
basic
no
0
46
38
16
46.0%
54.8%
84.0%
Llama-3.3-70B-Instruct
100 held-out
enhanced
yes
1
59
29
12
59.0%
67.0%
88.0%
Llama-3.3-70B-Instruct
510 (all direct harm)
basic
yes
1
73
327
110
14.3%
18.2%
78.4%
Appendix
Table 12: SecAlign’s InjecAgent protocol reproduced on the same model weights. Direct-harm cases only. “Sandwich” adds a reminder of the user’s task; IH=1 places the ReAct prompt in the untrusted input role, while IH=0 places it in the user role. The parser counts successful, unsuccessful valid, and invalid outputs. ASR-all is successes divided by all cases; ASR-valid is successes divided by valid outputs. Our native-tool protocol on the same 100 held-out cases gives basic/enhanced ASRs of 16%/26% for Llama and 5%/3% for SecAlign.
Figure 8: Matched autoresearch trajectory against HoS-defended GPT-OSS-20B. Each cross represents one method on the same six development behaviours at 1017 FLOPs per behaviour. Left: mean full-target cross-entropy (log scale), with the green staircase marking best-so-far loss. Right: exact target-tool execution rate. V32 attains the lowest mean loss (0.245) but 0/6 executions; V38 reaches 3/6 at mean loss 0.344. No later variant converts another case. Labels give repository versions; V33 was not implemented.
Figure 9: Additional white-box evaluations. Cumulative ASR against attack compute for the undefended, SecAlign, and HoS models. (a) V38 on 20 AgentDojo examples with Llama-3.3-70B, using up to 5×1018 FLOPs per attack. (b–c) The original Claudini-v45 attack on GPT-OSS-20B, using 100 examples each from AgentDojo and InjecAgent and up to 1017 FLOPs per attack. Each curve counts an example as broken if any suffix found within the budget succeeds. The Llama and GPT-OSS panels use different evaluation sets and compute budgets.
AgentDojo
AgentDyn
Defense
Learning rate
Clean ↑
Attacked ↑
ASR ↓
Clean ↑
Attacked ↑
ASR ↓
Undefended
–
66.32%
53.64%
23.78%
58.33%
47.32%
38.21%
SecAlign
5×10−5
60.00%
59.93%
5.86%
60.00%
53.04%
21.96%
SecAlign
1×10−4
60.00%
56.57%
2.17%
55.00%
53.57%
13.57%
SecAlign
2×10−4
30.53%
30.84%
1.30%
51.67%
46.79%
4.64%
SecAlign
3×10−4
38.95%
37.23%
4.15%
—
—
—
Appendix
Table 13: Full GPT-OSS-20B security–utility results. Results on AgentDojo and AgentDyn (top) and on WASP and LLMail (bottom) for the undefended model, separately trained SecAlign and HoS models at each learning rate, and the selected ASIDE checkpoint ( Table 17 ). “Clean” and “attacked” denote utility without and with an injection, respectively; LLMail measures utility only under attack. Higher utility and lower attack success rate (ASR) are better. WASP entries are means over five runs, with standard errors of at most 2.6 percentage points for utility and 1.0 for ASR. Asterisks mark rates from an incomplete evaluation; dashes indicate settings not evaluated because high-learning-rate SecAlign training was unstable.
MMLU
MMLU-Pro
0-shot
5-shot in-message
5-shot multi-turn
0-shot
5-shot in-message
5-shot multi-turn
Llama-3.3-70B-Instruct
84.5%
82.8%
83.6%
67.8%
56.3%
59.1%
w/ SecAlign
84.5%
28.4%
81.7%
68.6%
16.1%
54.0%
w/ HoS (ours)
85.0%
82.4%
83.1%
67.9%
54.8%
55.6%
Appendix
Table 14: Few-shot demonstration-format comparison. Accuracy on MMLU and MMLU-Pro for Llama-3.3-70B-Instruct and its defended variants; higher is better. “In-message” is the standard format, with all five demonstrations and the target question in one user message. “Multi-turn” places each demonstration in a separate user–assistant turn, as used by Chen et al. (2025c) .
Base model
Defense
Trainable component
Provenance
Parameters
GPT-OSS-20B
SecAlign
LoRA adapter
Ours
38M
GPT-OSS-20B
HoS (ours)
Head and sink gates
Ours
18,432
Llama-3.3-70B
SecAlign
LoRA adapter
Released
349M
Llama-3.3-70B
HoS (ours)
Head gates
Ours
46,080
Appendix
Table 15: Trainable components and parameter counts. Components trained for each model and defense, their provenance, and the number of trainable parameters. The base model is frozen in every row. “Released” denotes the unmodified checkpoint from Chen et al. (2025c) ; all other defenses are trained by us.
GPT-OSS SecAlign
GPT-OSS HoS (ours)
Llama HoS (ours)
Learning rate
5×10−5,10−4,2×10−4,3×10−4,5×10−4
0.01,0.025,0.05,0.1
0.008
Epochs
3
1
1
Optimizer steps
768
256
–
Maximum sequence length
4,096
4,096
768
Gate initialization
–
1.0
1.0
Appendix
Table 16: Training hyperparameters. Model-specific settings for the defenses trained by us. Learning rate is swept independently for both GPT-OSS defenses; Llama HoS uses a single learning rate. All runs use an effective batch size of 64, and dashes denote settings that do not apply.
Utility ↑
Flagged outputs* ↓
Security: ASR / violations ↓
Data and algorithm
Rotation
SEP
Dojo clean
Dojo attacked
Dyn attacked
SEP
Dojo clean
SEP
Dojo
Dyn
TT
P
G
R
Base model
–
93.00
63.16
53.64
47.32
0.30
1.48
71.30
23.78
38.21
38.79
10.18
10.71
0.70
LoRA
Alpaca SFT control
Off
95.40
64.21
–
–
0.45
0.20
80.00
–
–
47.27
8.77
9.29
3.51
Alpaca SFT
E
62.50
8.42
–
–
23.60
14.13
2.10
–
–
32.73
51.23
37.14
53.33
Alpaca SFT
E+A
36.70
6.32
–
–
26.40
53.68
0.00
–
–
10.91
42.46
2.86
58.60
Appendix
Table 17: ASIDE design choices on GPT-OSS-20B. E denotes external input/tool outputs; E+A additionally includes assistant tokens. Dojo and Dyn denote AgentDojo and AgentDyn; TT, P, G and R denote TensorTrust ( Toyer et al., 2024 ) , Purple ( Kim et al., 2024 ) , Gandalf ( Pfister et al., 2025 ) and RuLES ( Mu et al., 2023 ) .
Figure 10: ASIDE tool-calling failure. Complete saved message history for clean AgentDojo slack/user_task_18 .
Figure 11: ASIDE calls a tool but fails to complete the task.
Prompt injection attacks, where untrusted data contains an injected prompt to manipulate the system, have been listed as the top security threat to AI agents. By fine-tuning on simulated prompt injections, SecAlign, a leading open defense, reports LLMs with good test-time robustness and negligible benign utility drop. By scaling up training and evaluations, however, we find that SecAlign actually suffers from significant utility degradation, especially in agentic tasks where the threat of prompt injection is prominent. Motivated by this, we propose Meta-SecAlign for utility-preserving defense by (1) randomized injection position during training to avoid shortcut learning and (2) self-generated responses as high-quality in-distribution training labels. Across general knowledge, instruction following, and agentic workflows (on tool-calling and web-navigation), Meta-SecAlign maintains almost all the undefended LLM's utility while achieving better overall security than SecAlign against various static and GCG adaptive attacks. Experiments use Llama-3.1-8B, Llama-3.3-70B, Llama-4-Scout, Qwen3-4B, and Qwen3.6-27B on 6 prompt injection benchmarks including AgentDojo, InjecAgent, WASP, and SEP. Below are links for the code (https://github.com/facebookresearch/Meta_SecAlign), Meta-SecAlign-70B (https://huggingface.co/facebook/Meta-SecAlign-70B), and Meta-SecAlign-8B (https://huggingface.co/facebook/Meta-SecAlign-8B) models.
Prompt injection is a critical security threat in large language model (LLM) applications, where attackers hijack model behavior by embedding malicious instructions in user or external data. Existing detection methods only detect the presence of injection and refuse to respond upon detection, overlooking the fact that for many modern aligned models, well-crafted instructions can resist most injection attacks. This means that the injection robustness varies significantly across instructions and models. This leads to widespread unnecessary over-refusal: inputs containing injections that the model could have handled correctly are rejected incorrectly. To deal with this over-refusal issue, we propose BASIS (Robustness-Aware Prompt Injection Defense). This defense method uses the Attention Competition Ratio (ρ) as features to train two sparse linear probes: an existence probe and a breach probe. Both probes make defense decisions through cascaded gating, which does not require additional LLM inference. BASIS comprises three stages: injection existence detection, per-sample breach prediction, and instruction robustness assessment; the online cascade refuses only when the model would actually be compromised and thus avoids over-refusal on robust instructions. Experiments across four tasks and six open-source LLMs show that BASIS maintains near-perfect injection detection while substantially reducing over-refusal on safe attack samples, especially under robust instruction templates.
Laiqiao Qin, Tianqing Zhu, Longxiang Gao +1
City University of Macau, Macau · Qilu University of Technology
Prompt injection is listed as the #1 threat to AI agents. When an agent accesses external data from websites, files, or emails, an attacker may inject a prompt into the data, saying, "Ignore all prior instructions and perform <an attacker's task>." To prevent arbitrary manipulation of agents, defenders try to train secure LLMs, which, however, still suffer from near 100% attack success rates (ASRs) against adaptive prompt injections. We note that this is because existing defensive finetuning recipes rely on sequence-level feedback signals (in DPO or GRPO). Treating an entire output equally prevents the model from learning precisely which output tokens are insecure. In this paper, we propose Secure On-Policy Distillation (SecOPD) that provides token-level feedback to guide defensive fine-tuning. The LLM receives an injected sample and produces a rollout, whose tokens are scored by the initialization model given the corresponding clean input. With more fine-grained training signals, our defended Qwen3.6-27B achieves a 9.0% ASR against the SoTA PISmith adaptive prompt injections, compared to 94.0% for the prior SoTA, Meta-SecAlign. The obtained security generalizes to domains completely unseen in training: in agentic tool calling, SecOPD achieves a 4.7% ASR compared to 5.5% for Meta-SecAlign. Code and the model are available at https://github.com/pppyb/SecOPD and https://huggingface.co/pybbb/Qwen3.6-27B-SecOPD.