Skip to content

Back to the catalog

RAISED: Self-Distillation for Robustness to Prompt Injection in LLM Agents
arxiv.orgpaper

RAISED: Self-Distillation for Robustness to Prompt Injection in LLM Agents

This arXiv paper introduces RAISED (Robust Attack Invariance through Self-Distillation), a training framework that defends tool-using LLM agents against indirect prompt injection without the utility degradation seen in prior training-based defenses. The authors show existing defenses cause distributional drift that makes models wrongly refuse legitimate tool-output-guided steps on benign tasks, then fix this by having a model generate its own tool-use scenarios and self-distill to match clean-context behavior on both clean and injected trajectory variants. Results show reduced attack success rate on injected tool responses while preserving performance on agentic and general-purpose benchmarks.

TL;DR A self-distillation training method (RAISED) cuts prompt-injection attack success in tool-using LLM agents without the usual utility hit of prior defenses.
Open original ↗
Document

What it is

RAISED (Robust Attack Invariance through SElf-Distillation) is a training framework from LIX/École polytechnique and Google DeepMind for defending tool-using LLM agents against indirect prompt injection, evaluated on Gemma, Nemotron, and a third open-weight model against undefended agents, prompting/filtering defenses, SFT, DPO-based SecAlign++, on benchmarks including AgentDojo, InjecAgent, and AgentDyn.

How it works

A frozen copy of the model acts as both data generator and teacher; a trainable copy is the student. Stage 1 (self-generation): the frozen model writes its own tool-use scenarios — emphasizing cases where finishing the task requires acting on legitimate tool-output guidance — solves them to get clean trajectories, then role-plays as attacker to insert injections into the tool outputs. Stage 2 (self-distillation): the student is trained via KL divergence to match the teacher's next-token distribution on the clean context, applied to both the clean trajectory and its injected variant. This makes the student insensitive to the injected span while still being pulled toward acting on legitimate guidance in the same tool output — unlike SFT/DPO, which only penalize following injected instructions without reinforcing the use of legitimate tool-output guidance.

Preconditions / Affected

Applies to agentic LLM setups where the model calls tools/APIs/MCP and tool outputs can carry attacker-controlled text (static attacks, i.e., injection text not adapted to the defended model). Training used 5K injection-augmented trajectories, rank-16 LoRA, 313 steps per the paper's controlled comparison.

What to do

Quantitative results reported: on AgentDojo with Gemma, RAISED cuts ASR from 39.5% to 1.0% and raises utility-under-attack from 64.0% to 80.5%, at a cost of only 1 point of benign utility — matching SecAlign++'s robustness without its 7-point utility drop. In a controlled drift comparison (same data/steps, objective varied), self-distillation gave forward/reverse KL of 0.0010 vs. base model, 44–55x lower than SFT and 29–30x lower than DPO, with 0.993 top-1 token agreement vs. 0.953 (SFT) / 0.968 (DPO). Teams building/evaluating agent defenses can use these numbers as a baseline to beat and the KL/top-1-agreement protocol as a reusable benign-drift diagnostic for their own defenses.

Caveats

Results are demonstrated on open-weight models (Gemma, Nemotron, one other) under static attacks only — adaptive attacks against the defended model itself aren't addressed. The drift analysis explicitly states it does not establish that output-distribution drift causes downstream utility loss, only that it correlates with it and that self-distillation avoids it.

Technologies named here

Techniques named here

Related in AI Security