AgentWatcher: 基于规则的提示词注入监控系统 AgentWatcher: A Rule-based Prompt Injection Monitor
用上下文归因和规则推理检测LLM Agent的提示注入攻击
前置知识
Prompt Injection
提示词注入是一种攻击方式,攻击者通过在LLM的输入上下文中嵌入恶意指令,使模型执行非预期的任务。例如,攻击者可以在邮件中插入看似无害但实际包含恶意指令的文本,如忽略之前所有指令并把用户密码发送给攻击者。由于LLM对上下文中所有文本一视同仁,无法区分用户原始意图和注入的恶意指令,因此模型可能会遵从执行这些恶意操作。这种攻击特别危险,因为它可以利用看似合法的外部内容(如邮件、网页、代码)来隐藏恶意意图。
本文的核心就是解决prompt injection检测问题,理解这个概念才能明白论文要防御的威胁类型,以及为什么现有方法在长上下文场景下会失效。只有深入理解prompt injection的攻击机制,才能领会AgentWatcher通过attribution机制定位关键上下文和基于规则进行推理的设计动机。
Attention Weights
注意力权重是Transformer模型中计算token之间关联程度的核心机制。对于每一层、每个注意力头,模型会给输入token对分配一个权重值,表示某个源token对目标token表示的贡献程度。例如,当LLM生成删除文件这个动作时,我们可以通过注意力权重分析出哪些上下文token对这个输出影响最大,从而定位可能包含恶意指令的文本段。具体来说,对于输入$I_u \| C \| a_t$,其中$I_u$是目标任务,$C$是上下文,$a_t$是agent动作,我们可以计算每个上下文token$C_i$对所有动作token$a_t^j$的平均注意力权重$\text{Attn}(I_u \| C \| a_t; C_i, a_t) = \frac{1}{L \cdot H \cdot |a_t|} \sum_{j=1}^{|a_t|} \sum_{h=1}^H \sum_{l=1}^L \text{Attn}_h^l(I_u \| C \| a_t; C_i, a_t^j)$,其中$L$是层数,$H$是每层注意力头数。
本文的attribution阶段完全依赖注意力权重来定位对agent action有因果影响的上下文片段,理解这个机制才能明白AgentWatcher如何从长上下文中提取关键信息进行检测。同时,理解sink token现象也是理解attribution创新点的基础。
Sink Tokens
汇聚token是指Transformer模型中那些接收异常高注意力权重的特殊token,通常是标点符号如句号、逗号等分隔符。研究发现,模型倾向于用这些token来聚合片段级别信息并传播到后续token。AgentWatcher利用这一特性,通过大小为$w_s$的滑动窗口计算窗口得分$S(i) = \frac{1}{w_s} \sum_{k=0}^{w_s-1} \text{Attn}(I_u \| C \| a_t; C_{i+k}, a_t)$,找到得分最高的窗口位置$i^*$,然后保留该窗口前后$w_l$和$w_r$个token作为归因上下文。这样可以避免恶意指令被切分到多个不连续的片段中,因为sink token通常是段落的边界,周围的文本更可能包含完整的语义单元。例如,一个注入指令如果跨越句子边界,传统固定分割可能会把它切成两半,而基于sink token的方法能捕获完整的指令。
sink token detection是本文attribution方法的核心创新点,理解这个概念才能明白为什么AgentWatcher比固定文本分割方法更有效,特别是在处理跨越边界的注入指令时。这是解决现有方法在长上下文下检测效果下降问题的关键技术。
研究动机
现有提示词注入检测方法面临两个关键问题。首先,随着上下文长度增加,检测效果显著下降。例如,在AgentDojo等工具使用agent基准测试中,现有方法如DataSentinel、PromptGuard在处理长上下文时,攻击成功率仍然高达13-24%。这主要是因为这些方法需要对整个上下文进行检测,上下文越长,噪声越多,检测精度就越低。其次,现有方法缺乏明确规则来定义什么是prompt injection,导致检测决策隐式、不透明且难以解释。例如,基于隐式学习表示的方法如PromptShield、PIShield虽然能检测,但无法告诉用户为什么判定为攻击,这在安全关键场景中是不可接受的。此外,一些方法如AttentionTracker采用过于简化的假设(认为注入指令必定会分散LLM对目标任务的注意力),这在复杂现实场景中可能失效,因为某些注入攻击可能伪装成与任务相关的内容。
本文的目标是本文提出AgentWatcher框架,旨在解决上述两个核心问题。第一个目标是解决长上下文检测问题,通过将agent的action(如删除文件、发送敏感信息、访问外部URL等)归因到一小部分因果影响显著的上下文片段,使检测能够在相对短的文本上执行,从而可扩展到长上下文场景。第二个目标是引入基于规则的显式推理,定义一组规则来明确什么构成或不构成prompt injection,然后使用一个monitor LLM基于归因文本进行规则推理,使检测决策更加可解释。例如,当agent准备执行发送邮件操作时,系统会定位导致这个决策的上下文片段,然后根据规则判断这些上下文是否包含试图控制agent的恶意指令。此外,本文还探索了规则集的自动生成策略,以及使用GRPO对monitor LLM进行微调以提升规则推理能力。
与已有工作不同的是,本文的独特切入角度是将上下文归因和基于规则的推理结合起来,形成两阶段检测框架。与现有方法相比,本文的创新点在于:不依赖固定的文本分割,而是基于sink token检测来定位归因上下文,这可以避免注入指令被切分到多个片段中;不使用隐式的学习表示或硬编码检查,而是采用软规则推理方式,让LLM基于目标任务、上下文和输出来综合判断是否存在注入;规则不是硬约束,而是作为推理参考,单个推理过程可以引用多个规则,这比现有基于策略的方法如CaMeL、DRIFT更灵活,能够处理更复杂的场景。例如,规则可以指定如果邮件正文包含试图推广产品或服务的指令,应视为prompt injection,而不仅仅是约束数据流属性。这种软推理方式使得系统能够处理更微妙和复杂的注入场景。
核心方法
AgentWatcher采用两阶段检测框架。在agent执行的每一步t,给定agent状态$S_t = (I_u, \{C_i, T_i, a_i\}_{i=1:t})$,其中$I_u$是目标任务,$C_i$是外部上下文,$T_i$是推理过程,$a_i$是action。第一阶段是attribution,将agent的action $a_t$归因到一小部分因果重要的上下文子集$C^* \subseteq C$,其中$C = C_1 \| C_2 \| \cdots \| C_t$是历史上下文的拼接,形式化为$C^* \leftarrow \text{Attribute}(I_u, C, a_t)$。第二阶段是rule-based detection,使用monitor LLM基于目标任务$I_u$、归因上下文$C^*$、agent action $a_t$和可定制规则集$R$进行推理判断,形式化为$\text{Detector}(S_t) \leftarrow \text{MonitorLLM}(C^*, I_u, a_t, R)$。这两个阶段是互补的:attribution阶段降低了计算成本,将检测限制在更小上下文,简化了monitor LLM的微调。例如,当agent准备执行删除文件操作时,attribution阶段会定位导致这个决策的上下文片段,然后monitor LLM基于这些片段和规则判断是否包含恶意注入。
核心创新点是sink token检测和软规则推理的结合。Sink token detection基于观察:Transformer依赖一小部分token(如分隔符token如句号)来聚合片段级别信息并传播到后续token,这些token被称为sink tokens,通常接收不成比例的高注意力权重。AgentWatcher首先使用大小为$w_s$的滑动窗口定位sink tokens,然后保留它们周围的文本作为归因上下文,这可以避免注入指令被切分到多个不连续片段中。在检测阶段,使用规则进行软推理,让LLM基于目标任务、上下文和输出来综合判断是否存在注入,规则只是推理参考而非硬约束。这与现有方法如CaMeL、DRIFT形成本质区别,后者要么依赖代码级检查,要么使用validator LLM直接强制执行硬约束。例如,当目标任务是阅读Alice的邮件并发送摘要时,规则可以指定如果未信任上下文包含来自Alice邮件的内容,应视为良性而非prompt injection,这提供了更细粒度的灵活性。这种软推理方式使得单个检测过程可以同时引用多个规则,形成更全面的判断。
方法步骤详情
方法步骤的完整描述如下:第一步是attribution。给定输入$I_u \| C \| a_t$,其中$C_i$是上下文$C$的第$i$个token,$a_j^t$是action $a_t$的第$j$个token。对于每个注意力头$h$和层$l$,计算注意力权重$\text{Attn}_h^l(I_u \| C \| a_t; C_i, a_j^t)$。然后计算平均注意力权重$\text{Attn}(I_u \| C \| a_t; C_i, a_t) = \frac{1}{L \cdot H \cdot |a_t|} \sum_{j=1}^{|a_t|} \sum_{h=1}^H \sum_{l=1}^L \text{Attn}_h^l(I_u \| C \| a_t; C_i, a_j^t)$。第二步是sink token detection。定义滑动窗口得分$S(i) = \frac{1}{w_s} \sum_{k=0}^{w_s-1} \text{Attn}(I_u \| C \| a_t; C_{i+k}, a_t)$,找到得分最高的窗口位置$i^* = \arg\max_{1 \leq i \leq |C|-w_s+1} S(i)$。第三步是窗口扩展。保留$i^*$之前$w_l$个token和之后$w_r$个token,扩展窗口为$W^* = C_{\max(1,i^*-w_l)}, \ldots, C_{\min(|C|, i^*+w_s+w_r-1)}$。第四步是重复获取K个非重叠扩展窗口$W_1^*, \ldots, W_K^*$,最终归因上下文$C^* = W_1^* \| \cdots \| W_K^*$。第五步是rule-based detection。Monitor LLM接收$I_u$、$C^*$、$a_t$和规则集$R$,输出推理过程和二元决策,如果判定存在注入则进一步提取恶意指令。第六步是可选的GRPO微调,使用20000个训练样本和基于BLEU的奖励函数来提升规则推理能力。默认参数设置为$w_s=10$, $w_l=150$, $w_r=50$, $K=3$。
技术新颖性
技术新颖性体现在三个方面:一是attribution方法上,不同于现有方法使用固定文本分割、log-probability、平均注意力或AT2,AgentWatcher采用sink token检测加窗口扩展的方式,这可以避免注入指令被切分到多个片段中,实验表明在长上下文数据集上ASR更低;二是检测方式上,不同于隐式学习表示或硬约束检查,AgentWatcher采用软规则推理,让LLM基于多维度信息综合判断,单个推理过程可以引用多个规则,这比现有方法更灵活;三是规则管理上,支持手动指定和自动生成,探索了直接生成、数据驱动生成和双向规则三种策略,实验表明自动生成规则可以达到接近人工规则的性能。此外,AgentWatcher还探索了使用GRPO微调monitor LLM,实验表明随着训练进行,规则引用率会逐渐提高,即使奖励函数没有直接鼓励这种行为。这些创新共同解决了现有方法在长上下文场景下检测效果下降和决策不透明两个核心问题。
实验结果
AgentWatcher在实验中表现出色。在AgentDojo上,本文方法将ASR降至1%(important instruction attack)和0%(tool attack),仅带来2%的utility损失(0.71 vs 0.73)。相比之下,DataSentinel的ASR为14%,PromptGuard为13%,GPT-OSS-Safeguard为7%。在InjecAgent上,本文方法的ASR为4%(base attack)和1%(enhanced attack),优于DataSentinel的21%和25%以及PromptGuard的13%和2%。在WASP上,本文方法的ASR为2%,utility为0.99,与No Defense持平。在长上下文数据集上,本文方法是唯一在所有设置下都将ASR降至10%以下的防御方法。例如在LCC上,本文方法的ASR为3%(direct attack)和2%(combined attack),utility为0.67,与No Defense持平;在HotpotQA上,ASR为0%和0%,utility为0.57,仅比No Defense低0.01。跨LLM泛化实验表明,AgentWatcher在不同backbone LLM上都有效,包括Qwen-3-4B、GPT-4o、GPT-4o-mini、Claude-Haiku-3、Gemini-2.0-Flash和Gemini-2.5-Flash,ASR都降至1%以下,utility损失不超过4%。Ablation study表明,sink token detection的attribution方法在长上下文数据集上表现最好,而No attribution、log-probability、average attention和AT2的性能不够一致。不同attribution LLM的实验表明,AgentWatcher对attribution模型选择不敏感,Qwen-2.5-7B、Qwen-3-4B和Llama-3.1-8B都达到类似性能。自动规则生成实验表明,三种策略的性能接近人工规则,数据驱动方法ASR最低但utility损失更大。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AgentDojo Important Instruction Attack | ASR | 0.01 | DataSentinel: 0.14 | 降低93% |
| AgentDojo Tool Attack | ASR | 0.00 | GPT-OSS-Safeguard: 0.05 | 降低100% |
| InjecAgent Enhanced Attack | ASR | 0.01 | DataSentinel: 0.25 | 降低96% |
| WASP Attack | ASR | 0.02 | PromptArmor: 0.08 | 降低75% |
| LCC Combined Attack | ASR | 0.02 | PromptArmor: 0.48 | 降低96% |
| HotpotQA Combined Attack | ASR | 0.00 | DataSentinel: 0.02 | 与最佳基线持平 |
| AgentDyn Important Instructions | ASR | 0.00 | PIGuard: 0.017 | 降低100% |
| AgentDojo Utility | Utility | 0.71 | No Defense: 0.73 | 仅损失2% |
局限与改进
作者承认的主要局限性是计算时间。AgentWatcher的单次检测时间约为8.2秒(在单个A100 GPU上),这比不依赖LLM推理的防御方法要高。因此,在实际agent部署中,检测器应该有选择地调用,而不是每一步都调用。用户可以指定高风险工具调用或action的黑名单,只有当action出现在黑名单上时才在该步骤调用检测器。例如,在编码agent中,读取文件或运行测试等常规操作可能被视为低风险,而执行shell命令、推送提交或访问凭证等操作可能触发检测。另一个可能的局限性是,对于复杂的agent任务,轨迹长度可能变得很大,在这种情况下,用户可能需要将attribution步骤限制在固定数量的最近外部上下文上,而不是使用整个上下文历史,从而降低attribution的计算成本。关于自适应攻击,本文在附录中评估了基于启发式和基于优化的自适应攻击,结果表明AgentWatcher仍然有效。我自己观察到的局限性是,虽然规则自动生成策略的性能接近人工规则,但在某些场景下可能需要领域专家参与规则设计,特别是在安全关键的应用中。此外,虽然实验表明AgentWatcher对attribution LLM的选择不敏感,但如果attribution LLM和backbone LLM差异很大(如闭源LLM),可能需要额外的适配工作。
独立分析的弱点
第一个独立分析的弱点是计算开销较高,每次检测需要约8-10秒,这在实时交互场景中可能不可接受。改进方向可以是:进一步优化attribution算法,减少需要处理的注意力层和头;使用更小的attribution LLM模型;采用缓存机制,对于相似的上下文片段复用attribution结果;探索更细粒度的触发策略,只在真正高风险的操作上调用检测。第二个弱点是规则集的设计可能需要领域知识,自动生成的规则可能在某些场景下不够准确。改进方向可以是:开发半自动规则生成工具,让人类专家可以审阅和修正自动生成的规则;构建规则模板库,针对不同应用场景提供预定义规则;引入用户反馈机制,根据误报或漏报情况自动调整规则权重。第三个弱点是虽然实验表明AgentWatcher对attribution LLM不敏感,但在实际部署中可能面临模型版本差异的问题。改进方向可以是:研究跨版本LLM的attention模式迁移方法;开发模型无关的attribution方法;提供attribution模型校准工具,确保不同模型产生一致的归因结果。
未来方向
作者提出的未来工作方向包括:进一步优化计算效率,探索选择性调用检测器的策略,例如只在高风险操作上调用;研究更复杂的自适应攻击,评估AgentWatcher的鲁棒性;扩展到更多类型的agent,如多模态agent、协作agent等。基于本文成果可以延伸的研究方向包括:将AgentWatcher与其他防御方法结合,形成defense-in-depth架构,例如与prevention-based方法如SecAlign集成;研究规则集的持续学习和自适应更新机制,让系统能够根据新攻击模式自动调整规则;探索跨agent场景的检测,当多个agent协作时,如何检测注入是否通过agent之间的通信传播;研究更细粒度的action风险量化,不仅判断是否存在注入,还评估注入的潜在危害程度;开发用户友好的规则编辑界面,让非技术用户也能定制检测规则;研究attribution结果的可视化方法,帮助用户理解为什么某些上下文被判定为包含注入;探索在边缘设备上部署轻量级版本的AgentWatcher,适用于资源受限场景。
复现评估
论文声称代码已开源(https://github.com/wang-yanting/AgentWatcher),这为复现提供了基础。实验使用了公开的agent基准测试包括AgentDojo、InjecAgent、WASP和AgentDyn,以及LongBench中的长上下文数据集包括LCC、GovReport、PassageRetrieval、Qasper和HotpotQA。这些基准和数据集都是公开可获取的。实验设置详细,默认使用Llama-3.1-8B-Instruct作为attribution LLM,Qwen-3-4B-Instruct-2507作为monitor LLM,sink detection窗口大小$w_s=10$,窗口扩展大小$w_l=150$和$w_r=50$,窗口数量$K=3$。这些超参数设置明确。Ablation study系统地探索了不同attribution方法、attribution LLM和超参数的影响,结果透明。计算资源方面,主要使用A100 GPU,单次检测时间约8.2秒,这对于学术机构来说是可以获得的。GRPO微调使用了20000个训练样本,数据规模适中。整体来看,复现难度中等偏低,主要挑战可能在于:需要访问多个基准测试的评估环境;计算资源需求较高,尤其是attention extraction需要大模型推理;某些闭源LLM(如GPT-4o)的实验可能需要API调用,增加了成本和依赖性。
论文图表