本文提出CachePrune方法，通过基于DPO损失的特征归因识别并修剪KV缓存中的关键神经元，防御间接提示注入攻击，同时保持模型响应质量。
Large Language Model, Safety, Robustness, Alignment, Human-AI Interaction
Rui Wang, Junda Wu, Yu Xia, Tong Yu, Ruiyi Zhang, Ryan Rossi, Lina Yao, Julian McAuley
Adobe Research, University of California San Diego, University of New South Wales, CSIRO’s Data61
Generated by grok-3-mini-latest
Background Problem
大型语言模型（LLMs）容易受到间接提示注入攻击的影响，这种攻击通过在提示上下文中注入指令来覆盖用户提供的指令，导致模型偏离预期行为。这源于LLMs在解析提示结构时的根本局限性，即无法有效区分数据和指令。现有防御方法要么涉及重新训练模型，计算开销巨大，要么通过修改提示格式或引入额外测试时工作流来优先处理用户请求，但这些方法往往防御效果有限、可能干扰模型输出质量或增加计算负担。本文的工作起点是解决这一数据与指令的误对齐问题，旨在开发一种高效的防御机制。
Method
- 核心思想： 通过特征归因（feature attribution）识别并修剪关键-值（KV）缓存中触发指令执行的神经元，强制LLM将输入上下文视为纯数据，而非指令。
- 如何实现： 基于直接偏好优化（DPO）目标的上界推导一个归因损失函数，仅需少量样本即可有效进行特征归因。具体步骤包括：
- 计算每个神经元的归因分数，使用公式 a_t^i = h_t^i imes rac{\partial \mathcal{L}^{attr}(\mathcal{Y}_x^c, \mathcal{Y}_x^p; \mathcal{X})}{\partial h_t^i} 来量化每个特征对损失的贡献。
- 聚合归因分数，取每个神经元在不同时间步的最大值：。
- 定义一个子集Φ，仅修剪那些在毒化响应中贡献更大且差异显著的神经元：。
- 应用掩码修剪：，其中τ是基于修剪比例p的阈值。
- 关键改进： 利用观察到的”触发效应”（triggering effect），仅对响应中的前几个token进行归因（默认k=1），提高效率，且不修改原始提示或引入额外测试时LLM调用。
Experiment
- 实验设置： 使用SQuAD、HotpotQA和WildChat数据集，模型包括Llama3-8B和Mistral-7B。指标涵盖攻击成功率（ASR, ↓）、F1分数（clean和attack, ↑）以及GPT-Score（↑）。实验设计合理，考虑了不同注入位置（上下文的开始、中间和结束），并与基线方法（如Vanilla、Delimiting、Datamarking、Encode_Base64）比较。消融实验探讨了修剪比例p、token数量k和掩码强度α的影响。
- 结果分析： CachePrune显著降低了ASR（例如，在SQuAD上，Llama3-8B的ASR从27.86%降至7.44%），同时保持了响应质量（F1 clean分数基本不变）。结果符合预期，证明了方法的有效性与高效性，尤其在不牺牲性能的情况下提高了防御能力。相比基线，CachePrune在ASR上表现出色，且修剪神经元后F1分数不下降，表明实验设置全面且合理。
- 为什么这样设计： 通过控制变量和多次试验（平均3次），确保结果鲁棒性；数据集选择覆盖不同任务（如问答和对话总结），验证泛化能力；指标设计直接评估攻击防御和输出质量，匹配研究目标。
Further Thoughts
这个方法强调了通过模型内部状态（如KV缓存）操纵来提升AI安全性的潜力，或许可以扩展到其他领域，如对抗样本防御或偏差缓解中；它突出了特征归因在解释性和鲁棒性构建中的作用，未来可与其他技术结合，例如与联邦学习整合以实现隐私保护的分布式防御，或应用于多模态模型中处理输入混淆的问题；此外，触发效应的发现可能启发更细粒度的模型行为分析，有助于开发更具适应性的AI系统。