# DriftNet: A Dual-Head Trajectory Transformer for Detecting and Localizing Prompt Injection in LLM Agents

arXiv: https://arxiv.org/abs/2609.10892 · HTML: https://mlsanity.com/paper/2609.10892/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Agent threats, guardrails and governance, one of 16 directions in LLM agents and tool use.
- Builds on 40 of its 45 references that are in the map.
- First posted on arXiv 9 Sep 2026.

## Placed in

- LLM agents and tool use / [Agent threats, guardrails and governance](https://mlsanity.com/agents/agent-threats-guardrails-and-governance/)
- LLM safety and security / [Prompt injection and agent security](https://mlsanity.com/safety/prompt-injection-and-agent-security/)

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Not what you've signed up for: Compromising Real-World LLM-Integrated Applications with Indirect Prompt Injection](https://mlsanity.com/paper/2302.12173/) (2023)
- [MPNet: Masked and Permuted Pre-training for Language Understanding](https://mlsanity.com/paper/2004.09297/) (2020)
- [Ignore Previous Prompt: Attack Techniques For Language Models](https://mlsanity.com/paper/2211.09527/) (2022)
- [InjecAgent: Benchmarking Indirect Prompt Injections in Tool-Integrated Large Language Model Agents](https://mlsanity.com/paper/2403.02691/) (2024)
- [Agent Security Bench (ASB): Formalizing and Benchmarking Attacks and Defenses in LLM-based Agents](https://mlsanity.com/paper/2410.02644/) (2024)
- [Formalizing and Benchmarking Prompt Injection Attacks and Defenses](https://mlsanity.com/paper/2310.12815/) (2023)
- [Defeating Prompt Injections by Design](https://mlsanity.com/paper/2503.18813/) (2025)
- [R-Judge: Benchmarking Safety Risk Awareness for LLM Agents](https://mlsanity.com/paper/2401.10019/) (2024)
- [Which Agent Causes Task Failures and When? On Automated Failure Attribution of LLM Multi-Agent Systems](https://mlsanity.com/paper/2505.00212/) (2025)
- [LlamaFirewall: An open source guardrail system for building secure AI agents](https://mlsanity.com/paper/2505.03574/) (2025)
- [GuardAgent: Safeguard LLM Agents by a Guard Agent via Knowledge-Enabled Reasoning](https://mlsanity.com/paper/2406.09187/) (2024)
- [AgenTracer: Who Is Inducing Failure in the LLM Agentic Systems?](https://mlsanity.com/paper/2509.03312/) (2025)
- [MELON: Provable Defense Against Indirect Prompt Injection Attacks in AI Agents](https://mlsanity.com/paper/2502.05174/) (2025)
- [Get my drift? Catching LLM Task Drift with Activation Deltas](https://mlsanity.com/paper/2406.00799/) (2024)
- [AgentDyn: Are Your Agent Security Defenses Deployable in Real-World Dynamic Environments?](https://mlsanity.com/paper/2602.03117/) (2026)
- [IPIGuard: A Novel Tool Dependency Graph-Based Defense Against Indirect Prompt Injection in LLM Agents](https://mlsanity.com/paper/2508.15310/) (2025)
- [PromptShield: Deployable Detection for Prompt Injection Attacks](https://mlsanity.com/paper/2501.15145/) (2025)
- [AttriGuard: Defeating Indirect Prompt Injection in LLM Agents via Causal Attribution of Tool Invocations](https://mlsanity.com/paper/2603.10749/) (2026)
- [ATBench: A Diverse and Realistic Agent Trajectory Benchmark for Safety Evaluation and Diagnosis](https://mlsanity.com/paper/2604.02022/) (2026)
- [The Landscape of Prompt Injection Threats in LLM Agents: From Taxonomy to Analysis](https://mlsanity.com/paper/2602.10453/) (2026)
- [Agent Drift: Quantifying Behavioral Degradation in Multi-Agent LLM Systems Over Extended Interactions](https://mlsanity.com/paper/2601.04170/) (2026)
- [TraceSafe: A Systematic Assessment of LLM Guardrails on Multi-Step Tool-Calling Trajectories](https://mlsanity.com/paper/2604.07223/) (2026)
- [TrajAD: Trajectory Anomaly Detection for Trustworthy LLM Agents](https://mlsanity.com/paper/2602.06443/) (2026)
- [TraceAegis: Securing LLM-Based Agents via Hierarchical and Behavioral Anomaly Detection](https://mlsanity.com/paper/2510.11203/) (2025)
- [Trajectory Guard -- A Lightweight, Sequence-Aware Model for Real-Time Anomaly Detection in Agentic AI](https://mlsanity.com/paper/2601.00516/) (2026)
- [WebSentinel: Detecting and Localizing Prompt Injection Attacks for Web Agents](https://mlsanity.com/paper/2602.03792/) (2026)
- [StepShield: When, Not Whether to Intervene on Rogue Agents](https://mlsanity.com/paper/2601.22136/) (2026)

## Abstract (arXiv)

When an indirect prompt injection succeeds against an LLM agent, the compromise is visible in the agent's own behavior: a benign prefix of tool calls, a poisoned observation, and a suffix of actions that serve the attacker. An operator needs three facts: where the attack entered, which steps it corrupted, and whether apparent poison was resisted. Existing systems return either a whole-trace verdict or a single unsafe index. We present DriftNet, a dual-head trajectory Transformer that reads a logged tool-call trajectory and answers all three questions in one forward pass: one head classifies the trajectory as compromised or not, and a second assigns every step one of four labels (benign, injection point, hijacked, failed injection). To our knowledge it is the first supervised detector to produce this joint output. A frozen sentence encoder and four identity-free world features embed each step; the trained trunk, under two million parameters and optimized with a class-weighted joint objective over both heads, needs no access to the agent's model. On the task-disjoint split of the AgentDrift benchmark (12,536 trajectories, 71,024 labeled steps), with a 20-configuration sweep bounding hyperparameter sensitivity to 0.011 F1 and the test part evaluated exactly once, DriftNet reaches trajectory-level F1 of 0.983, exact injection-point recovery on 98.7% of attacked trajectories, hijacked-span IoU of 0.979, zero flags on 218 resisted attacks, and 2.9% flags on hard negatives. A surface baseline retrained on the identical split recovers 11.1% of partial hijacks and 17.1% of delayed executions; DriftNet reaches 98.6% and 93.2% while lowering every false-alarm rate. Reading all 26 residual errors shows that most misses trace to trajectories whose labeled injection observation carries no legible instruction, and we report the benchmark's measured world-identity regularity alongside the results.