AttnTrace: Op aandacht gebaseerde contextterugspoeling voor LLM's met lange context

AttnTrace: Attention-based Context Traceback for Long-Context LLMs

August 5, 2025

Auteurs: Yanting Wang, Runpeng Geng, Ying Chen, Jinyuan Jia

cs.AI

Samenvatting

Lang-context grote taalmodellen (LLM's), zoals Gemini-2.5-Pro en Claude-Sonnet-4, worden steeds vaker gebruikt om geavanceerde AI-systemen te versterken, waaronder retrieval-augmented generation (RAG) pijplijnen en autonome agents. In deze systemen ontvangt een LLM een instructie samen met een context—vaal bestaande uit teksten die zijn opgehaald uit een kennisdatabase of geheugen—en genereert een reactie die contextueel onderbouwd is door de instructie te volgen. Recente studies hebben oplossingen ontworpen om een subset van teksten in de context te traceren die het meest bijdraagt aan de reactie die door het LLM wordt gegenereerd. Deze oplossingen hebben tal van praktische toepassingen, waaronder het uitvoeren van forensische analyse na een aanval en het verbeteren van de interpreteerbaarheid en betrouwbaarheid van LLM-outputs. Hoewel er aanzienlijke inspanningen zijn geleverd, leiden state-of-the-art oplossingen zoals TracLLM vaak tot hoge rekenkosten, bijvoorbeeld duurt het TracLLM honderden seconden om een traceback uit te voeren voor een enkel reactie-context paar. In dit werk stellen we AttnTrace voor, een nieuwe context traceback-methode gebaseerd op de aandachtswaarden die door een LLM worden geproduceerd voor een prompt. Om aandachtswaarden effectief te benutten, introduceren we twee technieken die zijn ontworpen om de effectiviteit van AttnTrace te vergroten, en we bieden theoretische inzichten voor onze ontwerpkeuze. We voeren ook een systematische evaluatie uit voor AttnTrace. De resultaten tonen aan dat AttnTrace nauwkeuriger en efficiënter is dan bestaande state-of-the-art context traceback-methoden. We laten ook zien dat AttnTrace state-of-the-art methoden kan verbeteren bij het detecteren van prompt-injectie onder lange contexten via het attributie-voor-detectie paradigma. Als een praktische toepassing demonstreren we dat AttnTrace effectief geïnjecteerde instructies kan aanwijzen in een paper die is ontworpen om LLM-gegenereerde reviews te manipuleren. De code is te vinden op https://github.com/Wang-Yanting/AttnTrace.

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.