Attention Gating: A New Defense — How HoS Makes Prompt Injection Attacks Fail

Prompt injection attacks have long been the number one security threat to LLM application deployment. Attackers embed malicious instructions in user data to trick the model into deviating from its original task, leaking information, or performing unauthorized actions. Traditional defenses either rely on fine-tuning to change model behavior or forcibly isolate instructions from data at the architectural level, but they often struggle to balance security and utility. A recent paper proposes Heads of Security (HoS), which enforces instruction-data separation through a set of learnable per-head attention gates, demonstrating significant advantages across multiple models and attack scenarios.

Core Idea: Teaching Each Attention Head to "Treat Different Inputs Differently"

The key design of HoS is that it does not modify the model's main weights, but instead introduces a small gating module for each attention head. This module learns how to reallocate, within that head, the attention weights assigned to "trusted tokens" (such as system instructions and explicit user requests) and "untrusted tokens" (such as external documents and tool return results). The training objective is straightforward: the model should be able to understand the information in untrusted content, but must never execute the instructions embedded within it.

Unlike defenses such as ASIDE that also perform instruction-data separation at the architectural level, HoS's gating is learned rather than hard-coded. This means the model can adaptively identify which context segments are more likely to carry injection risk and dynamically weaken their instructional influence at the attention level.

Multi-head attention mechanism: Each layer of a Transformer consists of multiple parallel attention heads, each independently computing query (Q), key (K), and value (V) projections and producing attention weights to capture different types of dependencies. HoS's "per-head gating" operates at exactly this level: it equips each attention head with a small gating network that reweights attention weights before or after the softmax, allowing that head's attention allocation between "trusted tokens" and "untrusted tokens" to be learned and adjusted. Since the gating parameters are small in scale and do not touch the main weights, training costs are far lower than full-parameter fine-tuning, and catastrophic forgetting is avoided. This design allows the model to retain the parallel expressive power of multi-head attention while gaining a fine-grained instruction-data separation pathway.

Experimental Performance: Clear Advantages Against Adaptive Attacks

The paper evaluates HoS on Llama-3.3-70B-Instruct and GPT-OSS-20B, comparing against baselines including the fine-tuning-based defense SecAlign++ and the architecture-based defense ASIDE.

Most notably is its robustness against adaptive attacks. Under the black-box RL-Hammer attack, HoS achieves an attack success rate (ASR) of only 10%, while the original GPT-OSS-20B model reaches 97% and SecAlign also reaches 85%. Against the white-box attack Claudini, attackers need to consume approximately 3x the FLOPs to achieve the same ASR on HoS as on the baseline. This shows that HoS not only defends against static attacks but also significantly raises the cost of adaptive attacks.

In terms of utility, HoS maintains the base model's performance on standard multiple-choice benchmarks and prompt settings, while SecAlign and ASIDE degrade in few-shot scenarios. On the agent benchmark under static attacks, HoS performs similarly to SecAlign and substantially outperforms the base model's security-utility trade-off. ASIDE, meanwhile, performs poorly in tool-use scenarios, with utility of only 10%.

Adaptive attacks and white-box/black-box threat models: Adaptive attacks refer to attackers optimizing attack strategies specifically against a known defense mechanism, rather than using fixed static attack templates. In the black-box setting, attackers can only query model inputs and outputs (e.g., RL-Hammer uses reinforcement learning to search for effective injections); in the white-box setting, attackers can also access gradient information (e.g., Claudini directly performs gradient optimization on inputs). Evaluating adaptive attacks is key to measuring a defense's true robustness, because defenses that only resist static attacks often quickly fail once attackers adjust their strategies. FLOPs consumption is used to quantify attack cost: if a defense forces attackers to expend several times the computation to achieve the same attack success rate, it has effectively raised the attack threshold.

ASR (Attack Success Rate) and the security-utility trade-off: ASR refers to the proportion of test cases containing injection attacks in which the model actually executes the malicious instructions — lower is safer. The security-utility trade-off measures whether a defense harms the model's performance on normal tasks while improving security — an ideal defense should achieve both high security and high utility. In the text, SecAlign's degradation in few-shot scenarios and ASIDE's utility of only 10% in tool-use scenarios are manifestations of utility loss; HoS maintaining base model performance on standard benchmarks and few-shot settings indicates that its gating mechanism interferes little with normal reasoning pathways.

Why It Matters: Security and Utility Need Not Be Mutually Exclusive

The significance of HoS lies in providing a security-hardening path different from fine-tuning. Fine-tuning may harm the model's general capabilities, while architectural hard isolation easily damages normal tasks. HoS's per-head gating achieves "understand but do not obey" at a fine-grained level, preserving the model's reasoning ability over untrusted content while cutting off the execution chain of injected instructions.

The authors point out that prompt injection robustness can be achieved and improved through targeted attention intervention. This provides a scalable direction for future defense research: rather than chasing and blocking at the output end, it is better to teach the model to distinguish "information" from "instructions" at the attention allocation stage.

Of course, HoS is currently validated only on specific models, and the training cost of the gating module and its cross-model generalization ability remain to be further observed. But based on the existing results, the ASR gap of 10% versus 97% is already sufficient to make the point.

Key Takeaways

• HoS achieves instruction-data separation through learnable per-head attention gating without changing the main weights.
• On Llama-3.3-70B-Instruct and GPT-OSS-20B, HoS's robustness against adaptive attacks significantly outperforms SecAlign++ and ASIDE.
• Under the RL-Hammer attack, HoS's ASR is only 10%, far lower than the baselines; white-box attacks require approximately 3x the FLOPs to achieve the same effect.
• HoS maintains base model utility on standard benchmarks and few-shot settings, while some baselines degrade.
• This work demonstrates that attention-level intervention is an effective path for improving prompt injection defense.