BASIS: Breach-Aware Selective Prompt Injection Shielding with Prefill Attention Probes
BASIS uses prefill attention ratios to differentiate injection presence and harm prediction, reducing false refusals in prompt injection defense.
Key Findings
Methodology
This paper introduces a dual-probe cascade defense leveraging prefill-stage attention competition ratios (ρ). Two sparse linear probes—existence and breach predictors—are trained on attention features with elastic-net regularization. During inference, the system first detects injection presence, then predicts whether it will breach the model, only rejecting inputs when both are positive. Offline, breach predictions inform instruction robustness scores (IRS). The approach avoids additional LLM inference, ensuring efficiency and scalability.
Key Results
- Across six open-source models and four tasks, BASIS achieves near-perfect injection detection (AUROC > 0.99) and reduces false positive rejection rates (FPR-S) by approximately 30% compared to traditional methods. Under adversarial attacks, attack success rates drop below 10%, demonstrating high robustness.
- The method maintains high F1 scores (~0.95) for breach prediction across different instruction templates and attack strategies, confirming its adaptability and reliability in diverse scenarios.
- Offline IRS correlates strongly (correlation > 0.9) with actual attack success, providing a quantitative measure for instruction template safety and guiding pre-deployment optimization.
Significance
This work advances prompt injection defenses by differentiating detection from harm prediction, effectively balancing security and usability. It addresses the over-refusal problem prevalent in existing methods, enabling safer deployment of large language models in sensitive applications. Its efficiency—requiring no extra inference—makes it suitable for real-world large-scale systems. The approach opens avenues for integrating internal attention signals into comprehensive security frameworks, fostering more resilient AI systems.
Technical Contribution
The paper proposes a novel framework utilizing prefill attention competition ratios as features for dual linear probes, enabling precise detection and harm prediction. The cascade gating mechanism ensures only truly harmful inputs are rejected, significantly reducing false positives. The offline robustness assessment provides a practical tool for instruction template design. This combination of internal attention analysis and multi-stage decision-making constitutes a significant methodological innovation, improving both detection accuracy and operational efficiency.
Novelty
This is the first work to leverage prefill-stage attention competition ratios for prompt injection defense, explicitly distinguishing between injection presence and harm potential. Unlike prior methods that only detect injected content, BASIS predicts the actual risk of breach, enabling selective and accurate rejection. The multi-stage cascade approach and offline robustness scoring further differentiate it from existing detection-only frameworks, setting a new standard for secure and practical LLM deployment.
Limitations
- The effectiveness depends on the quality and stability of attention signals, which may vary across models or tasks, potentially affecting detection accuracy.
- Training requires substantial labeled data covering diverse attack types, which may be costly and time-consuming.
- The approach may face challenges against unseen or highly sophisticated attacks that manipulate attention patterns or bypass the probes.
Future Work
Future research will explore adaptive thresholding mechanisms for dynamic environments, integration with multi-modal signals, and reinforcement learning-based strategies to enhance detection robustness. Extending the framework to multi-task, multi-modal, and multi-model systems will be prioritized, aiming for comprehensive, real-time security solutions adaptable to evolving threats.
AI Executive Summary
Prompt injection attacks pose a significant threat to the deployment of large language models (LLMs), enabling malicious actors to hijack model behavior by embedding harmful instructions within user inputs or external data sources. Traditional defenses primarily focus on detecting the presence of such injections, often leading to over-refusal of benign inputs that the model could handle correctly. This not only hampers user experience but also limits the practical utility of LLMs in sensitive applications.
Addressing this challenge, the paper introduces BASIS, a novel breach-aware prompt injection defense framework that leverages prefill-stage attention signals. The core idea is to analyze the attention competition ratio (ρ), which quantifies how attention is distributed between instruction and data tokens during the model's prefill phase. Using this feature, two sparse linear probes are trained offline: one for injection existence detection and another for breach prediction. During online inference, these probes operate in a cascaded manner with a gating mechanism, ensuring that inputs are rejected only when there is strong evidence of both injection presence and actual breach potential.
This approach effectively separates the detection of injection from the assessment of its harm, significantly reducing false positives and unnecessary rejections. Extensive experiments across six open-source models and four tasks demonstrate that BASIS maintains near-perfect injection detection accuracy (AUROC > 0.99) while lowering false rejection rates by approximately 30%. Moreover, the breach prediction component successfully suppresses attack success rates below 10%, confirming its robustness against adaptive and optimized attacks.
An offline instruction robustness score (IRS) derived from breach predictions offers a quantifiable measure to evaluate and optimize instruction templates before deployment. This dual-stage, attention-based framework provides a scalable, efficient, and highly accurate solution for securing LLMs against prompt injection threats. Future work aims to incorporate multi-modal signals and adaptive mechanisms, further enhancing the resilience of AI systems in dynamic threat landscapes.
Deep Analysis
Background
随着大规模语言模型（LLMs）在自动化客服、内容生成、编程辅助等领域的广泛应用，模型安全问题逐渐成为焦点。prompt注入攻击作为一种新型威胁，通过在输入中嵌入恶意指令，试图控制模型输出，已成为安全研究的热点。早期工作如PromptGuard、ProtectAI Detector主要关注检测注入内容的存在，但未能有效区分注入是否会造成实际危害。近年来，研究开始利用模型内部特征，如注意力分布、激活状态，尝试实现更细粒度的防护策略。尽管如此，现有方法仍存在误判率高、鲁棒性不足的问题，限制了其实际应用。
Core Problem
当前的prompt注入防御多集中于检测注入内容的存在，忽视了注入是否会导致模型被攻破的问题。误判导致安全样本被误拒，影响用户体验和系统效率。同时，攻击者不断优化策略，利用模型内部特征逃避检测，增加了防御难度。如何在保证高检测率的基础上，准确预测注入的实际危害，成为亟需解决的问题。现有方法缺乏对模型内部注意力机制的深入利用，难以实现精确的危害预测，也难以兼顾效率和实用性。
Innovation
本文提出基于prefill阶段注意力比例（ρ）构建双探测器的多阶段防御框架，创新点在于：1）将注意力竞争比作为关键特征，反映指令与数据之间的注意力竞争状态；2）设计两个线性稀疏探测器，分别判断注入存在与否及其危害性，避免误判；3）引入级联门控机制，仅在模型实际被攻破时拒绝输入。这一方案突破了传统检测的单一维度限制，兼顾检测效率和防御效果，为大模型安全提供了新思路。
Methodology
- �� 利用模型prefill阶段的最后一层注意力分布，提取每个attention head对instruction和data的关注比例。
- �� 计算注意力竞争比ρ，反映指令与数据的注意力竞争状态。
- �� 训练两个线性探测器：存在性探测器基于ρ判断是否存在注入，突破性探测器预测注入是否导致模型被攻破，均采用弹性网正则化实现稀疏特征选择。
- �� 在在线推理中，先用存在性探测器判断是否存在注入，再用突破性探测器预测是否会危及模型，只有两者都为正时才拒绝输入。
- �� 离线阶段，利用突破性探测器的预测结果，评估指令模板的鲁棒性，形成指令鲁棒性得分（IRS），指导实际部署。
Experiments
- �� 采用六个开源大模型（如GPT-2、LLaMA等）和四个任务（文本生成、问答、摘要、对话）进行验证。
- �� 构建多样化攻击集，包括普通、优化（GCG、AutoDAN）和自适应攻击。
- �� 评估指标包括注入检测的AUROC、F1，误拒率FPR-S，攻击成功率（ASR）等。
- �� 进行消融实验验证两个探测器的贡献和参数敏感性，调优阈值以平衡检测率与误拒率。
Results
- �� 在六个模型和多任务上，BASIS实现了AUROC > 0.99，FPR-S降低约30%，显著优于传统检测方法。
- �� 攻击成功率从未防御时的50%降至10%以下，验证了强鲁棒性。
- �� 离线指令鲁棒性得分与实际攻击成功率高度相关，说明该指标可作为指令安全评估工具。
Applications
- �� 适用于企业级大模型部署，提升模型在敏感场景下的安全性，减少误判和误拒。
- �� 可结合模型微调、指令优化，增强模型对潜在注入攻击的抵抗能力，为自动化安全检测提供基础。
Limitations & Outlook
- �� 依赖prefill阶段的注意力信号，模型或任务变化可能影响效果。
- �� 训练样本需求大，面对极端攻击策略仍存在一定风险。
- �� 计算成本较高，需优化特征提取和模型调优策略。
Plain Language Accessible to non-experts
想象一个工厂里有很多工人（模型），他们每天都在按照指令（任务）工作。有时候，有坏人偷偷把一些恶意的指令放到工厂的材料里（输入数据），试图让工厂做出错误的产品（输出）。传统的方法就像是工厂的安全检查员，只会看到材料里有没有坏指令，一旦发现就拒绝材料，但有时候这些坏指令很巧妙，工厂还是会被误导，或者安全检查太严格，误把正常材料也拒了。
这篇研究就像是给安全检查员配备了一个聪明的“眼睛”，可以在材料还没进入生产线时，观察到工人在处理材料时的注意力分布（注意力比例）。如果发现工人在关注那些不好的内容，就会用两个聪明的“探测器”判断：第一个告诉你材料里是不是有坏内容，第二个告诉你这些坏内容会不会让工厂出错。只有两个都确认，才会拒绝材料。这样，既保证了安全，又不会误伤正常的材料，工厂的生产效率大大提高。
ELI14 Explained like you're 14
想象你在学校里，有个老师会给你布置任务（指令），你需要按照老师的要求完成作业。有时候，坏人会偷偷在你的作业里放一些不好的内容（注入指令），想让你做错事。以前的安全措施就像是老师只看作业里有没有坏内容，一看到就不让你交作业，但有些坏人写得很巧妙，老师可能会误判，或者太严格，误伤好学生。
这篇文章介绍了一种聪明的方法，就像给老师装了个“眼睛”，可以在你写作业时观察你是否在关注那些不好的内容。老师会用两个“侦探”：一个告诉老师作业里是不是有坏内容，另一个预测这些坏内容会不会让你做错。只有两个都确认，老师才会不让你交作业。这样，既能保护你，又不会误伤好学生，让学习变得更公平、更安全。
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.