PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Asif Pinjari, Mithun Paul Saint-Germain

arXiv:2609.10892v1cs.CRcs.AIcs.LG

TL;DR

Indirect prompt injection can compromise an agent through retrieved content, while existing systems do not jointly identify the attack entry, corrupted steps, and resisted injections. DriftNet reads logged trajectories with a dual-head Transformer and achieves strong detection and localization on a task-disjoint benchmark, though the synthetic single-generator corpus limits transfer claims.

  • Problem

    Existing detectors provide either whole-trace verdicts or single unsafe indices, leaving joint trajectory classification and injection-specific step labeling unresolved.

  • Method

    DriftNet uses a frozen sentence encoder, four identity-free world features, and a small dual-head Transformer to classify trajectories and label every step from logs alone.

  • Results

    Trajectory-level F1 reaches 0.983, exact injection-point recovery 98.7%, hijacked-span IoU 0.979, and DriftNet outperforms the retrained surface baseline on partial hijacks and delayed executions.

  • Takeaways & Limitations

    Joint sequence reading supports high-fidelity detection and localization, including distinguishing resisted injections, while remaining model-agnostic and log-only.

  • Takeaways & Limitations

    The corpus is synthetic and single-generator, so the results establish performance on that generator’s distribution rather than transfer to other generators or real agent traffic.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Indirect prompt injection can make an agent’s behavior drift after a poisoned observation, but existing detectors lack joint trajectory verdicts and step-level localization. DriftNet addresses this gap with a small log-only dual-head Transformer and a trust-oriented evaluation protocol.

  • Motivation: A poisoned observation can shift later tool calls toward an attacker’s goal, making injection detection and localization necessary for rollback, auditing, and source distrust.
  • Gap: Existing systems provide either whole-trace verdicts or single unsafe indices, which cannot represent delayed execution, recovery, or resisted injections.
  • Method: DriftNet embeds serialized steps with a frozen sentence encoder and four identity-free world features, then uses an under-two-million-parameter Transformer with two heads.
  • Evaluation: The evaluation uses task-disjoint training, a 20-configuration sweep within a 0.011 validation-F1 band, one held-out test evaluation, and an identical-split retrained baseline.
  • Contribution: DriftNet jointly classifies each logged trajectory and labels every step as benign, injection point, hijacked, or failed injection.Its output locates attack entry, corrupted actions, and resisted injections in one forward pass.
  • Results: The paper reports high-fidelity detection and localization on AgentDrift, strict localization metrics, pattern breakdowns, error analysis, and benchmark-artifact analysis.

II. RELATED WORK

Related work spans indirect-injection benchmarks, live-loop defenses, model-internal probes, and trajectory detectors. DriftNet differs by producing a trajectory verdict plus dense four-way step labels from logged behavior alone.

  • Indirect Prompt Injection and Agent Attack Benchmarks: Indirect prompt injection exploits retrieved content that agents may treat as instructions, and attack benchmarks measure substantial compromise across agent settings.
  • Live-Agent Defenses: Live-loop defenses can pause, replay, or veto actions, while DriftNet operates after the fact on recorded trajectories and world context without model internals.
  • Internal-State Methods: Model-internal approaches detect task drift or exposure using activations and attention, requiring access to the agent’s internal states.
  • Trajectory Detectors: Trajectory-level judges and guards commonly collapse safety into binary judgments or anchor supervision to the first unsafe action.
  • Positioning: Prior systems generally output a verdict or a single index or span, whereas DriftNet jointly labels entry point, corrupted span, and resisted injections.
  • Terminology: The paper distinguishes its AgentDrift benchmark from two unrelated works sharing the same name.

B. Threat Model

The threat model treats indirect prompt injection as a joint trajectory-classification and step-labeling problem from recorded agent behavior, without access to model internals. The labels distinguish benign steps, injection entry, hijacked actions, and resisted injections, enabling actionable localization.

  • Threat model: The attacker plants instructions in tool-returned observations, while the defender observes only the recorded trajectory and world context.This excludes access to agent internals and re-execution.
  • Threat model: DriftNet predicts both whether a trajectory is compromised and a four-way label for every step.The step labels are benign, injection_point, hijacked, and failed_injection.
  • Threat model: The joint output identifies where compromise entered, which actions were corrupted, and whether the agent resisted the injection.This makes the detector’s output operationally useful for rollback and source distrust.
  • Metrics: Trajectory detection uses precision, recall, F1, and category-specific flag rates, while step labeling uses per-class metrics with padding excluded.Localization additionally measures injection-point exact match and hijacked-span Jaccard overlap on attacked trajectories.

E. The Drift Hypothesis

