通过注意力权重精准追踪大模型响应背后的上下文来源,提升安全检测效率。
AttnTrace: Contextual Attribution of Prompt Injection and Knowledge Corruption
- 基于大模型生成时的注意力权重进行上下文溯源
- 比现有方法快数十倍,单次追踪耗时不足1秒
- 适合需要快速验证模型输出可信性的安全研究人员
长上下文大语言模型(如Gemini-2.5-Pro和Claude-Sonnet-4)广泛应用于检索增强生成(RAG)与自主智能体系统中。这类系统接收指令与上下文(来自知识库或记忆),生成受上下文约束的回复。已有研究致力于追溯影响输出的关键上下文片段,用于攻击溯源与提升模型可解释性。但当前最优方法TracLLM需数百秒完成一次追踪,计算成本极高。本文提出AttnTrace,利用模型生成时的注意力权重实现高效上下文溯源。我们引入两项优化技术并提供理论支持,系统评估显示AttnTrace在准确率和效率上均优于现有方法。进一步实验表明,AttnTrace可借助‘先溯源再检测’范式,显著提升长上下文下提示注入攻击的识别能力。实际应用中,成功定位了设计用于操纵模型生成评论的注入指令。代码已开源。
原文摘要 · Abstract (English)
Long-context large language models (LLMs), such as Gemini-2.5-Pro and Claude-Sonnet-4, are increasingly used to empower advanced AI systems, including retrieval-augmented generation (RAG) pipelines and autonomous agents. In these systems, an LLM receives an instruction along with a context--often consisting of texts retrieved from a knowledge database or memory--and generates a response that is contextually grounded by following the instruction. Recent studies have designed solutions to trace back to a subset of texts in the context that contributes most to the response generated by the LLM. These solutions have numerous real-world applications, including performing post-attack forensic analysis and improving the interpretability and trustworthiness of LLM outputs. While significant efforts have been made, state-of-the-art solutions such as TracLLM often lead to a high computation cost, e.g., it takes TracLLM hundreds of seconds to perform traceback for a single response-context pair. In this work, we propose AttnTrace, a new context traceback method based on the attention weights produced by an LLM for a prompt. To effectively utilize attention weights, we introduce two techniques designed to enhance the effectiveness of AttnTrace, and we provide theoretical insights for our design choice. We also perform a systematic evaluation for AttnTrace. The results demonstrate that AttnTrace is more accurate and efficient than existing state-of-the-art context traceback methods. We also show that AttnTrace can improve state-of-the-art methods in detecting prompt injection under long contexts through the attribution-before-detection paradigm. As a real-world application, we demonstrate that AttnTrace can effectively pinpoint injected instructions in a paper designed to manipulate LLM-generated reviews. The code is at https://github.com/Wang-Yanting/AttnTrace.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。