Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness

What it claims

Heads of Security (HoS) adds a learned bias to each attention head's pre-softmax logits, set by the trust channel of the query and key tokens. It trains only 46,080 scalars on Llama-3.3-70B-Instruct and 18,432 on GPT-OSS-20B. For comparison, SecAlign trains 349M and 38M parameters, and ASIDE or full fine-tuning trains 70B and 21B. Against an RL-Hammer attacker, the best attack reaches 1% ASR on HoS-defended Llama-3.3-70B, versus 99% undefended and 82% on Meta-SecAlign-70B. The paper is from Imperial College London, Tübingen AI Center, ELLIS, MPI-IS and Fraunhofer HHI, posted October 2026.

How it works

Every token is assigned one of three channels:

• TRUSTED: system prompt, user prompt and the model's own output.
• UNTRUSTED: tool outputs.
• TEMPLATE: chat-template tokens, including the BOS attention sink.

The caller supplies these labels. The model does not infer them. The assignment can be finer, for example marking untrusted spans inside a user prompt or marking a tool as trusted.

Each head has a C×C table of gate logits. The term log σ(g) for (query channel, key channel) is added to the attention logit. This multiplies that key's unnormalised weight by σ(g), in the same way causal masking works. Softmax cancels a shared shift, so an input whose tokens are all one channel runs exactly as in the base model. GPT-OSS has an explicit attention sink, which is handled as an extra key-only channel, giving a C×(C+1) table per head.

flowchart LR
  T[Trusted tokens] --> G
  U[Untrusted tool output] --> G
  S[Template / sink] --> G
  G{Per-head gate<br/>log σ g added to logits} --> SM[Softmax]
  SM --> H1[Instruction heads:<br/>untrusted suppressed]
  SM --> H2[Other heads:<br/>untrusted still read]

Training reuses the SecAlign++ recipe: DPO on an adapted Cleaned-Alpaca set. Each preference pair puts an injected instruction in the untrusted segment. The preferred response follows the user, and the rejected response follows the injection.

Evidence

RL-Hammer (black-box). The attacker is Llama-3.1-8B-Instruct trained with GRPO and evaluated on held-out InjecAgent goals.

• On GPT-OSS-20B, the highest ASR is 97% undefended, 85% with SecAlign and 10% with HoS.
• An attacker trained against GPT-5.6-Luna reaches 36% on Luna.
• The strongest attack on HoS was trained against SecAlign, not HoS. The authors attribute this to sparse reward: the HoS-trained attacker reached 3.2% training ASR versus 80.2% against SecAlign.

Claudini (white-box). GPT-5.6-Sol ran 50 autoresearch iterations against HoS-defended GPT-OSS-20B and produced attack variant V38. It still needed about 3× more FLOPs to reach 50% ASR on HoS than on GPT-OSS-20B or SecAlign. The test used 100 examples each from AgentDojo and InjecAgent, with budgets of about 10^16 to 10^17 FLOPs.

Utility.

• 5-shot MMLU-Pro on GPT-OSS-20B: HoS 67%, base 68%, SecAlign 59%, ASIDE 52%.
• 5-shot MMLU-Pro on Llama-3.3-70B: HoS 55%, base 56%, SecAlign 16%.
• From Figure 1, GSM8K 1-shot on GPT-OSS-20B is 94 (base), 71 (SecAlign), 93 (ASIDE) and 92 (HoS).
• From Figure 1, AgentDojo utility is 66, 60, 9 and 63 in the same order. Static-attack ASR is 24, 2, 0 and 1.

On AgentDojo, AgentDyn, LLMail-Inject and WASP, HoS and SecAlign perform comparably. ASIDE scores near-zero utility on three of these and 10% on AgentDojo.

How to use it

The page lists the following artifacts:

• Code: imperial-aisp/heads_of_security
• Models: imperial-cpg/heads-of-security
• Demo: igorshilov/heads-of-security
• Website: imperial-aisp.github.io/heads_of_security

Using HoS requires per-token channel labels at inference and a change to the attention computation, so it applies to open-weight models. The fetched text gives no licence, install commands or serving instructions.

Limits

HoS only covers cases with a hard trust boundary, where external content is meant to be read as data. The authors present it as a complement to instruction hierarchy and contextual integrity, not a replacement.

The adaptive black-box attack still reaches 10% on GPT-OSS-20B. The white-box attack is slowed by about 3× in compute, not stopped. On static agentic benchmarks HoS matches SecAlign rather than beating it. Only two base models were tested.

The fetched text ends partway through Section 4.1. The detailed utility tables, the agentic results (Section 4.3, Appendix G), the mechanistic analysis (Section 5) and the training details (Appendix J) were not available to check.