Language model agents improve security by learning from attack failures
Self-Reflection Fine-Tuning: Enhancing Agent Security against Prompt Injection Attacks from Failure Experience
Machine LearningCryptography and Security
Summary
Language models that help with tasks can be tricked by bad instructions that make them do the wrong thing. The authors created a way for these models to learn from their own mistakes when attacked, by thinking carefully about what went wrong and how to stay on track. This helps the models recognize harmful instructions and keep doing what the user originally wanted. Their method worked well even when the bad instructions changed to try to fool the model again.
What this means in practice
- •For ai security teams: Develop safer AI assistants that resist being manipulated by harmful user instructions during operation.
- •For software developers: Build more robust AI tools that automatically detect and correct unsafe command injections during interactive tasks.
Authors
Zixuan Wang, Hao Li, Fengyu Gao, G. Edward Suh, Yi Zeng, Yevgeniy Vorobeychik, Ning Zhang, Chaowei Xiao
Abstract
Large language model (LLM) agents are increasingly deployed in tool-augmented environments, but their reliance on external inputs makes them highly vulnerable to prompt injection attacks that can hijack task objectives. Existing safety alignment methods rely on static expert trajectories or preference optimization, limiting their ability to generalize to adaptive attack patterns. In this work, we propose Self-Reflection Fine-Tuning (SRFT), a training framework that enables agents to improve robustness by learning from their own failure experiences under adversarial conditions. Instead of passively imitating expert behaviors, SRFT exposes the agent to compromised trajectories constructed via injected attacks, and leverages an expert model to generate structured self-reflection reasoning that contrasts unsafe and optimal actions. This reflective supervision teaches the agent to identify malicious instructions, reason about their consequences, and maintain alignment with the original user intent. We instantiate this framework in SR-Agent, built on Llama-3.1-8B-Instruct and Qwen3-8B, and evaluate it on both static and adaptive prompt injection benchmarks. Experimental results show that SRFT substantially reduces attack success rates while preserving task performance, and demonstrates strong generalization under adaptive attacks. These findings suggest that learning from failure via self-reflection is a promising direction for building robust and secure LLM agents. Our code is released at https://github.com/Eden-Wang1710/srft-repo.