Can self-distillation make LLM agents robust to prompt injection without losing utility?

The tradeoff RAISED is built to break

Tool-using agents read untrusted content from webpages, emails, documents and tool outputs, so instructions hidden in that content can hijack them away from the assigned task. Defenses face a two-sided constraint: the model must resist adversarial instructions in tool outputs while still acting on the task-relevant guidance those same outputs may legitimately contain.

The paper argues that existing training-time defenses buy robustness at a cost. Supervised fine-tuning on targets drawn from an external source rather than the base model's own output distribution can shift the model away from its original behavior, and preference-based defenses such as DPO optimize preferred responses relative to rejected ones on injected contexts without directly matching the model's predictive distribution on clean contexts. The paper reports that training-based defenses induce substantial drift in the model's output distribution, altering behavior even in benign settings, and treats that drift as a potential mechanism for utility degradation.

A failure mode where agents stop acting on tool output

Beyond distribution drift, the paper identifies a specific failure mode: on benign tool-use tasks, the model executes up to a point and then refrains from an authorized step needed to finish the task, particularly when that step is indicated by a tool output. The task is left unfinished even when the model was explicitly instructed to complete it end to end without further user involvement.

The paper's explanation is an asymmetry in the training objective. The objective consistently pushes the model away from responses that follow injected directions in tool outputs, while providing much less signal to follow legitimate actionable information from the same channel. The observed failures match a shortcut under that asymmetry: the model uses tool outputs as information but avoids acting on their guidance, even when doing so is necessary to complete the authorized task.

How RAISED trains invariance instead of a new policy

RAISED combines self-generation and self-distillation. A frozen copy of the model being defended serves as both data generator and teacher, while a trainable copy serves as the student. In the self-generation stage, the frozen model creates tool-use scenarios emphasizing cases where task completion requires acting on legitimate guidance from tool outputs, and the same model generates both the clean task-solving trajectories and the injections for their tool outputs.

In the self-distillation stage, the student minimizes a KL divergence between the teacher's next-token distribution and the student's distribution at generated token positions, on both clean and injected variants of the same trajectory. On clean contexts this keeps the student close to the frozen model's original distribution, limiting drift. On injected examples the same teacher targets train the student to recover the original prediction despite the added adversarial content. Because the legitimate tool content is left unchanged, the objective encourages resisting the injection while still acting on legitimate guidance needed to complete the user's task. For efficiency the method stores only the teacher's top-K token probabilities plus one residual bucket; the paper reports the residual bucket generally holds little probability mass, on the order of 10^-5.

Robustness gains that do not require giving up utility

The paper's central claim is that 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. The comparison is drawn against SecAlign++ at its 313-step checkpoint, the checkpoint the paper uses for the main comparison.

The paper reports that RAISED matches or exceeds SecAlign++'s robustness while better preserving benign utility, and induces far less distribution drift than SFT or DPO. The framing the authors draw from this is that when the desired behavior already exists in the clean model, defensive training should preserve it under attack rather than learn a new defensive policy.

Longer training lowers attack success and hides capability loss

The paper shows why checkpoint choice matters for how these defenses are judged. Extending SecAlign++ training continues to improve robustness: on Nemotron, AgentDojo ASR falls from 4.0±0.7% at 313 steps to 0.7±0.0% at 625 steps, and AgentDyn ASR falls from 7.8±0.3% to 3.0±0.2%. Judged on security benchmarks alone, further training looks beneficial.

The general-capability evaluation tells a different story. From 313 to 625 steps, MMLU drops from 48.3±0.7 to 42.8±0.3, GPQA from 43.8±1.7 to 39.4±1.1, Minerva Math from 75.0±0.1 to 66.1±0.9, and MBPP from 64.6±1.0 to 13.4±2.7. On benign agentic benchmarks the effect is milder but uneven: ToolTalk falls from 43.2±0.6 to 42.7±0.6, BFCL from 33.0±0.4 to 31.8±0.8, MMT-Bench from 43.4±0.3 to 41.9±0.3, and AsyncTool from 34.5±0.3 to 32.3±0.8, while WorkBench increases. The paper's reading is that prolonged training produces not a uniform collapse but a broader and increasingly uneven drift away from the original checkpoint's capabilities, and that selecting a checkpoint solely by attack success can conceal substantial degradation elsewhere. RAISED, by contrast, remains close to the Vanilla checkpoint across the same suite.

Ablations separate the data from the training objective

