AttnTrace: Трассировка контекста на основе внимания для языковых моделей с длинным контекстом

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

August 5, 2025

Авторы: Yanting Wang, Runpeng Geng, Ying Chen, Jinyuan Jia

cs.AI

Аннотация

Модели с длинным контекстом (LLMs), такие как Gemini-2.5-Pro и Claude-Sonnet-4, всё чаще используются для создания продвинутых ИИ-систем, включая конвейеры генерации с использованием извлечения данных (RAG) и автономных агентов. В таких системах LLM получает инструкцию вместе с контекстом — часто состоящим из текстов, извлечённых из базы знаний или памяти — и генерирует ответ, который контекстуально обоснован, следуя инструкции. Недавние исследования предложили решения для отслеживания подмножества текстов в контексте, которые в наибольшей степени влияют на ответ, сгенерированный LLM. Эти решения имеют множество практических применений, включая проведение анализа после атак и повышение интерпретируемости и доверия к выводам LLM. Несмотря на значительные усилия, современные решения, такие как TracLLM, часто приводят к высоким вычислительным затратам — например, TracLLM требуется сотни секунд для выполнения отслеживания для одной пары ответ-контекст. В данной работе мы предлагаем AttnTrace — новый метод отслеживания контекста, основанный на весах внимания, генерируемых LLM для промпта. Для эффективного использования весов внимания мы вводим две техники, предназначенные для повышения эффективности AttnTrace, и предоставляем теоретические обоснования для нашего выбора. Мы также проводим систематическую оценку AttnTrace. Результаты показывают, что AttnTrace более точный и эффективный, чем существующие современные методы отслеживания контекста. Мы также демонстрируем, что AttnTrace может улучшить современные методы обнаружения инъекций промптов в длинных контекстах с использованием парадигмы "атрибуция перед обнаружением". В качестве практического применения мы показываем, что AttnTrace может эффективно выявлять вставленные инструкции в статье, предназначенной для манипуляции отзывами, сгенерированными LLM. Код доступен по адресу 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.