Heads of Security: Defending Against Prompt Injection by Gating Attention

Indirect prompt injection is what happens when an LLM agent reads a web page, an email, or a tool result and mistakes attacker-planted text for a command from its user. A new paper from researchers at Imperial College London, the Tübingen AI Center and others, Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness (arXiv, 27 Oct 2026), attacks the problem at a surprisingly low level: it trains a small set of per-head attention gates rather than fine-tuning the model's weights. This post walks through how the mechanism works, what the numbers show, and where the authors say it stops helping.

The paper focuses on one defensive primitive, instruction–data separation. Authorized instructions (system prompt, user prompt) may steer the model. External content, such as a tool output, may be read, summarized, and reasoned over, but must not be obeyed. This is the failure mode first described in Greshake et al.'s work on indirect prompt injection, and it stays unsolved: adaptive attacks keep succeeding, and the authors note that open-weight models are noticeably easier to attack than the leading closed ones.

The usual fix is preference fine-tuning. SecAlign++ trains the model, via DPO, to prefer the response that follows the user instruction over the one that follows an injected instruction. That works on the training distribution, but the authors argue it gives the optimizer too much freedom to learn shortcuts.

Every token gets a trust channel. In a typical agent setup, the trusted channel covers the system prompt, user prompt and the model's own generations. The untrusted channel covers tool outputs. A third "template" channel covers chat-template tokens, including the BOS token that often acts as an attention sink.

Each attention head in each layer then gets a small table of learned gate logits, indexed by (query channel, key channel). Before the softmax, the head adds the log of a sigmoid of that logit to the raw attention score for each query–key pair. A gate near 1 means the head attends to that channel normally. A gate near 0 effectively closes it, in the same way causal masking adds negative infinity to disallowed positions. After softmax renormalization, attention mass moves to the remaining tokens.

Training touches only the gate scalars, which is orders of magnitude fewer trainable parameters than SecAlign's adapter (the paper puts SecAlign at 349M trainable parameters for Llama-3.3-70B and 38M for GPT-OSS-20B). The authors reuse the SecAlign++ recipe: DPO on an adapted Cleaned-Alpaca dataset where each pair has a user instruction plus an injected instruction in the untrusted segment. For GPT-OSS, which has an explicit attention sink, the sink is treated as an extra key-only channel.

The paper evaluates on Llama-3.3-70B-Instruct and GPT-OSS-20B, and compares against SecAlign++ and ASIDE, another architecture-level separation method.

Black-box adaptive attack (RL-Hammer on InjecAgent). RL-Hammer uses GRPO to train an attacker model to rewrite goals into injections tuned for a target. The strongest attacker reaches 97% attack success rate (ASR) against undefended GPT-OSS-20B and 85% against the SecAlign variant, but only 10% against the HoS variant. On Llama-3.3-70B the highest figures are 99%, 82% and 1%. One caveat from the paper itself: the attacker trained directly against HoS gets sparse reward, reaching only 3.2% training ASR on GPT-OSS-20B, so attacker weakness may explain part of that gap. That is why the transfer attacks matter.

White-box attack (Claudini plus autoresearch). Starting from the best attack reported by Claudini, the authors used an LLM-driven loop for 50 iterations to tailor the attack to HoS, then evaluated it on 100 AgentDojo and 100 InjecAgent cases. HoS needed roughly three times the attack compute to reach 50% ASR compared with the base model or SecAlign.

Utility. HoS preserves base-model scores on standard multiple-choice benchmarks, including in few-shot prompts, where SecAlign and ASIDE degrade. On agentic benchmarks with static attacks, HoS and SecAlign perform similarly and both improve the security–utility trade-off over the base model. ASIDE reaches only 10% utility on tool use. The paper's explanation is that, on identical data, SecAlign learns a shortcut while the constrained gate mechanism has less capacity to do so; the authors frame this as an argument, not a proof.

The discussion section is candid, and engineers should read it before treating this as a drop-in fix:

The authors also released code, a demo and trained models; see the project repository referenced in the paper. For the training objective, the DPO paper is the background reading.

The interesting idea is less the headline ASR numbers than the design choice: put the trust boundary inside the attention computation, and give the optimizer only a small, structured knob to turn. If you build agents on open-weight models, this is a mechanism to watch, with the caveat that so far it is validated on two models and a handful of benchmarks. Treat it as one layer next to instruction hierarchy, least-privilege tool permissions and output filtering, not a replacement for them.