DriftNet is motivated by a drift hypothesis: successful injections create an ordered benign prefix, poisoned observation, and attacker-serving suffix. The benchmark tests this against resisted attacks and hard negatives under a task-disjoint split.

  • E. The Drift Hypothesis: The drift hypothesis models successful attacks as a benign prefix, an injection point, and subsequent actions drifting toward the attacker’s goal.The model must therefore use the whole ordered sequence and distinguish inert poison from suspicious legitimate content.
  • Benchmark: AgentDrift contains 12,536 synthetic trajectories across five domains with 71,024 labeled steps and categories for benign, successful, failed, and hard-negative executions.Hard negatives contain legitimate content that superficially resembles an attack.
  • Benchmark: The task-disjoint split uses 9,081 training, 1,733 validation, and 1,722 test trajectories, preventing shared task templates from inflating generalization.Template holdout separates generalization from memorization.
  • Benchmark: The test part contains 9,796 labeled steps: 7,048 benign, 775 injection_point, 1,755 hijacked, and 218 failed_injection.These category counts provide denominators for later per-category rates.

V. DRIFTNET

DriftNet combines frozen textual step embeddings with identity-free world features and a compact bidirectional Transformer. Two heads jointly produce a trajectory verdict and per-step localization labels.

  • Architecture: DriftNet uses a frozen sentence encoder, four world-grounded identity-free features, and a compact Transformer with two prediction heads.The design preserves world context without encoding recipient identities.
  • Input representation: Each step is serialized with tool, thought, and argument fields before deterministic frozen sentence encoding.Embeddings are precomputed and cached per split.
  • World features: The world features capture internal versus external email and URL targets, so the model reads where actions go rather than only what text says.An example shows clean injection-step features followed by external-recipient features on the hijacked action.
  • Transformer trunk: The Transformer conditions each step on the full bidirectional trajectory, while masked pooling produces the trajectory verdict and step readouts produce localization scores.This links later hijacked actions to earlier poisoned observations.
  • Model size: The sweep-selected model has 1,845,637 trainable parameters, keeping the complete detector under two million parameters.The reference two-layer configuration has 1,318,533 parameters.

C. Right-Sizing

Because trajectories contain only 3–11 steps, DriftNet favors a shallow encoder while retaining self-attention for bidirectional contextual localization. The key requirement is sequence context rather than long-range memory.

  • Right-Sizing: Trajectories average 5.67 steps, so deep encoders risk overfitting while self-attention remains useful for bidirectional context.Localization depends on connecting later actions with earlier injection points and task context.

D. Training Objective

DriftNet jointly trains trajectory-level compromise detection and four-way step labeling with a class-weighted objective, masking padded steps and emphasizing rare failed injections. Training uses AdamW, scheduled learning rates, early stopping, and validation checkpoint selection.

  • D. Training Objective: DriftNet optimizes trajectory and step heads jointly with class-weighted binary and categorical cross-entropy losses.The trajectory loss weights the compromised class, while the step loss excludes padding and uses inverse-frequency class weights.
  • D. Training Objective: 1,845,637 trainable parameters are used in the selected three-layer Transformer configuration.The model freezes the sentence encoder and world-feature extractor, training only the projection, Transformer stack, and two heads.
  • D. Training Objective: The failed_injection class receives heavy weighting because it represents only 2.1% of steps and is safety-relevant.Validation uses the same weighted formulas as training, making the two loss curves directly comparable.
  • D. Training Objective: Training uses AdamW with triangular warmup and decay, gradient clipping at 1.0, and up to 40 epochs with patience-8 early stopping.The checkpoint with the best validation trajectory F1 is retained.

F. Deployment View

DriftNet is positioned as a post-hoc audit system that scores logged trajectories without model access or re-execution. Its evaluation protocol combines task-disjoint development, a 20-configuration sweep, one held-out test run, and a retrained surface baseline.

  • F. Deployment View: DriftNet audits completed or in-progress tool-call logs in one forward pass and returns per-step labels for operational triage.The labels identify attack entry, corrupted actions, and resisted injections.
  • F. Deployment View: The protocol touches the held-out test part exactly once, after development and robustness tuning.The selected configuration is evaluated without subsequent tuning, threshold adjustment, or model selection.
  • F. Deployment View: The sweep trains 20 configurations from scratch on the task-disjoint training split and selects using validation F1.The test part remains untouched during the sweep.
  • F. Deployment View: The surface baseline is retrained on the identical split for a like-for-like comparison.It uses six trajectory-level surface features with logistic regression and no sequence modeling or semantic embedding.

VII. RESULTS

On the held-out test set, DriftNet achieves high trajectory detection, step labeling, and localization performance, including exact injection-point recovery and strong hijacked-span overlap. Its selected configuration remains stable across the sweep, while resisted attacks are not flagged.

  • VII. RESULTS: Validation F1 spans only 0.011 across the 20-configuration sweep, indicating that the selected run is representative rather than exceptional.All configurations fall between 0.9827 and 0.9934 validation F1.
  • VII. RESULTS: 0.983 trajectory-level F1 is achieved on the single held-out test evaluation, with precision 0.982 and recall 0.985.The result comes from 763 true positives, 14 false positives, 12 false negatives, and 933 true negatives.
  • VII. RESULTS: 0.000 of 218 failed attacks are flagged, while hard negatives have a 0.029 flag rate.Benign trajectories have a 0.015 flag rate.
  • VII. RESULTS: Step-level F1 reaches 0.994 for benign, 0.981 for injection_point, 0.983 for hijacked, and 1.000 for failed_injection labels.The confusion is concentrated near hijack boundaries and among trajectories missed outright.
  • VII. RESULTS: 98.7% of attacked trajectories receive exact injection-point recovery, while hijacked-span IoU reaches 0.979.The span is recovered perfectly in 97.0% of attacked trajectories.
  • VII. RESULTS: Full hijacks reach 99.8% exact localization, partial hijacks 98.6%, and delayed executions 94.9%.The corresponding hijacked-span IoUs are 0.998, 0.964, and 0.936.

