AttnTrace : Rétrocontexte basé sur l'attention pour les LLM à contexte étendu

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

August 5, 2025

papers.authors: Yanting Wang, Runpeng Geng, Ying Chen, Jinyuan Jia

cs.AI

papers.abstract

Les modèles de langage à grand contexte (LLMs), tels que Gemini-2.5-Pro et Claude-Sonnet-4, sont de plus en plus utilisés pour renforcer les systèmes d'IA avancés, y compris les pipelines de génération augmentée par récupération (RAG) et les agents autonomes. Dans ces systèmes, un LLM reçoit une instruction accompagnée d'un contexte—souvent composé de textes extraits d'une base de connaissances ou d'une mémoire—et génère une réponse contextualisée en suivant l'instruction. Des études récentes ont conçu des solutions pour retracer un sous-ensemble de textes dans le contexte qui contribuent le plus à la réponse générée par le LLM. Ces solutions ont de nombreuses applications pratiques, notamment l'analyse médico-légale post-attaque et l'amélioration de l'interprétabilité et de la fiabilité des sorties des LLM. Bien que des efforts significatifs aient été déployés, les solutions de pointe comme TracLLM entraînent souvent un coût de calcul élevé, par exemple, il faut à TracLLM des centaines de secondes pour effectuer un retraçage pour une seule paire réponse-contexte. Dans ce travail, nous proposons AttnTrace, une nouvelle méthode de retraçage de contexte basée sur les poids d'attention produits par un LLM pour une invite. Pour utiliser efficacement les poids d'attention, nous introduisons deux techniques conçues pour améliorer l'efficacité d'AttnTrace, et nous fournissons des insights théoriques pour notre choix de conception. Nous effectuons également une évaluation systématique d'AttnTrace. Les résultats démontrent qu'AttnTrace est plus précis et efficace que les méthodes de retraçage de contexte existantes de pointe. Nous montrons également qu'AttnTrace peut améliorer les méthodes de pointe dans la détection d'injection d'invites dans des contextes longs grâce au paradigme d'attribution avant détection. Comme application pratique, nous démontrons qu'AttnTrace peut identifier efficacement les instructions injectées dans un article conçu pour manipuler les critiques générées par un LLM. Le code est disponible à l'adresse 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.