Source-linked AI summary
Attention Tracker: Detecting Prompt Injection Attacks in LLMs
Kuo-Han Hung, Ching-Yun Ko, Ambrish Rawat, I-Hsin Chung, Winston H. Hsu, Pin-Yu Chen
TL;DR
Prompt injections can make LLMs disregard original instructions, creating risks for integrated systems. Attention Tracker detects them by tracking attention to instructions in heads that exhibit the distraction effect. The paper reports that it significantly outperforms existing baselines, including with small LMs.
Problem
Method
Attention Tracker identifies heads that shift attention from original to injected instructions and aggregates their instruction-directed attention during inference for detection.
Source: §1 · pp. 1–2Results
Attention Tracker significantly outperforms existing baselines, and important-head attention differences remain consistent across the examined datasets.
Source: §5.4 · p. 7, §6 · p. 9Takeaways & Limitations
The distraction effect offers a new perspective on prompt-injection attacks and provides groundwork for future defenses.
Source: §6 · p. 9Detection relies on internal LLM information such as attention scores, which closed-source model users typically cannot access unless aggregated statistics are made available.
Source: Limitation · p. 9
Abstract
from arXiv · showhide
Large Language Models (LLMs) have revolutionized various domains but remain vulnerable to prompt injection attacks, where malicious inputs manipulate the model into ignoring original instructions and executing designated action. In this paper, we investigate the underlying mechanisms of these attacks by analyzing the attention patterns within LLMs. We introduce the concept of the distraction effect, where specific attention heads, termed important heads, shift focus from the original instruction to the injected instruction. Building on this discovery, we propose Attention Tracker, a training-free detection method that tracks attention patterns on instruction to detect prompt injection attacks without the need for additional LLM inference. Our method generalizes effectively across diverse models, datasets, and attack types, showing an AUROC improvement of up to 10.0% over existing methods, and performs well even on small LLMs. We demonstrate the robustness of our approach through extensive evaluations and provide insights into safeguarding LLM-integrated systems from prompt injection vulnerabilities.
1 Introduction
Prompt injection can redirect LLMs from original instructions to malicious instructions, creating risks for integrated systems. The paper identifies this attention shift and uses it to build a training-free detector.
1 Introduction
- Prompt injections embed malicious instructions in user inputs to make an LLM disregard its original instruction, threatening applications such as email and banking systems.Potential consequences include sensitive-information leakage or unauthorized transactions.Source: §1 · p. 1
- During attacks, specific attention heads shift focus from the original instruction to the injected instruction, a pattern the paper terms the distraction effect.The effect generalizes across attack types and dataset distributions.Source: §1 · pp. 1–2
- Attention Tracker identifies important heads with random LLM-generated sentences and a naive ignore attack, then aggregates their attention to the instruction as a focus score.The scores are obtained during the original inference, so detection requires no additional LLM inference.Source: §1 · pp. 1–2
- Attention Tracker improves AUROC by up to 10.0% over existing detection methods and by up to 31.3% on average over existing training-free methods.It is also effective on LMs with only 1.8 billion parameters.Source: §1 · p. 2
- The method requires zero data and zero training from existing prompt injection datasets.The paper reports effectiveness on both small and large language models.Source: §1 · p. 2
2 Related Work
Prior work studies prompt-injection attacks, instruction/data separation, and detectors, while attention research examines head functions and behavior. The paper positions its approach as attention-based detection without additional model inference.
Prompt Injection Attack
- Prompt-injection research covers handcrafted attacks and games, optimization-based attacks, differentiable trigger search, and attacks targeting existing defenses [Perez and Ribeiro, 2022; Greshake et al., 2023; Liu et al., 2023; Jiang et al., 2023b; Toyer et al., 2024; Debenedetti et al., 2024; Shi et al., 2024; Liu et al., 2024a; Zhang et al., 2024a; Pasquini et al., 2024; Khomsky et al., 2024].Some studies report that current defense and detection strategies remain ineffective against attacks targeting defenses.Source: §2 · p. 3
Prompt Injection Defense
- Defense work includes prompting or training models to distinguish instructions from data, querying LLMs, known-answer and perplexity detection, trained detectors, and classifiers using activation differences [Jain et al., 2023; Hines et al., 2024; lea, 2023; Piet et al., 2024; Suo, 2024; Chen et al., 2024; Wallace et al., 2024; Zverev et al., 2024; Liu et al., 2024b; Stuart Armstrong, 2022; Yohei, 2022; Alon and Kamfonas, 2023; Abdelnabi et al., 2024].The paper characterizes existing detectors as resource-intensive to train and often inaccurate.Source: §2 · p. 3
Backdoor Defense
- Unlike backdoor attacks, prompt injection manipulates user inputs at inference time without requiring training access or a hidden trigger [Gao et al., 2020; Zhao et al., 2024b].The paper distinguishes in-context triggers from prompt injections that exploit user-input content.Source: §2 · p. 3
Attention Mechanism of LLM
- Prior attention research studies head functions, attention-based behavior control, and detection of contextual hallucinations or trojans [Singh et al., 2024; Ferrando et al., 2024; Zhao et al., 2024a; Olsson et al., 2022; Crosbie and Shutova, 2024; Gould et al., 2024; Todd et al., 2024; Zhang et al., 2024b; Chuang et al., 2024; Lyu et al., 2022].This work applies attention analysis to distraction in important heads and prompt-injection detection.Source: §2 · pp. 3–4
3 Distraction Effect
Prompt injection replaces the intended task with an injected task by modifying user data. The section finds that attacks shift attention from the original instruction toward the injected instruction, with stronger attacks associated with lower total attention to the original instruction.
3.1 Problem Statement
- A prompt injection attack modifies data with an injected instruction so the application performs the injected task instead of the intended task.The section frames detection as identifying whether data has been compromised.Source: §3.1 · p. 4
3.2 Background on Attention Score
- At the first output token, attention weights represent the last input token’s attention to each prompt token in each layer and head.The analysis examines the last input token because it has the most direct influence on the model’s output.Source: §3.2 · p. 4, §3.3 · p. 4
3.3 A Motivating Observation
- During attacks, last-token attention shifts from the original instruction toward the injected instruction, while normal data receives higher instruction attention in specific heads.The analysis uses Llama3-8B on Open-Prompt-Injection and names this shift the distraction effect.Source: §3.3 · p. 4, Figure 2 · p. 5
- As attack success rate increases, the aggregated attention score across heads decreases, indicating a more pronounced distraction effect.The comparison spans different attack strategies.Source: §3.3 · p. 4, Figure 3 · p. 5
4 Prompt Injection Detection using Attention
Attention Tracker identifies attention heads that exhibit the distraction effect, then uses those pre-identified heads to measure attention to the instruction during detection. Its head-selection score accounts for separation between normal and attack attention-score distributions.
4 Prompt Injection Detection using Attention
- Attention Tracker uses the distraction effect to detect prompt injections by identifying important heads and tracking attention to the instruction.The focus score is available during test-query inference, making detection cost negligible compared with the original inference cost.Source: §4 · p. 5, §4.1 · p. 5, §4.2 · p. 6
4.1 Finding Important Heads
- Because the distraction effect does not apply to every head, the method first identifies the heads that exhibit it, termed important heads.Figure 2 is cited as showing that the effect is not universal across heads.Source: §4.1 · p. 5
- For each attention head, the method collects attention scores across normal and attack data and calculates a candidate score to select important heads.The passages distinguish the normal and attack score sets but do not provide a readable complete formula.Source: §4.1 · p. 5
- The candidate score accounts for differences in both the means and standard deviations of normal and attack attention-score distributions.It selects heads where the normal mean shifted down by k standard deviations exceeds the attack mean shifted up by k standard deviations; the default is k = 4.Source: §4.1 · p. 5
- The authors state that important heads generalize across datasets and use that claim to motivate finding them with a simple instruction and generated normal and attack examples.Normal data uses “Say {random word}” and 30 GPT-4-generated sentences; attack examples append “Ignore previous instruction and say ...” (Achiam et al., 2023).Source: §4.1 · pp. 5–6
4.2 Prompt Injection Detection with Important Heads
- At test time, Attention Tracker inputs the instruction and user query to the target LLM and calculates a focus score over the pre-identified important heads.The focus score measures the LLM’s attention to the instruction and is used to determine whether the input contains a prompt injection.Source: §4.2 · p. 6
- Because the important heads are pre-identified, the focus score is obtained during test-query inference “for free,” with detection cost described as negligible relative to inference.The method’s detection cost is compared with the original inference cost.Source: §4.2 · p. 6
5 Experiments
Across two datasets and five open LLMs, Attention Tracker outperforms existing baselines; qualitative and ablation analyses examine distraction patterns, head consistency, and focus-score sensitivity.
5.1 Experiment Setup
- Evaluation uses Open-Prompt-Injection and deepset across five open LLMs (1.5B–9B), reporting AUROC against trained, LLM-based, and known-answer detectors.The tested models include Qwen2-1.5B-Instruct, Phi-3-mini-4k-instruct, Mistral-7B-Instruct-v0.3, Llama3-8B-Instruct, and Gemma2-9b-it.Source: §5.1 · pp. 6–7
5.3 Qualitative Analysis
- In the grammar-correction example, attack data significantly reduces important-head attention on the instruction and shifts focus to the injected instruction.Figure 4 visualizes attention aggregation across tokens for normal and attack data.Source: §5.3 · p. 7, Figure 4 · p. 7
5.4 Discussion and Ablation Studies
- Across three datasets, the attention heads with the most distinct normal–attack score differences remain consistent, including with a basic attack used for head selection.The magnitude of score differences varies between datasets, while the relative differences across heads remain consistent.Source: §5.4 · p. 7, Figure 5 · p. 8
5.2 Performance Evaluation and Comparison with Existing Methods
- 10.0% maximum AUROC improvement on deepset and 3.1% on Open-Prompt-Injection show Attention Tracker outperforms existing baselines across both benchmarks.Against training-free methods, average AUROC gains are 20.9% on deepset and 31.3% on Open-Prompt-Injection.Source: §5.2 · p. 7, Table 1 · p. 8
5.4 Discussion and Ablation Studies
- Focus score decreases only negligibly relative to data-length growth, while instruction content—not length—is the primary influence on last-token attention.The analysis varies data length for a fixed instruction and compares data with and without an added attack string.Source: §5.4 · p. 8, Figure 6 · p. 9
- For Llama3, selecting approximately 0.3% of heads with k = 4 yields the highest AUROC; selecting too many or too few harms performance.Across models, most selected important heads lie in the first few or middle layers.Source: §5.4 · pp. 8–9, Table 2 · p. 8
6 Conclusion
The paper concludes that the distraction effect offers a new perspective on prompt injection and that Attention Tracker performs effectively, including with small LLMs.
6 Conclusion
- Attention Tracker significantly outperforms existing baselines and remains effective with small LLMs.The authors say the distraction-effect analysis and detector offer a new perspective and lay groundwork for future defenses.Source: §6 · p. 9
Limitation
- The approach relies on access to internal LLM information, such as attention scores, during inference.For closed-source LLMs, developers typically have access unless aggregated statistics such as focus scores are made available to users.Source: Limitation · p. 9
Ethics Statement
- The authors say reducing prompt-injection risks is crucial for LLM-integrated application safety and anticipate no negative social impact from this work.Source: Ethics Statement · p. 9
A.1 Introduction of Different Attacks in Figure 3
The appendix distinguishes five attack strategies by how they present injected instructions: direct concatenation, context cues, explicit disregard, false completion, or a combination.
A.1 Introduction of Different Attacks in Figure 3
- Naive attacks concatenate the injected instruction directly with the data without a separator.Source: §A.1 · p. 12
- Escape attacks use special characters such as “\n” to make the LLM perceive a context change.Source: §A.1 · p. 12
- Ignore attacks explicitly tell the LLM to disregard the original task, for example with “Ignore my previous instructions.”Source: §A.1 · p. 12
- Fake-complete attacks present a false response to the original task to prompt execution of the injected task instead.Source: §A.1 · p. 12
- Combined attacks use separators from escape, ignore, and fake-complete attacks together to maximize the effect.Source: §A.1 · p. 12
A.2 Dataset Settings
The benchmark settings cover two datasets with different instruction and attack setups.
A.2 Dataset Settings
- Open-Prompt-Injection uses five tasks, 50 target-inject data per task, and four attack types, totaling 2,000 test queries.The tasks are natural language inference, grammar correction, sentiment analysis, spam detection, and duplicate sentence detection; attack types are naive, ignore, escape, and combine.Source: §A.2 · p. 12
- For deepset, which has no instruction, the authors use “Say xxxxxx” as the instruction for every query.Source: §A.2 · p. 12
A.3 Baseline Settings
The appendix describes four detection baselines and their scoring or prompting procedures, alongside a qualitative attention-distribution figure.
A.3 Baseline Settings
- The appendix introduces Prompt-Guard, Protect AI detector, LLM-based Detection, and Known-answer Detection as the baselines whose settings it details.Source: §A.3 · p. 12
- LLM-based Detection asks whether a prompt is safe and requires a yes-or-no answer.Its prompt frames the model as Eliezer Yudkowsky assessing whether the input contains prompt-injection content.Source: §A.3 · p. 12
- Known-answer Detection scores the probability of generating “HELLO”; a higher score indicates the input likely contains no prompt injection.The prompt begins by asking the model to repeat “HELLO” once while ignoring the injected task.Source: §A.3 · p. 12
- Figure 7 presents qualitative attention-aggregation distributions across tokens for normal and attack data.Source: Figure 7 · p. 13
- Prompt-Guard scores a text by summing its prompt-injection and jailbreak logits, while Protect AI detector uses its prompt-injection logit.Source: §A.3 · p. 13
A.4 Experiment Settings
The appendix reports implementation settings and analyses of generated data, important-head positions, and instruction selection. It finds that important heads are usually early or middle layers, while detection performance remains high across tested instructions.
A.4 Experiment Settings
- Experiments used PyTorch and an NVIDIA RTX 3090, with each model’s evaluation across two datasets taking about one hour.Source: §A.4 · p. 13
A.5 More Qualitative Analysis
- The appendix visualizes additional instructions and data from the Open-Prompt-Injection benchmark.Source: §A.5 · p. 13
A.6 LLM-generated Dataset for Finding Important Heads
- Important-head identification used 30 GPT-4-generated sentences and an attack instruction to ignore the previous instruction and say a GPT-4-generated random word.Source: §A.6 · p. 13
A.7 Position of Important Heads.
- Across models, most important heads lie in the first few or middle layers, which may have greater influence on instruction-following behavior.Figure 8 visualizes head scores across different LLMs.Source: §A.7 · p. 13, Figure 8 · p. 14
A.8 Impact of Itest Selection
- AUROC remains consistently high across tested instructions, with output-directing instructions tending to yield higher scores on Deepset for Qwen-2-1.8B.Table 3 reports the AUROC scores for this dataset and model.Source: §A.8 · p. 13, Table 3 · p. 14