E. Comparison with the Surface Baseline

Against a retrained surface baseline on the identical task-disjoint split, DriftNet improves detection of partial and delayed hijacks while reducing false alarms. Residual misses concentrate in trajectories whose labeled injection observations contain no legible instruction.

  • E. Comparison with the Surface Baseline: 98.6% of partial hijacks and 93.2% of delayed executions are recovered by DriftNet, versus 11.1% and 17.1% for the surface baseline.DriftNet simultaneously lowers every false-alarm rate, showing the largest gains on behaviorally subtle attacks.
  • E. Comparison with the Surface Baseline: 11 of 12 false negatives come from partial or delayed executions, where one or two corrupted steps are embedded in otherwise normal trajectories.Only one full hijack is missed.
  • E. Comparison with the Surface Baseline: 11 of 12 misses receive predicted attack probability below 0.12, so threshold adjustment does not recover them.The misses are confident rather than marginal.
  • E. Comparison with the Surface Baseline: Nine of the 12 misses have no legible injected instruction in the observation labeled injection_point, consistent with residual corpus generation noise.In those cases, the attack’s textual trace appears only later or is absent from the labeled observation.

B. The Fourteen False Alarms

DriftNet’s residual errors are confident and concentrated in specific corpus and feature blind spots, while sequence-aware detection improves recall on stealthy attacks and localization remains strong.

  • B. The Fourteen False Alarms: 12 of 14 false alarms receive probability above 0.99, and Fig. 8 places errors at confident extremes rather than near the decision threshold.The false alarms include benign and hard-negative trajectories, especially in coding and web domains.
  • B. The Fourteen False Alarms: The world-feature channel confuses legitimate external browsing with attacks, while routine benign coding operations can resemble post-compromise activity.These patterns explain concentrated false alarms in web and coding cases.
  • A. Support for the Drift Hypothesis: 0.986 versus 0.111 partial-hijack recall and 0.932 versus 0.171 delayed-execution recall show DriftNet’s advantage over the surface baseline where sequence reading is required.The comparison uses the identical task-disjoint test split.
  • A. Support for the Drift Hypothesis: Most residual misses involve labeled injection observations with no legible instruction, indicating a benchmark labeling or signal boundary rather than an ordinary surface-statistics failure.The detector’s advantage concentrates on partial hijacks, delayed executions, resisted attacks, and hard negatives.
  • B. Localization Turns Flags into Actions: 98.7% exact injection-point recovery and 0.979 hijacked-span IoU make DriftNet’s localization dependable for auditing poisoned reads and corrupted actions.It also flags zero of 218 resisted attacks, separating injected text from actual compromise.
  • X. LIMITATIONS: The evaluation remains bounded by synthetic single-generator data, short 3–11-step trajectories, frozen embeddings, assumed world grounding, and non-adaptive attacks.The corpus also contains measured world-identity regularity that may affect trajectory-level results.

XI. FUTURE WORK

Future work targets leakage-controlled evaluation, stronger representations, broader data, and longer or streaming operation, while the current conclusion is supported by strong results under a bounded benchmark protocol.

  • XI. FUTURE WORK: Jointly fine-tuning the frozen step encoder could improve representation quality, but would sacrifice cached-embedding determinism and exact input reproducibility.This is identified as the most direct representation-quality extension.
  • XI. FUTURE WORK: Multi-generator corpora and logged real agent traffic are proposed to test generator robustness and practical transfer beyond the current benchmark.An LLM-judge comparison is also proposed to quantify trade-offs against zero-shot flexibility.
  • XI. FUTURE WORK: Scoring growing prefixes could turn completed-log analysis into streaming monitoring, but longer-horizon delayed execution remains untested.The open question is whether bidirectional attention over hundreds of steps preserves current delayed-execution performance.
  • XII. CONCLUSION: DriftNet achieves 0.983 trajectory-level F1, 98.7% exact injection-point recovery, and 0.979 hijacked-span IoU under a task-disjoint, single-test evaluation protocol.The protocol includes a 20-configuration sweep within a 0.011 validation-F1 band.
  • XII. CONCLUSION: 98.6% versus 11.1% partial-hijack recall and 93.2% versus 17.1% delayed-execution recall favor DriftNet over the retrained surface baseline while every false-alarm rate decreases.The pattern supports sequence reading rather than sharper surface statistics as the relevant advantage.
  • XII. CONCLUSION: The main empirical boundary is synthetic single-generator data with measured world-identity regularity, so the results do not establish transfer to real traffic or detector-aware attackers.The paper recommends anonymized, world-held-out, and generator-disjoint evaluation.
Loading 2609.10892v1…