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

• URL: http://arxiv.org/abs/2508.03793v1
• Date: Tue, 05 Aug 2025 17:56:51 GMT
• Title: AttnTrace: Attention-based Context Traceback for Long-Context LLMs
• Authors: Yanting Wang, Runpeng Geng, Ying Chen, Jinyuan Jia,
• Abstract summary: We propose AttnTrace, a new context traceback method based on the attention weights produced by an LLM for a prompt.<n>The results demonstrate that AttnTrace is more accurate and efficient than existing state-of-the-art context traceback methods.
• Score: 30.472252134918815
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• TracLLM: A Generic Framework for Attributing Long Context LLMs [34.802736332993994]
  We develop TracLLM, the first generic context traceback framework tailored to long context LLMs.<n>Our framework can improve the effectiveness and efficiency of existing feature attribution methods.<n>Our evaluation results show TracLLM can effectively identify texts in a long context that lead to the output of an LLM.
  arXiv  Detail & Related papers  (2025-06-04T17:48:16Z)
• KnowTrace: Bootstrapping Iterative Retrieval-Augmented Generation with Structured Knowledge Tracing [64.38243807002878]
  We present KnowTrace, an elegant RAG framework to mitigate the context overload in large language models.<n>KnowTrace autonomously traces out desired knowledge triplets to organize a specific knowledge graph relevant to the input question.<n>It consistently surpasses existing methods across three multi-hop question answering benchmarks.
  arXiv  Detail & Related papers  (2025-05-26T17:22:20Z)
• Towards Effectively Leveraging Execution Traces for Program Repair with Code LLMs [13.708569727719434]
  Large Language Models (LLMs) show promising performance on various programming tasks.<n>We aim to remedy this potential blind spot by augmenting standard APR prompts with program execution traces.
  arXiv  Detail & Related papers  (2025-05-07T14:12:41Z)
• Large Language Models as Realistic Microservice Trace Generators [54.85489678342595]
  Workload traces are essential to understand complex computer systems' behavior and manage processing and memory resources.<n>This paper proposes a first-of-a-kind approach that relies on training a large language model to generate synthetic workload traces.<n>Our model adapts to downstream trace-related tasks, such as predicting key trace features and infilling missing data.
  arXiv  Detail & Related papers  (2024-12-16T12:48:04Z)
• TraceMesh: Scalable and Streaming Sampling for Distributed Traces [51.08892669409318]
  TraceMesh is a scalable and streaming sampler for distributed traces. It accommodates previously unseen trace features in a unified and streamlined way. TraceMesh outperforms state-of-the-art methods by a significant margin in both sampling accuracy and efficiency.
  arXiv  Detail & Related papers  (2024-06-11T06:13:58Z)
• Get my drift? Catching LLM Task Drift with Activation Deltas [55.75645403965326]
  Task drift allows attackers to exfiltrate data or influence the LLM's output for other users.<n>We show that a simple linear classifier can detect drift with near-perfect ROC AUC on an out-of-distribution test set.<n>We observe that this approach generalizes surprisingly well to unseen task domains, such as prompt injections, jailbreaks, and malicious instructions.
  arXiv  Detail & Related papers  (2024-06-02T16:53:21Z)
• Unsupervised Information Refinement Training of Large Language Models for Retrieval-Augmented Generation [128.01050030936028]
  We propose an information refinement training method named InFO-RAG. InFO-RAG is low-cost and general across various tasks. It improves the performance of LLaMA2 by an average of 9.39% relative points.
  arXiv  Detail & Related papers  (2024-02-28T08:24:38Z)
• Harnessing Explanations: LLM-to-LM Interpreter for Enhanced Text-Attributed Graph Representation Learning [51.90524745663737]
  A key innovation is our use of explanations as features, which can be used to boost GNN performance on downstream tasks. Our method achieves state-of-the-art results on well-established TAG datasets. Our method significantly speeds up training, achieving a 2.88 times improvement over the closest baseline on ogbn-arxiv.
  arXiv  Detail & Related papers  (2023-05-31T03:18:03Z)
• Origin Tracing and Detecting of LLMs [46.02811367717774]
  We propose an effective method to trace and detect AI-generated contexts. Our proposed method works under both white-box and black-box settings. We construct extensive experiments to examine whether we can trace the origins of given texts.
  arXiv  Detail & Related papers  (2023-04-27T10:05:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.