Learned attention gates cut prompt-injection success in two open-weight models

Heads of Security, an arXiv preprint, tests a prompt-injection defence inside the model’s attention calculation. It lets an open-weight model read untrusted material while limiting how that material influences its response. The authors report sharply lower adaptive attack success on GPT-OSS-20B and Llama-3.3-70B-Instruct, with less utility loss than several tested defences.[1]

The method assigns every token to a channel. System and user instructions, model output and chat-template tokens can be treated as trusted or structural; tool results are normally marked untrusted. Each attention head receives a small learned table that adjusts attention between channel pairs before the softmax. The base model is frozen. Only the gates are trained, amounting to 18,432 scalar parameters for GPT-OSS-20B and 46,080 for Llama-3.3-70B-Instruct in the authors' implementation.[1][2]

A security boundary inside the model

The design resembles a data-flow rule more than a stronger warning in the prompt. A query token can still attend to an untrusted key, so the model can extract facts from a web page or email. The learned gate changes the weight given to that channel when the model decides what to do. This is meant to preserve the informational value of tool output without treating every sentence inside it as an instruction with the same authority as the user's request.[1]

After an automated search strengthened the white-box Claudini attack against the defence, the attack required about three times as much computation to reach 50% success as against the base or SecAlign models. This is an increased attack cost under the paper’s setup, not a guarantee against an attacker with a larger budget.[1]

The caller has to know what not to trust

The central limitation is architectural and operational. The model does not infer which tokens came from an untrusted source. The surrounding application supplies the labels. If a malicious document is accidentally placed in the trusted user channel, or a compromised tool is classified as trusted, the attention gates cannot restore a boundary that the application failed to express. Fine-grained systems may also need to distinguish trusted metadata from untrusted content within the same message.[1][2]

That requirement is a strength when provenance is clear. Agent frameworks already know whether text came from a system instruction, a user, a database result or a remote webpage. Making that provenance explicit at inference is more auditable than hoping the model will infer authority from prose. It is harder when data has been copied, summarised, concatenated or passed through several tools without a reliable lineage record.[1]

Deployment is also limited to systems whose attention implementation can be modified. The released package supports the two tested open-weight models through Hugging Face Transformers and vLLM. Hosted proprietary APIs generally do not expose the per-head operations required. Organisations adopting the method would still need tool permissions, sandboxing, output validation and monitoring because prompt injection is only one route to unsafe agent behaviour.[2][3]

Heads of Security does not settle the prompt-injection problem, but it clarifies where a durable defence may live. Instructions and data are different security principals even when they are presented to the same model as one sequence of tokens. Encoding that distinction in attention produced a better reported security-utility trade-off than relying only on alignment training. The next test is whether applications can maintain trustworthy provenance labels across messy, long-running workflows—the condition on which the model-level protection depends.[1][2]

Report an error or suggest a correction