AttnTrace: Rastreamento de Contexto Baseado em Atenção para LLMs de Contexto Longo

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

August 5, 2025

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

cs.AI

Resumo

Modelos de linguagem de longo contexto (LLMs), como o Gemini-2.5-Pro e o Claude-Sonnet-4, estão sendo cada vez mais utilizados para capacitar sistemas avançados de IA, incluindo pipelines de geração aumentada por recuperação (RAG) e agentes autônomos. Nesses sistemas, um LLM recebe uma instrução junto com um contexto—frequentemente composto por textos recuperados de uma base de conhecimento ou memória—e gera uma resposta contextualmente fundamentada ao seguir a instrução. Estudos recentes têm projetado soluções para rastrear um subconjunto de textos no contexto que mais contribuem para a resposta gerada pelo LLM. Essas soluções têm diversas aplicações no mundo real, incluindo a realização de análises forenses pós-ataque e a melhoria da interpretabilidade e confiabilidade das saídas dos LLMs. Embora esforços significativos tenham sido feitos, soluções de ponta, como o TracLLM, frequentemente resultam em um alto custo computacional—por exemplo, o TracLLM leva centenas de segundos para realizar o rastreamento de um único par resposta-contexto. Neste trabalho, propomos o AttnTrace, um novo método de rastreamento de contexto baseado nos pesos de atenção produzidos por um LLM para um prompt. Para utilizar efetivamente os pesos de atenção, introduzimos duas técnicas projetadas para aumentar a eficácia do AttnTrace, e fornecemos insights teóricos para nossa escolha de design. Também realizamos uma avaliação sistemática do AttnTrace. Os resultados demonstram que o AttnTrace é mais preciso e eficiente do que os métodos de rastreamento de contexto de ponta existentes. Também mostramos que o AttnTrace pode melhorar métodos de ponta na detecção de injeção de prompt em contextos longos por meio do paradigma de atribuição antes da detecção. Como uma aplicação do mundo real, demonstramos que o AttnTrace pode efetivamente identificar instruções injetadas em um artigo projetado para manipular revisões geradas por LLMs. O código está disponível em 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.