RAISED: Self-Distillation for Robustness to Prompt Injection in LLM Agents
RAISED is a training framework for tool-using LLM agents that targets indirect prompt injection by having the model generate tool-use scenarios and then self-distill so its behavior on injected tool outputs matches its clean-context behavior. The neat idea is framing robustness as attack-invariant preservation of the model’s own benign trajectory distribution, especially for cases where legitimate tool output must be followed, though the summary leaves open how well it generalizes beyond the evaluated injection patterns, tools, and teacher-model behaviors.
Abstract
Tool-using language-model agents are vulnerable to indirect prompt injection because they must act on untrusted external content. Existing training-time defenses can reduce attack success rates, but often at the cost of general capabilities. We show that training-based defenses induce substantial drift in the model's output distribution, altering its behavior even in benign settings and providing a potential mechanism for utility degradation. We further identify a failure mode of these defenses: On benign tool-use tasks, the model refrains from a step needed to finish an authorized task, particularly when that step is indicated by a tool output. To address these limitations, we introduce RAISED (Robust Attack Invariance through Self-Distillation), a training framework that combines self-generation and self-distillation. The model first generates its own tool-use scenarios, with an emphasis on cases where task completion requires acting on legitimate guidance from tool outputs. Then, through self-distillation, the student is trained to match the teacher's clean-context behavior on both clean and injected variants of the same trajectory. RAISED substantially reduces the attack success rate of prompt injections in tool responses while, unlike prior training-based defenses, preserving utility on both agentic and general-purpose benchmarks.