Prompt Injection Attack to Tool Selection in LLM Agents

原文：https://arxiv.org/pdf/2504.19793

代码：

摘要：

本文提出了一种名为 ToolHijacker 的新型提示注入攻击，专门针对大语言模型（LLM）智能体中的工具选择（tool selection）过程。LLM智能体选择工具通常分为“检索”和“选择”两个步骤。ToolHijacker攻击通过向工具库中注入一个精心制作的恶意工具文档，来操纵智能体的工具选择过程，迫使智能体在执行特定目标任务时，总是优先选择攻击者指定的恶意工具。

本文将恶意工具文档的生成过程形式化为一个优化问题，并提出了一种两阶段优化策略来解决该问题，分别针对工具选择的检索和选择两个阶段进行优化。这种攻击被设计用于“无盒”（no-box）场景

工具选择的原理：

工具选择分为“检索”和“选择”两个核心阶段：

2. 检索阶段 (Retrieval)
   该过程利用一个双编码器（Dual-Encoder）架构进行高效检索。首先，任务描述编码器 (f_q) 与工具文档编码器 (f_d) 分别将用户任务 q 和工具库 D 中的每个文档 d_j 映射为高维嵌入向量 f_q(q) 和 f_d(d_j)。随后，系统通过相似度函数（Sim(·,·)，如余弦相似度）计算任务向量与各工具向量间的相关性分数，并筛选出得分最高的 top-k 个工具文档，构成候选集 D_k。
   

3. 选择阶段 (Selection)
   在选择阶段，用户的任务描述 q 及检索到的候选集 D_k 将一同提交给大语言模型 E。该模型负责对候选集进行最终裁定，从中选出最适于执行任务 q 的工具 d*。
   

威胁模型

攻击者的能力：
无盒场景：目标任务描述、检索器、LLM 以及包括 top-k 设置在内的工具库都是不可访问的

攻击目标： 攻击者选定一个目标任务（可以通过多种语义表达），只要用户查询目标任务时，恶意工具会被选择来执行目标任务，从而绕过其他正常、合法的工具

攻击方法：

ToolHijacker攻击将攻击分解为两个阶段，以对应工具选择的“检索”和“选择”两个步骤。

并将恶意工具的描述 dt_des 被拆分为两个子序列：R 和 S (dt_des = R ⊕ S)，分别进行优化。

第一阶段：优化 R 以攻破检索 (Optimizing R for Retrieval)

• 目标：确保恶意工具在检索阶段能被高分召回，进入top-k候选列表。

利用无梯度方法得到初始化的R，再通过梯度下降优化

• 方法：最大化子序列 R 与一组“影子任务描述”的语义相似度。论文提出了两种实现方式：

  • 无梯度方法 (Gradient-Free)：利用一个LLM，根据多个影子任务描述自动生成一段功能性描述作为R。这段描述自然地包含了任务的核心语义，因此能有效提升检索排名。下面是模板prompt：
    

  • 基于梯度的方法 (Gradient-Based)：利用影子检索器的梯度信息，通过梯度下降（如HotFlip算法）来直接优化R中的每个词元（token），以最大化其与任务描述的相似度得分。
    

观察到的结论： 不同检索模型学到的语义模式有很大的重叠，使优化后的R能有效的迁移到目标检索器

第二阶段：优化 S 以攻破选择 (Optimizing S for Selection)

优化选择条件S：

• 目标：在恶意工具被成功检索到候选列表后，确保LLM最终会选择它。

• 方法：优化子序列 S，使其成为一个强大的、不可抗拒的指令，迫使LLM做出选择。同样提供了两种方式：

  • 无梯度方法 (Gradient-Free)：采用类似“攻击树”（Tree-of-Attack）的策略。由一个“攻击者LLM”生成S的多个变体，然后在“影子LLM”上进行测试和评估，保留效果最好的变体，并在此基础上进行多轮迭代优化。
    步骤一：初始化：初始化一个S
    步骤二：生成变体： 以当前的叶节点列表Leaf_curr的每个变体为基础再生成B个变体构建下一个叶节点列表Leaf_next，每个变体可表示为：

      （p_attack为systemp prompt，Feed表示前一次的迭代反馈）**步骤三：评估效果**：用影子LLM E' 对每个Leaf_next中的变体 产生一个响应，并通过正则化匹配来验证Leaf_next中的变体对影子任务描述的响应是否与恶意工具匹配**步骤四：剪枝与选择**：如果一个节点满足匹配成功的个数>m，则被认为是成功的，则优化结束；否则对Leaf_next进行剪枝，将节点个数限制在最大宽度W内，在进行迭代![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/f33942b7861e422db53e6c7730791c2e.png)
    

  • 基于梯度的方法 (Gradient-Based)：定义一个包含对齐损失（让输出与目标工具名对齐）、一致性损失（强化工具名本身）和困惑度损失（保证生成文本的流畅性）的综合损失函数。通过梯度下降优化S，以最大化“影子LLM”生成恶意工具名称的概率。
    将损失函数定义为
    
    其中：
    对齐损失L1
    
    一致性损失L2
    
    困惑度损失L3
    

采用JudgeDeceiver中的算法，优化过程包括两个关键部分：1) 位置自适应优化： 对于每个任务-检索对 {qi, D̃(i) ∪ {dt(S)}}，我们通过将 dt(S) 放置在影子检索工具文档集内的不同位置来优化 S；2) 步进式优化： 我们不一次性优化所有配对，而是逐步将任务-检索对纳入优化过程。这种渐进式方法有助于稳定优化过程。

实验：

Agent：

数据集：

• Metatool: 该基准测试专注于大语言模型（LLM）在工具使用方面的能力。它包含 21,127 个实例，涉及 199 个源自 OpenAI 插件的良性工具文档，并涵盖了四个子任务的查询。

• Toolbench: 该基准测试旨在通过 126,486 个指令微调样本来增强开源 LLM 的工具使用能力，利用了来自 RapidAPI 的 16,464 个工具文档。在移除重复工具和空描述后，其工具库包含 9,650 个良性工具文档。

指标：

• ACC (准确率): 在没有攻击的情况下，从工具库中为目标任务正确选择合适工具的可能性。通过对每个目标任务评估 100 个任务描述（即 m = 100）来计算。

• ASR (攻击成功率): 在注入恶意工具文档后，从工具库中选择该恶意工具的可能性。同样通过对每个目标任务评估 100 个任务描述（即 m = 100）来计算。

• HR (命中率): 对于一个目标任务，至少有一个正确工具出现在 top-k 检索结果中的比例。

• AHR (攻击命中率): 恶意工具文档 d_t 出现在 top-k 检索结果中的比例。

实验结果

​