Because RAISED differs from SecAlign++ in both training objective and training data, the paper runs controlled comparisons on Gemma-4-12B-it. Holding the Glaive Function Calling v2 trajectories, injection placements and rendered inputs fixed, self-distillation preserves more benign agentic utility than SFT or DPO at similarly low attack success: on AgentDojo it reaches 80.1 utility versus 79.1 for DPO and 73.2 for SFT, and on AgentDyn 55.5 versus 48.6 and 17.5. SFT attains slightly lower ASR than self-distillation but at a much larger utility cost.

Holding the objective fixed, self-generated trajectories add a further utility gain. With DPO, replacing Glaive with self-generated data raises benign utility from 79.1 to 82.2 on AgentDojo and from 48.6 to 55.6 on AgentDyn while reducing ASR from 1.8 to 1.0 and from 2.3 to 1.6. Self-distillation shows the same trend, rising from 80.1 to 83.5 and from 55.5 to 60.0 while maintaining low ASR. Disabling the high-level structural patterns in generation while keeping the rest of the pipeline unchanged drops utility from 83.5 to 80.8±0.4 on AgentDojo and from 60.0 to 56.7±2.1 on AgentDyn, close to the 80.1 and 55.5 obtained with Glaive trajectories under self-distillation. The paper concludes the benefit comes largely from targeting scenarios that exercise the behaviors it aims to preserve, rather than from synthetic data generation alone.

Internal states shift less under injection

The paper checks whether the training objective shows up inside the model. On 88 clean/injected context pairs from AgentDojo, it measures the injection-induced change in hidden states at each of 48 decoder blocks, plus the distance between each trained model and the base model on clean contexts. The injection produces a similar early-block pattern across models followed by a peak around block 31, where the relative L2 change is 0.45 for the base model versus 0.23 for RAISED.

A behavioral check on 72 pairs accompanies this: RAISED starts the injected tool call in 12.5% of checked pairs, versus 62.5% for the base model and 13.9% for SecAlign++. RAISED also stays closer to the base model on clean contexts than SecAlign++ (0.12 versus 0.19 at block 47), and its KL to the base model's clean prediction is 0.011 nats versus 0.062 for SecAlign++ and 0.338 for the base model. The paper reads these measurements as consistent with the objective's intent: reduce sensitivity to injected content while limiting changes to clean-context behavior.

Where the robustness claim stops

The threat model is static indirect prompt injection, where adversarial content is chosen independently of the defended model. That matches the evaluation protocol of InjecAgent, AgentDojo and AgentDyn, and it is the setting for which the paper makes its robustness claims. The authors also tried an adaptive black-box attack using PISmith, but report that the optimization reward remained identically zero for most model-defense configurations, leaving the optimizer without signal. They explicitly state that a failed optimization run is not evidence of adaptive robustness, and restrict their empirical claims to the static threat model; robustness to model-adaptive prompt injection remains open.

Two further boundaries matter for interpreting the utility results. Clean-context self-distillation preserves the model's own behavior rather than improving its correctness, so if the clean model picks the wrong tool, follows a poor plan or hallucinates, the objective preserves that behavior under injection instead of correcting it; the defended model's quality is bounded by the clean-context teacher. And the evaluation covers three model families with different tokenizers, post-training recipes and baseline robustness, but not frontier models in the trillion-parameter regime, so whether the same security-utility behavior persists at that scale remains an empirical question. The paper also positions RAISED as a defense-in-depth mechanism that complements access control, sandboxing, capability restriction and least-privilege tool design rather than replacing them.

Compare checkpoints on capability, not attack success alone

The most decision-relevant result here is not a single attack success number but the pairing of two measurements. On Nemotron, extending SecAlign++ from 313 to 625 steps pushes AgentDojo ASR from 4.0±0.7% down to 0.7±0.0% while MBPP falls from 64.6±1.0 to 13.4±2.7 and Minerva Math from 75.0±0.1 to 66.1±0.9. A checkpoint chosen on attack success alone would look strictly better while carrying capability losses that only appear on the general-capability suite.

That is the comparison RAISED is meant to win: at similarly low attack success, it stays close to the Vanilla checkpoint on both agentic and general-purpose benchmarks, and the ablation attributes the gain to two separable choices, clean-context distribution matching rather than a DPO-style preference objective, and self-generated trajectories that specifically exercise acting on legitimate tool-output guidance. The Qwen results extend the pattern to a model that already starts from a low-ASR regime, where RAISED still reduces attack success on AgentDojo and AgentDyn without reducing benign tool-use capability. For readers weighing a training-time injection defense, the paper's own evidence says to evaluate the defended checkpoint on capability benchmarks alongside attack success, and to check whether the defense still lets the agent act on legitimate instructions arriving through the same tool-output channel it is being trained to distrust.

Sources used in this answer