Source-linked AI summary
Attention-based Classification for Indirect Prompt Injection
Amelia Kawasaki, Stefan Lee
TL;DR
Indirect prompt injection embeds adversarial instructions in external data processed by LLM agents, while evidence for Attention Tracker's dataset-agnostic detection in realistic agentic settings is limited. This paper evaluates that method on AgentDojo and introduces AttnDA, which classifies per-head focus scores with Fisher Linear Discriminant Analysis. The evaluation finds that dataset-agnostic transfer does not fully hold, and complex covariance variants can degrade sharply in cross-domain transfer.
Problem
The paper evaluates whether Attention Tracker's claimed dataset-agnostic detection transfers to realistic agentic settings where injections arrive through external content.
Source: §1 · p. 2Method
The paper evaluates Attention Tracker on AgentDojo and introduces AttnDA, which replaces its scalar threshold with Fisher Linear Discriminant Analysis over head-level focus scores.
Source: §1 · p. 2, §3.1 · p. 4Results
AgentDojo evaluation finds that Attention Tracker's dataset-agnostic claim does not fully hold, while full and block-layer covariance variants show chance-level cross-domain transfer.
Source: §1 · p. 2Takeaways & Limitations
Training-set mismatch with downstream tasks can result in poor performance, especially in multi-turn agentic settings.
Source: §3.1 · p. 5The AgentDojo evaluation is limited to banking and Slack and contains 51 total samples, restricting conclusions about absolute performance there.
Source: §7.3 · p. 17
Abstract
from arXiv · showhide
Indirect prompt injection (IPI) attacks, in which adversarial instructions are embedded in external data processed by a language model, pose a significant threat to the integrity of deployed LLM agents. Attention Tracker (Hung et al., 2025) is a recently proposed white-box detection method that identifies prompt injection attacks by monitoring how a model's attention is distributed between user task tokens and data tokens, claiming dataset-agnostic detection via a scalar focus score threshold. In this paper, we provide the first evaluation of Attention Tracker in a realistic agentic IPI setting using the AgentDojo benchmark, finding that the dataset agnosticism claim does not hold. We find that in agentic environments, long structured tool outputs subvert the fundamental assumptions Attention Tracker makes about attention that the method depends on and produces near-chance detection on AgentDojo across most model families. We propose AttnDA, an extension replacing the scalar threshold with a Fisher Linear Discriminant Analysis classifier over the per-head focus score vector, and evaluate both methods in a cross-dataset setting spanning four benchmarks. AttnDA improves detection performance across agentic and tool-integrated benchmarks, though transfer to AgentDojo remains challenging regardless of classifier choice. A covariance ablation further shows that removing inter-head correlation structure substantially contributes to performance, with complex covariance assumptions causing significant degradation on cross-domain data. These results indicate that the attentional distraction signal underlying focus-score-based detection does not transfer reliably to agentic environments, and that learned linear classifiers over head-level attention features offer improved detection performance, provided training data from the target deployment domain is available.
1. Introduction
Indirect prompt injection hides adversarial instructions in external content that agents process, leaving practical defenses an open problem. This paper evaluates Attention Tracker in agentic settings, introduces AttnDA, and reports that dataset-agnostic transfer does not fully hold.
1. Introduction
- IPI attacks place malicious instructions in external content an agent reads, making them difficult to defend against in deployed systems.Existing defenses may require multiple forward passes or rely on brittle text-level classifiers.Source: §1 · pp. 1–2
- Attention Tracker detects attention shifts from legitimate instructions toward injected content using a scalar focus score, and claims its selected heads transfer across domains without retraining.Source: §1 · p. 2
- The first AgentDojo evaluation finds that Attention Tracker's dataset-agnostic claim does not fully hold in structured, multi-step agent environments.AgentDojo embeds prompt injection attacks in these environments.Source: §1 · p. 2
- AttnDA replaces Attention Tracker's scalar threshold with Fisher Linear Discriminant Analysis to account for varying head importance and find more discriminative boundaries.Source: §1 · p. 2
- Complex covariance assumptions significantly degrade detection, with cross-domain transfer collapsing to chance for full and block-layer variants.The degradation also occurs on in-domain data.Source: §1 · p. 2
- A matched benign dataset for LLMail-Inject attacks enables binary-classification evaluation.Source: §1 · p. 2
2. Related Works
Indirect prompt injection reaches agents through environmental content they retrieve and act on, sometimes without user oversight. Defenses include structural interventions and internal-state detection, with trade-offs in inference cost and robustness to subtle attacks.
2.1. Prompt Injection Attacks
- IPI places adversarial instructions in environmental content that agents retrieve and act on, creating an attack surface outside direct user prompts.Source: §2.1 · pp. 2–3
2.2. Defenses Against Prompt Injection
- Structural defenses can incur inference overhead, while surface-level classifiers may be brittle to semantically subtle attacks.Source: §2.2 · p. 3
- MELON re-executes each agent step with the user prompt masked and requires two forward passes; SmoothLLM uses 10–50 passes per query.Spotlighting marks untrusted data at prompt construction, while text classifiers can be brittle to novel or subtle attacks.Source: §2.2 · p. 3
- Internal-state detectors train classifiers on activations or representations, and prior work reports that benign/adversarial distinctions generalize to unseen attacks.This paper extends that direction to attention heads and structured environment data retrieved through tool calls.Source: §2.2 · pp. 3–4
3. Preliminary: Attention Tracker
Attention Tracker uses attention shifts between user instructions and task data for single-pass detection, claiming transfer without domain-specific retraining. Its head-selection and inference procedures rely on labeled data, and the paper notes that training-data mismatch can impair performance, especially in multi-turn agentic settings.
3. Preliminary: Attention Tracker
- Attention Tracker is a white-box, single-pass detector that uses attention drifting from the user task toward injected data and claims dataset-agnostic transfer without retraining.Source: §3 · p. 4
3.1. Training: Head Selection
- Head selection uses labeled benign and malicious instruction–data pairs, selecting heads by class-wise focus-score separation and assuming attacks lower attention to task instructions.The paper notes that despite the “training-free” label, a training set estimates head-selection parameters; mismatch with downstream tasks can impair performance, especially in multi-turn agentic settings.Source: §3.1 · pp. 4–5
3.2. Inference: Aggregate Focus Score
- At inference, Attention Tracker averages selected-head focus scores equally and flags a sample as an injection when its score falls below a calibrated threshold.Higher aggregate scores are taken to indicate benign samples.Source: §3.2 · p. 5
4. Improving Attention Tracker for Agentic Use Cases
AttnDA replaces Attention Tracker’s unweighted scalar average with a Fisher Linear Discriminant classifier over per-head focus scores, using diagonal covariance. The section describes its training and inference and contrasts how the two methods weight head-level evidence.
4. Improving Attention Tracker for Agentic Use Cases
- AttnDA replaces Attention Tracker’s unweighted average with a Fisher Linear Discriminant classifier over the per-head focus-score vector, assuming diagonal covariance.The diagonal-covariance choice uses the same per-head means and variances as Attention Tracker.Source: §4 · p. 5
4.1. Training: Fisher’s Linear Discriminant
- Training computes benign and malicious class means and diagonal covariance matrices from labeled examples’ per-head focus-score vectors.Diagonal covariance is chosen to reduce parameter count and overfitting; ε regularizes the shrunk covariance estimate.Source: §4.1 · pp. 5–6
4.2. Inference: Classifier Score
- At inference, the weighted focus-score vector yields a class-evidence score, and an attack is detected when that score exceeds a threshold calibrated on training data.Larger classifier scores indicate a higher likelihood that the example is malicious.Source: §4.2 · p. 6
4.3. Comparing AttnDA and Attention Tracker
- AttnDA weights all head focus scores by cross-class spread, whereas Attention Tracker uses binary weights and ignores negative malicious-indicating heads.Attention Tracker considers positive benign evidence only for heads above its importance threshold, without accounting for its relative strength.Source: §4.3 · p. 6
5. Experimental Settings
The experiments cover single-turn and agentic or tool-use injection datasets, with within-dataset and cross-dataset evaluation. AgentDojo detection is evaluated on the step where the agent processes the tool result containing injected content.
Datasets
- Evaluation spans single-turn deepset and agentic or tool-use AgentDojo, InjecAgent, and LLMail-Inject datasets, divided into training and evaluation sets for generalization to new instances.Cross-dataset generalization is also considered.Source: §5 · pp. 6–7
Datasets and Agentic Evaluation Procedure
- AgentDojo† uses the banking and slack suites, with 25 and 26 samples, respectively, and embeds injections in realistic structured environment data.In the agentic procedure, detection is assessed when the agent processes the tool result containing the injection; the user task is the instruction span and the tool result is the data span.Source: §5 · p. 7
Datasets
- InjecAgent covers 16 user task types and 1,008 samples, while LLMail-Inject has 272 attack emails paired with 272 benign rewrites that remove the injection.The “WebBrowerNavigateTo” task type is excluded because it returns no environmental data.Source: §5 · p. 7
Models and Metrics
- Experiments use five models and report held-out AUROC, where higher values indicate a greater likelihood of ranking a random attack above a random benign input.The models are Llama-3-8B-Instruct, Qwen2 1.5B, Gemma 2 9B, Phi-3 mini, and Mistral 7B.Source: §5 · p. 7
6. Results and Findings
Attention-based detection does not transfer reliably to agentic data: long structured inputs blur the focus-score distinction, and synthetic training performs poorly on AgentDojo. AttnDA improves several comparisons, but results depend on training domain and covariance structure, and AgentDojo transfer remains challenging.
6.1. RQ1: Examining Focus Score in the Agentic Setting
- AgentDojo’s long inputs have largely overlapping benign and attack focus scores, unlike deepset, undermining Attention Tracker’s assumed distraction signal.The AgentDojo inputs are longer, reflecting structured environment data in tool results.Source: §6.1 · p. 8
6.2. RQ2: Generalization from Synthetic Data
- Synthetic training yields 0.97 average AUROC on deepset but 0.57 average AUROC on AgentDojo.With synthetic training, AttnDA improves over AT by +0.05 average AUROC on AgentDojo and +0.11 on InjecAgent, while decreasing by -0.02 on deepset and -0.05 on LLMail-Inject.Source: Table 1 · p. 9, §6.2 · p. 9
6.3. RQ3: In-Domain Performance
- In-domain training raises AgentDojo AttnDA to 0.94 average AUROC versus AT’s 0.62; across datasets, AttnDA averages 0.92 versus AT’s 0.87.AT outperforms AttnDA on the single-turn deepset and LLMail-Inject tasks; both reach approximately 0.95 average AUROC on InjecAgent.Source: Table 2 · p. 10, §6.3 · p. 10
6.4. RQ4: Cross-Domain Performance
- Across 12 cross-domain pairs, AttnDA matches or beats AT in 10 and averages 0.79 AUROC versus 0.75, though AgentDojo transfer remains challenging.AttnDA underperforms on LLMail-Inject when trained on deepset or AgentDojo.Source: §6.4 · pp. 10–11, Table 3 · p. 11
6.5. Comparison to a Feature-level Classifier
- Feature-based logistic regression averages 0.58 AUROC across datasets, underperforming attention methods, though it is in-domain competitive except on InjecAgent.The authors speculate that higher feature dimensionality and instance-specific information may contribute to its poorer cross-dataset generalization.Source: Table 3 · p. 11, §6.5 · pp. 11–12
6.6. Covariance Structure Ablation of AttnDA
- For Qwen trained on AgentDojo, diagonal AttnDA transfers best (0.81 deepset, 0.86 LLMail-Inject, 0.84 InjecAgent), while full and block-layer covariance overfit.Full covariance produces the strongest in-domain performance across models, illustrating the tradeoff between in-domain results and transfer.Source: Table 4 · p. 12, §6.6 · pp. 12–13
6.7. Per-Head Weight Similarity
- AgentDojo’s learned weights are nearly orthogonal to those from other datasets (cosine 0.06–0.17); positive-weight heads also matter, but AT discards them.For Qwen2-1.5B, 41.7% of AgentDojo heads have positive weights, and AT’s in-domain AUROC is 0.44 versus AttnDA’s 0.88.Source: §6.7 · pp. 13–14
6.8. Quantifying Variability and Uncertainty
- For AgentDojo-trained AttnDA, bootstrapped average AUROC ranges from 0.80 to 0.94, with confidence intervals as low as 0.72 that overlap AT’s cross-domain averages.The bootstrap characterizes uncertainty from the small evaluation set.Source: §6.8 · p. 14, Table 6 · p. 15
7. Discussion
The claimed dataset agnosticism of Attention Tracker breaks down in agentic settings, where long structured inputs differ from single-turn data. AttnDA adds no measurable cost over Attention Tracker, but the conclusions remain limited by small and narrow evaluations.
7.1. Dataset Agnosticism Revisited
- Attention-based detectors trained on synthetic data do not reliably generalize to agentic inputs; deployment should include agentic samples in training.The authors identify input length and composition as the critical difference between single-turn and agentic environments.Source: §7.1 · p. 15
7.2. Computational Overhead
- AttnDA adds no measurable cost over Attention Tracker, with identical peak memory and runtime overheads within a few milliseconds of baseline prefill.At 2048 tokens, both attention-based methods use 9.6 GB, compared with 1.7 GB for baseline prefill; the authors attribute their shared memory cost to eager attention.Source: §7.2 · p. 15
- Figure 3 maps cross-dataset AUROC by model and method, with diagonal entries in-domain and off-diagonal entries showing zero-shot transfer.Synthetic is used for head selection only, without classifier training for LDA.Source: §7.2 · p. 16, Figure 3 · p. 16
7.3. Limitations
- AgentDojo† evaluation covers only banking and Slack, totaling 51 samples, limiting conclusions about its absolute performance.The authors note that results on a larger agentic dataset may differ.Source: §7.3 · p. 17
- The covariance ablation uses only Qwen trained on AgentDojo†, and weight-similarity analysis is also Qwen-only; bootstrap intervals do not capture uncertainty from fitting on a small training split.The covariance advantage and cross-dataset weight patterns may differ for other models or datasets.Source: §7.3 · p. 17
- The agentic benchmarks do not cover all deployment scenarios, and generalization to code-execution, retrieval-augmented, or multi-agent settings is unknown.LLMail-Inject’s benign class was synthetically generated, so the dataset should not be treated as a realistic sample of agentic data.Source: §7.3 · p. 17
Appendix A. Bootstrap Intervals and Procedure
The appendix estimates variability in AttnDA’s learned weights by stratified resampling of training data, then distinguishes that uncertainty from evaluation-set uncertainty in AUROC.
Appendix A. Bootstrap Intervals and Procedure
- For each dataset, the authors draw 200 stratified bootstrap resamples from the training split and refit AttnDA’s weight vector on each.Class proportions are preserved, and no test data enters the fits.Source: Appendix A · p. 20
- Within-dataset cosine similarities compare weights from independent resamples of the same dataset; cross-dataset similarities compare weights fit on different datasets.Both quantities are averaged over replicates.Source: Appendix A · p. 20
- The reported 5th–95th percentile intervals quantify training-resample uncertainty in learned weights, unlike Section 6.8 intervals, which quantify AUROC uncertainty by resampling evaluation data with fitted models fixed.Source: Appendix A · p. 20
Appendix B. Double BOS Token in Original Implementation
The appendix corrects tokenization behavior that misaligned Attention Tracker’s instruction and data spans, and reports all results using the corrected implementation.
Appendix B. Double BOS Token in Original Implementation
- The original code could prepend a second BOS token after the chat template, shifting token positions and misaligning instruction and data spans.Default tokenizer encoding also inflated length counts, further corrupting span boundaries used for head selection and focus scores.Source: Appendix B · p. 20
- The implementation disables added special tokens in all three tokenizer calls, leaving BOS handling to the chat template.Source: Appendix B · p. 20
- All reported results use the corrected implementation; Qwen2 avoids sequence-level misalignment because its template has no explicit BOS token, but its length inflation issue remains.Source: Appendix B · p. 21
Appendix C. LLMail-Inject Benign Sample Generation
Because LLMail-Inject contains attack emails without naturally occurring benign counterparts, the authors synthesize matched benign examples for binary classification.
Appendix C. LLMail-Inject Benign Sample Generation
- LLMail-Inject consists entirely of attack emails and has no naturally occurring benign counterparts.Source: Appendix C · p. 21
- For each attack email, an LLM removes the injected payload and returns the legitimate surrounding prose as a benign sample.Source: Appendix C · p. 21
- The resulting paired samples differ by the presence of an injected instruction and form a class-balanced evaluation set.Source: Appendix C · p. 21
C.1. Generation Procedure
The generation procedure uses GPT-4o-mini to remove injection-specific content while preserving legitimate email text. Model-control tokens are replaced with [TOKEN] and treated as injection.
C.1. Generation Procedure
- GPT-4o-mini removes injection-specific content while leaving legitimate email text unchanged.The sanitizer is instructed not to pad, expand, or invent content; fragments remain fragments and empty output is allowed when no legitimate text surrounds the attack.Source: §C.1 · p. 21
- Before sanitization, model-control tokens such as <|end|>, <|assistant|>, and [INST] are replaced with [TOKEN].Source: §C.1 · p. 21
- The prompt directs the model to remove each [TOKEN] marker and everything it introduces as injection.Source: §C.1 · p. 21
C.2. Generation Prompt
The generation prompt instructs a text sanitizer to remove injected attacks from synthetic email samples while preserving legitimate prose and its original structure.
C.2. Generation Prompt
- Treat each [TOKEN] marker and the content after it as injection until normal email prose resumes.The prompt says model-control tokens are represented by the [TOKEN] placeholder.Source: §C.2 · pp. 21–22
- Remove content intended to hijack an AI system, including injected commands, exfiltration requests, and tags wrapping injection instructions.The prompt also identifies fake conversation turns and API request payloads as injection content.Source: §C.2 · p. 22
- Preserve ordinary email prose, contextual email references, named details, and the original formatting without adding greetings or sign-offs.Fragments, unusual capitalization, and the absence of a greeting should remain unchanged.Source: §C.2 · p. 22
- Return only the sanitized text, even when removing an attack leaves just one or two legitimate words.Do not pad or invent text to fill the gap.Source: §C.2 · p. 22