AttnTrace: Aufmerksamkeitsbasierte Kontextrückverfolgung für LLMs mit langem Kontext

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

Langkontext-große Sprachmodelle (LLMs), wie Gemini-2.5-Pro und Claude-Sonnet-4, werden zunehmend eingesetzt, um fortschrittliche KI-Systeme zu stärken, einschließlich Retrieval-Augmented-Generation (RAG)-Pipelines und autonomer Agenten. In diesen Systemen erhält ein LLM eine Anweisung zusammen mit einem Kontext – oft bestehend aus Texten, die aus einer Wissensdatenbank oder einem Speicher abgerufen wurden – und generiert eine Antwort, die kontextuell fundiert ist, indem die Anweisung befolgt wird. Aktuelle Studien haben Lösungen entwickelt, um auf eine Teilmenge der Texte im Kontext zurückzugreifen, die am meisten zur vom LLM generierten Antwort beitragen. Diese Lösungen haben zahlreiche praktische Anwendungen, einschließlich der Durchführung von forensischen Analysen nach Angriffen und der Verbesserung der Interpretierbarkeit und Vertrauenswürdigkeit von LLM-Ausgaben. Obwohl erhebliche Anstrengungen unternommen wurden, führen state-of-the-art Lösungen wie TracLLM oft zu hohen Berechnungskosten, z.B. benötigt TracLLM Hunderte von Sekunden, um eine Rückverfolgung für ein einzelnes Antwort-Kontext-Paar durchzuführen. In dieser Arbeit schlagen wir AttnTrace vor, eine neue Methode zur Kontextrückverfolgung, die auf den Aufmerksamkeitsgewichten basiert, die ein LLM für einen Prompt erzeugt. Um Aufmerksamkeitsgewichte effektiv zu nutzen, führen wir zwei Techniken ein, die darauf abzielen, die Wirksamkeit von AttnTrace zu verbessern, und wir liefern theoretische Einblicke für unsere Designentscheidung. Wir führen auch eine systematische Bewertung für AttnTrace durch. Die Ergebnisse zeigen, dass AttnTrace genauer und effizienter ist als bestehende state-of-the-art Methoden zur Kontextrückverfolgung. Wir zeigen auch, dass AttnTrace state-of-the-art Methoden bei der Erkennung von Prompt-Injection unter langen Kontexten durch das Attribution-before-Detection-Paradigma verbessern kann. Als praktische Anwendung demonstrieren wir, dass AttnTrace effektiv injizierte Anweisungen in einem Papier identifizieren kann, das darauf abzielt, LLM-generierte Bewertungen zu manipulieren. Der Code ist unter https://github.com/Wang-Yanting/AttnTrace verfügbar.

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.