AttnTrace: Rastreo de Contexto Basado en Atención para LLMs de Contexto Largo

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

August 5, 2025

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

cs.AI

Resumen

Los modelos de lenguaje de gran contexto (LLMs, por sus siglas en inglés), como Gemini-2.5-Pro y Claude-Sonnet-4, se utilizan cada vez más para potenciar sistemas de IA avanzados, incluyendo pipelines de generación aumentada por recuperación (RAG) y agentes autónomos. En estos sistemas, un LLM recibe una instrucción junto con un contexto—que a menudo consiste en textos recuperados de una base de conocimiento o memoria—y genera una respuesta que está contextualmente fundamentada siguiendo la instrucción. Estudios recientes han diseñado soluciones para rastrear un subconjunto de textos en el contexto que contribuyen más a la respuesta generada por el LLM. Estas soluciones tienen numerosas aplicaciones en el mundo real, incluyendo la realización de análisis forenses post-ataque y la mejora de la interpretabilidad y confiabilidad de las salidas de los LLM. Aunque se han realizado esfuerzos significativos, las soluciones de vanguardia, como TracLLM, a menudo conllevan un alto costo computacional; por ejemplo, TracLLM tarda cientos de segundos en realizar el rastreo para un solo par respuesta-contexto. En este trabajo, proponemos AttnTrace, un nuevo método de rastreo de contexto basado en los pesos de atención producidos por un LLM para un prompt. Para utilizar eficazmente los pesos de atención, introducimos dos técnicas diseñadas para mejorar la efectividad de AttnTrace y proporcionamos insights teóricos para nuestra elección de diseño. También realizamos una evaluación sistemática de AttnTrace. Los resultados demuestran que AttnTrace es más preciso y eficiente que los métodos de rastreo de contexto de vanguardia existentes. También mostramos que AttnTrace puede mejorar los métodos de vanguardia en la detección de inyección de prompts en contextos largos a través del paradigma de atribución antes de la detección. Como aplicación en el mundo real, demostramos que AttnTrace puede identificar eficazmente instrucciones inyectadas en un artículo diseñado para manipular reseñas generadas por LLM. El código está disponible en 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.