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

2026-06-13

AI Summary

Existing defense methods exhibit limited generalization against adaptive prompt injection attacks due to their reliance on fixed attack patterns and insufficient diversity in training data. To address this, this work proposes RETA, the first defense framework that integrates task-oriented chain-of-thought reasoning into the mitigation mechanism. RETA verifies the alignment between each tool-calling action and the user’s intent at every reasoning step, while enhancing adversarial training coverage through red-teaming with diversity-promoting rewards derived from dictionary learning. By synergistically combining chain-of-thought reasoning, multi-objective reinforcement learning, and diverse adversarial example generation, RETA achieves average attack success rates (ASR) of only 2.92% and 3.75% under six black-box adaptive attacks—both well below 10%—while maintaining high task utility on both clean and adversarially perturbed inputs.

Abstract

Indirect prompt injection attacks hijack LLM-based agents by embedding malicious instructions in third-party data that the agent retrieves during task execution. Existing defenses report near-zero attack success rate on static benchmarks, yet recent adaptive evaluations show that these results collapse once the attacker is allowed to optimize against the deployed defense. In this work, we trace this collapse to two failure modes. First, existing defense methods are confined to recognizing specific attack patterns, rather than assessing whether the intent of every embedded instruction is relevant to the user task. Second, training-based defenses, which otherwise offer the strongest safety-utility trade-off, assemble their adversarial examples from a handful of hand-crafted templates, and the resulting defender fails to generalize outside that narrow strategy distribution. To address these gaps, we propose RETA, a training-based method that grounds defense decisions on the user tasks rather than attacker-controlled data. At each tool-output step, the defender undertakes chain-of-thought reasoning verifying that its actions are consistent with the user task. Leveraging red-teaming, a simulated attacker synthesizes adversarial training data and receives a dictionary-learning diversity reward, achieving broad coverage of injection-reformulation strategies. Together, these allow the defender to be optimized via multi-objective reinforcement learning and achieve better safety-utility trade-off. Across six black-box adaptive attacks, RETA keeps every per-attack ASR below 10%, with average ASR of 2.92% and 3.75% on the two target models, while preserving most utility under attack and on clean inputs.

Problem

• Adaptive Prompt Injection
• LLM Security
• Task Alignment
• Adversarial Robustness
• Indirect Prompt Injection

Innovation

• Adaptive Prompt Injection
• Task Alignment
• Chain-of-Thought Reasoning
• Red-Teaming
• Multi-Objective Reinforcement Learning