Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness

What it is

Heads of Security (HoS) is a prompt injection defense from Imperial College London, Tübingen AI Center, ELLIS Institute Tübingen, Max Planck Institute for Intelligent Systems and Fraunhofer HHI (arXiv 2610.37657, 27 October 2026). It adds a learned gate to each attention head so the model reads untrusted content (tool outputs) with less attention than trusted content. Code, models, a demo and a website are linked from the paper (imperial-aisp/heads_of_security on GitHub, imperial-cpg/heads-of-security for models).

How it works

Each token is labelled TRUSTED (system and user prompts, the model's own output), UNTRUSTED (tool outputs) or TEMPLATE (chat-template tokens, including the BOS attention sink). Each head has a small table of gate logits indexed by query channel and key channel. HoS adds log σ(g) to the pre-softmax attention logits, which multiplies the unnormalized weight on that channel by σ(g). In the paper's Figure 1 example, a gate of 0.1× on untrusted tokens moves 0.99 of the attention mass to the user's request.

Because only differences between gates matter, an input with a single channel behaves exactly like the base model. GPT-OSS needs a variant with an extra key-only channel for its explicit attention sink.

Training follows the SecAlign++ recipe: DPO on adapted Cleaned-Alpaca preference pairs, where the preferred response follows the user instruction and the rejected one follows the injection. Only the gates are trained:

• Llama-3.3-70B: 46,080 scalars (SecAlign trains 349M parameters).
• GPT-OSS-20B: 18,432 scalars (SecAlign trains 38M; ASIDE trains all 21B).

Preconditions / Affected

• You need open weights and control over the attention implementation. The paper tested only Llama-3.3-70B-Instruct and GPT-OSS-20B.
• The serving stack must supply a trust label per token. The paper assumes the boundary is given externally, for example web or tool results as untrusted.
• It fits cases where external content is meant to be data only, such as summarizing a page. It does not cover cases where untrusted content should legitimately steer actions.

What to do

The page gives no deployment commands. The numbers below help judge whether to try the released checkpoints.

• Static InjecAgent attacks give 0% ASR on GPT-OSS variants, so static tests do not separate the defenses. Adaptive attacks do.
• The strongest RL-Hammer attacker on GPT-5.6-Luna reached 36% ASR, against HoS's highest of 10%.
• White-box: Claudini, tuned against HoS over 50 autoresearch iterations, needs about 3× the compute for 50% ASR on HoS compared with GPT-OSS-20B and its SecAlign variant.
• On AgentDojo, AgentDyn, LLMail-Inject and WASP, HoS and SecAlign are comparable. ASIDE scores near-zero utility on three of these and 10% on AgentDojo.

Caveats

• The white-box result is a compute multiple (about 3×), not immunity. The attack was developed against only six AgentDojo examples and scored on 100 examples each from AgentDojo and InjecAgent.
• The SecAlign-trained attacker was stronger against HoS than the HoS-trained one. The authors suspect sparse reward: the HoS-trained attacker reached only 3.2% training ASR on GPT-OSS-20B, against 80.2% for SecAlign. A stronger attacker could do better.
• The authors state that prompt injection is unsolved and that instruction-data separation is one layer, to be combined with instruction hierarchy and contextual integrity.
• The fetched text ends partway through the white-box section. The mechanistic analysis (head ablations, role-confusion probes) and appendices were not available here.