BASIS: Breach-Aware Selective Prompt Injection Shielding with Prefill Attention Probes

2026-08-08

AI Summary

This work addresses the over-rejection problem in existing prompt injection defenses, which often erroneously block inputs that language models can safely process. The authors propose BASIS, a cascade detection mechanism based on the attention competition ratio (ρ), which performs two-stage probing—existence and privilege escalation—to reject responses only when the model is genuinely vulnerable, without incurring additional inference overhead. BASIS introduces, for the first time, an instruction robustness-aware strategy that leverages attention patterns from the prefilling phase to construct a lightweight, sparse linear probe, combined with cascaded gating for precise online decision-making. Experiments across four tasks and six open-source large language models demonstrate that BASIS achieves near-perfect injection detection rates while substantially reducing over-rejection of safe inputs under robust instructions.

Abstract

Prompt injection is a critical security threat in large language model (LLM) applications, where attackers hijack model behavior by embedding malicious instructions in user or external data. Existing detection methods only detect the presence of injection and refuse to respond upon detection, overlooking the fact that for many modern aligned models, well-crafted instructions can resist most injection attacks. This means that the injection robustness varies significantly across instructions and models. This leads to widespread unnecessary over-refusal: inputs containing injections that the model could have handled correctly are rejected incorrectly. To deal with this over-refusal issue, we propose BASIS (Robustness-Aware Prompt Injection Defense). This defense method uses the Attention Competition Ratio ($ρ$) as features to train two sparse linear probes: an existence probe and a breach probe. Both probes make defense decisions through cascaded gating, which does not require additional LLM inference. BASIS comprises three stages: injection existence detection, per-sample breach prediction, and instruction robustness assessment; the online cascade refuses only when the model would actually be compromised and thus avoids over-refusal on robust instructions. Experiments across four tasks and six open-source LLMs show that BASIS maintains near-perfect injection detection while substantially reducing over-refusal on safe attack samples, especially under robust instruction templates.

Problem

• prompt injection
• over-refusal
• large language models
• security
• robustness

Innovation

• Prompt Injection Defense
• Attention Competition Ratio
• Over-refusal Mitigation
• Robustness-aware Shielding
• Sparse Linear Probes