Technical Summary: Defending against Adaptive Prompt Injection Attacks via Reasoning-enabled Task Alignment

1. Problem Statement

The paper addresses Indirect Prompt Injection (PI), a security vulnerability in LLM-based agents where adversaries embed malicious instructions in third-party data (e.g., web pages, API responses) retrieved by the agent. Unlike jailbreaking, where the user input is adversarial, PI involves a trusted user task ($t_u$) being hijacked by untrusted data ($o_t$) containing an attack task ($t_{adv}$).

The authors identify a critical gap in current defenses: while existing methods achieve low Attack Success Rates (ASR) on static benchmarks, their robustness collapses under adaptive attacks where the attacker optimizes the injection against the deployed defense. The paper traces this failure to two specific modes:

1. Observation-Level Failure: Existing defenses (prompting, filtering, secret-knowledge, and training-based) decide based on whether an observation resembles known injection patterns. They fail to verify if the agent's next action is actually aligned with the trusted user task. This allows attackers to rephrase injections (e.g., "Lexical Camouflage" or "Phantom Continuation") to bypass pattern matching while preserving the attack goal.
2. Adversarial Training Failure (Mode Collapse): Training-based defenses that use adversarial examples often suffer from "mode collapse." Even when using adaptive attackers (e.g., RL-Hammer), the generated examples tend to converge on a few successful reformulation strategies. Consequently, the defender learns to block specific templates rather than generalizing to the broader space of injection strategies, leaving it vulnerable to unseen reformulation tactics.

2. Methodology: RETA

The authors propose RETA (REasoning-enabled Task Alignment), a two-stage, training-based defense designed to enforce task alignment through explicit reasoning and diverse adversarial training.

Design Principles

1. Reason about Task Alignment, Not Patterns: Instead of detecting injection patterns, the defender must reason whether its planned tool calls ($c_{t+1}$) follow logically from the trusted user task ($t_u$).
2. Train for Strategy Coverage: Adversarial training must reward the discovery of diverse strategies for reformulating the attack, not just successful instances of a single strategy.
3. Adversarial Reinforcement Learning (RL): Defense is a trajectory-level problem. The defender must optimize over complete rollouts to ensure both security (rejecting $t_{adv}$) and utility (completing $t_u$).

Stage I: Adversarial Example Generation via Red-Teaming

This stage generates a diverse set of adversarial examples to train the defender.

• Attacker Policy ($\pi_A$): An attacker model generates injection strings ($d_{adv}$) and a reasoning trace explaining the reformulation strategy.
• Dictionary-Learning Diversity Reward: To prevent mode collapse, the system maintains a rolling strategy dictionary.
  • An auxiliary LLM extracts a "strategy record" (name and explanation) from each successful injection.
  • A dictionary learning algorithm (using sparse coding) learns primitives from these records.
  • A new injection receives a high diversity reward if its strategy record is poorly explained by the current dictionary (i.e., it represents a novel attack strategy).
  • The dictionary is refreshed periodically to include under-explained successful strategies.
• Outcome: A dataset ($D_{RT}$) of successful, strategy-diverse attacks is collected.

Stage II: Adversarial Reinforcement Learning

This stage trains the defender policy ($\pi_D$) using the dataset from Stage I.

• Initialization: The defender is first fine-tuned (SFT) on "boundary cases" where it learns to emit Task-Alignment Reasoning ($\rho_t$) before tool calls. This reasoning must satisfy two conditions:
  • C1: Explicitly identify content in the observation that attempts to redirect the agent.
  • C2: Justify that the planned tool calls are derivable solely from the trusted user task $t_u$.
• RL Optimization: The defender is optimized using Group reward-Decoupled Normalization Policy Optimization (GDPO). The reward function combines:
  • Security Reward ($r_{sec}$): 1 if the attack task is not completed.
  • Utility Reward ($r_{utl}$): 1 if the user task is completed.
  • Reasoning Reward ($r_{CoT}$): Scored by a judge on whether the generated reasoning satisfies C1 and C2.
• Goal: The policy learns to reject injected tasks while maintaining utility, guided by the explicit reasoning trace.

3. Key Contributions

1. Empirical Evidence of Static Benchmark Overestimation: The authors demonstrate that static benchmarks significantly overestimate defense robustness. Across four defense paradigms (Prompting, Filtering, Secret-Knowledge, Training) and two models, the ASR increased by an average of 15.8 percentage points when moving from static to adaptive attacks.
2. Identification of Structural Failure Modes: The paper isolates two root causes for this gap:
   • Observation-level decisions that ignore task alignment.
   • Adversarial training distributions that lack strategy coverage due to mode collapse.
3. Proposal of RETA: A novel defense framework that enforces task alignment via reasoning and utilizes dictionary-learning-based diversity rewards to ensure broad coverage of injection strategies during training.

4. Results

The authors evaluated RETA on the AgentDojo benchmark using Qwen3-4B-Instruct and Llama-3.1-8B-Instruct models against six black-box adaptive attack families (including TAP, Strategy, Genetic Search, and RL-Hammer).

• Attack Success Rate (ASR): RETA reduced the ASR for all evaluated adaptive attacks to below 10%.
  • Average ASR on Qwen3-4B: 2.92%.
  • Average ASR on Llama-3.1-8B: 3.75%.
• Utility Preservation:
  • Utility-under-Attack (UA): RETA preserved most utility under attack, significantly outperforming baselines that often dropped the user task to achieve security.
  • Benign Utility (BU): RETA maintained high performance on clean inputs, showing minimal degradation compared to undefended models.
• Ablation: Experiments confirmed that adding Task-Alignment Reasoning (CoT) to adversarial training significantly reduced held-out ASR (e.g., from 40.0% to 27.5% against Genetic Search) and recovered benign utility.

5. Significance and Claims

The paper claims that static robustness is insufficient against adaptive attackers in the context of indirect prompt injection. It argues that the field must shift from "pattern recognition" to "task alignment verification."

The significance of RETA lies in its ability to:

• Generalize beyond specific injection templates by training on strategy-diverse adversarial examples.
• Enforce task integrity by requiring the agent to justify its actions against the trusted user task, rather than simply filtering suspicious text.
• Achieve a superior safety-utility trade-off, successfully blocking adaptive attacks without sacrificing the agent's ability to complete legitimate tasks.

The authors modestly frame their work as a solution to the specific failure modes of current defenses, noting that while RETA performs well against the tested adaptive attacks, the problem of PI remains an active area of research where attackers may continue to evolve.