PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses

1. Problem Statement

Prompt injection attacks pose a critical security risk to Large Language Model (LLM) applications, particularly autonomous agents. Attackers embed malicious instructions into the input context (e.g., web pages, documents) to manipulate the LLM's behavior. While numerous defenses (filter-based and prevention-based) have been proposed, their robustness against adaptive attacks remains insufficiently evaluated.

Existing red-teaming methods often fail to stress-test these defenses effectively:

• Static attacks rely on fixed templates and lack adaptability.
• Search-based attacks (e.g., TAP, PAIR) optimize per-instance but incur high computational costs and lack a transferable attack policy.
• Standard Reinforcement Learning (RL) approaches, such as applying Group Relative Policy Optimization (GRPO) directly, suffer from extreme reward sparsity. When facing strong defenses, the attack LLM generates mostly failed prompts (reward = 0). This leads to two failure modes:
  1. Entropy Collapse: The policy overfits to the few rare successful prompts, causing the policy entropy to drop to near zero. The model stops exploring new strategies and converges to a repetitive, ineffective mode.
  2. Gradient Dilution: Even if exploration continues, the gradient signal from the few successful rollouts is drowned out by the overwhelming majority of failures, preventing efficient learning.

2. Methodology: PISmith

The authors propose PISmith, a reinforcement learning-based red-teaming framework designed to overcome reward sparsity and entropy collapse. The core objective is to train an attack policy ($\pi_\phi$) to generate injected prompts ($p$) that maximize the success of an injected task ($g$) against a defended target LLM ($M_\theta$).

PISmith introduces two novel mechanisms to address the limitations of standard GRPO:

A. Adaptive Entropy Regularization

To prevent the policy from collapsing into a narrow mode before discovering effective attacks, PISmith replaces the standard KL-divergence regularization with a dynamic entropy bonus.

• Mechanism: An entropy cap ($H_{cap}$) is set. If the policy's entropy ($H(\pi_\phi)$) drops below this cap, an entropy bonus is applied.
• Adaptivity: The strength of the entropy regularization ($\beta$) is conditioned on the current mean reward ($\bar{r}$).
  • When $\bar{r}$ is low (sparse rewards), $\beta$ is maximized to force high exploration.
  • As $\bar{r}$ increases, $\beta$ decays, allowing the policy to shift toward exploitation.
• Goal: Sustain exploration long enough to discover rare successful attack patterns without causing the generated text to degenerate into incoherent noise.

B. Dynamic Advantage Weighting

To ensure the policy learns efficiently from the scarce successful rollouts, PISmith amplifies the gradient contribution of successful samples.

• Mechanism: The advantage ($A_i$) of a successful rollout ($r_i=1$) is multiplied by a dynamic factor $\gamma(\bar{r})$.
• Adaptivity:
  • When $\bar{r}$ is low (successes are rare), $\gamma(\bar{r}) > 1$, significantly amplifying the weight of the successful gradient.
  • As $\bar{r}$ improves, $\gamma(\bar{r})$ decays to 1, reverting to standard GRPO behavior.
• Goal: Ensure that the few successful signals drive the policy update rather than being diluted by the majority of failures.

Combined Objective: The final loss function combines the clipped surrogate loss (using the weighted advantage) and the adaptive entropy regularization term.

3. Key Contributions

1. PISmith Framework: A novel RL-based red-teaming method that introduces adaptive entropy regularization and dynamic advantage weighting to solve the reward sparsity problem in prompt injection red teaming.
2. Comprehensive Benchmarking: Evaluation across 13 diverse benchmarks covering Question Answering (QA), Retrieval-Augmented Generation (RAG), and Long-Context tasks.
3. Systematic Defense Evaluation: A rigorous comparison of 8 state-of-the-art defenses (both filter-based and prevention-based) against 7 baselines (static, search-based, and RL-based).
4. Agentic Setting Validation: Demonstration of PISmith's effectiveness in complex agent environments (InjecAgent and AgentDojo) against both open-source and closed-source models (including GPT-4o-mini and GPT-5-nano).

4. Experimental Results

A. Attack Performance (RQ1)

• Superior Success Rates: PISmith achieved an average ASR@10 of 1.0 and ASR@1 of 0.87 against Meta-SecAlign-8B (a state-of-the-art defense), significantly outperforming the strongest baseline, RL-Hammer (0.70/0.48).
• Generalization: Trained on only 100 samples from the Dolly Closed QA dataset, PISmith generalized effectively to 12 unseen benchmarks, including difficult long-context tasks where other methods failed completely (e.g., HotpotQA-Long: 0.99 ASR vs. 0.04 for RL-Hammer).
• Efficiency: PISmith required fewer total queries to the target model than search-based methods (32k vs. 82k–194k) while achieving higher success rates.

B. Utility-Robustness Trade-off (RQ2)

• Fundamental Tension: The study reveals a critical trade-off: No existing defense can simultaneously maintain high utility in benign settings and resist adaptive attacks.
  • Defenses that preserve utility well are easily broken by PISmith.
  • Defenses that are robust to PISmith often severely degrade the model's performance on benign tasks (high false positives).
• Implication: Current defenses are insufficient for secure real-world deployment; a new paradigm is needed to balance security and utility.

C. Agentic Settings (RQ3)

• PISmith achieved high attack success rates (ASR $\ge$ 0.95) on InjecAgent and AgentDojo.
• It successfully attacked GPT-5-nano and GPT-4o-mini, where static attacks and search-based methods failed or performed poorly.
• In AgentDojo (multi-step workflows), PISmith significantly outperformed template-based strategies, demonstrating its ability to navigate complex tool-use scenarios.

D. Ablation Studies

• Removing either the adaptive entropy regularization or the dynamic advantage weighting caused significant performance drops (ASR@1 fell from 0.87 to 0.57 or lower).
• This confirms that both mechanisms are essential: entropy ensures the discovery of rare successes, while advantage weighting ensures those successes are learned effectively.

5. Significance and Conclusion

• Security Implications: The paper demonstrates that current "state-of-the-art" prompt injection defenses are fragile against adaptive, RL-driven attackers. The false sense of security provided by static evaluations is dangerous.
• Research Direction: The findings highlight an urgent need for defense mechanisms that can maintain high utility while resisting adaptive attacks, a gap that current filter-based and prevention-based methods have not yet bridged.
• Responsible Disclosure: The authors emphasize that PISmith is a research tool designed to stress-test defenses. By identifying these vulnerabilities, the work aims to accelerate the development of more resilient LLM applications.

Code Availability: The implementation is open-sourced at https://github.com/albert-y1n/PISmith.