PaperlayerView exact revision on arXiv

Source-linked AI summary

Identify, Classify, Act: Per-Instruction Controlled Compliance for LLMs over External Data

Zesheng Shi, Xue Li, Chenrui Dai, Jing Li

arXiv:2610.22874v1cs.CLcs.CR

TL;DR

External data can interleave useful instructions with harmful prompt injections, while existing defenses often suppress or trust external directives wholesale. EDIC trains a judge-guided response model to classify and act on each instruction, improving safety and utility across two backbones and five benchmarks. The authors identify limited scene coverage and fixed injection triggers among the method’s limitations.

  • Problem

    External data can mix useful and harmful instructions, while existing defenses largely treat embedded directives uniformly rather than distinguishing them per instruction.

    Source: §1 · p. 1

  • Method

    EDIC synthesizes mixed-instruction contexts and trains a generative judge to assign per-instruction compliance rewards for a response model optimized by supervised fine-tuning and judge-guided reinforcement learning.

    Source: §1 · pp. 1–2

  • Results

    Across two backbones and five external benchmarks, EDIC improves both safety and utility over the strongest prior defenses, exceeding ReasAlign by 7.7pp on Qwen3-4B harmful-rejection.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    The results support per-instruction identification, classification, and action as a way to improve safety and utility together on the evaluated external-data benchmarks.

    Source: §3 · p. 3, §7 · p. 8

    Training covers four scenes and uses a fixed catalog of benchmark injection triggers, leaving generalization to other settings and adaptive attackers unvalidated.

    Source: Limitations · p. 9

Abstract

from arXiv · show

LLMs over external data face a per-instruction tension: a single document or tool return often interleaves legitimate directives with adversarial prompt injections, and the model must follow the former while rejecting the latter. Existing defenses suppress or trust all external instructions, collapsing the safety-utility frontier and forcing a choice between over-refusal and over-compliance. We argue the missing capability is \textbf{per-instruction credit attribution} and formalize it as External-Data Instruction Classification, an identify-classify-act decision over each embedded directive. We propose \textbf{EDIC}, a decomposed reward framework for per-instruction compliance. A scene-grounded pipeline injects mixed beneficial and harmful instructions with counterfactual filtering, while a generative judge is initialized by SFT and refined with GRPO to attribute compliance at the instruction level. The judge then produces a reward that separately encourages beneficial compliance and penalizes harmful compliance for RL training. Across two backbones and five benchmarks, EDIC improves both safety and utility over the strongest prior defense, lifting harmful-rejection by up to 7.7\% and Arena-Hard-Auto utility by 12.5\% on Qwen3-4B. The code repository is hosted at https://github.com/yingyingxia666/EDIC.

1 Introduction

External data can contain both useful and harmful instructions, while existing defenses often reject or trust all embedded directives. EDIC frames per-instruction classification as the alternative and uses judge-guided training to improve safety and utility.

1 Introduction

  • External data mixes task-relevant beneficial instructions with harmful injections, but defenses that suppress external directives can cause over-refusal and degrade legitimate task performance.The alternative of fully trusting embedded instructions can allow attacks to succeed.Source: §1 · p. 1, Figure 1 · p. 2
  • A generative judge assigns per-instruction compliance, and a response model is trained with supervised fine-tuning followed by judge-guided reinforcement learning.The judge separately rewards following beneficial instructions and penalizes following harmful ones.Source: §1 · pp. 1–2
  • On Qwen3-4B, EDIC achieves 91.2% harmful-rejection and 89.5% benign-follow, outperforming the strongest reasoning-based baseline ReasAlign by 7.7pp on safety and 12.5pp on Arena-Hard-Auto utility.The same trend holds on Llama-3.1-8B, with 92.8% harmful-rejection and 91.7% benign-follow.Source: §1 · p. 2
  • EDIC formalizes the task as identifying, classifying, and acting on each embedded instruction in mixed external data.Its data-synthesis pipeline builds realistic mixed-instruction contexts with counterfactual quality checks.Source: §1 · pp. 1–2

2 Related Work

Prior prompt-injection defenses largely treat external directives as threats to reject wholesale. EDIC extends process-aware supervision to mixed external contexts where useful and malicious instructions coexist.

2.1 Prompt injection and external instruction safety

  • Existing defenses and related work rarely model settings where beneficial and harmful external instructions coexist in one document, the gap addressed by fine-grained classification.The related methods include instruction prioritization, goal prioritization, safety-aware decoding, self-reminders, and ReasAlign.Source: §2.1 · p. 2

2.2 Reasoning supervision and alignment for multi-step decisions

  • EDIC extends process-aware supervision to attribute behavior per instruction in external contexts containing both useful and malicious directives.Prior process-supervision work rewards intermediate reasoning in hard reasoning regimes, but does not address this mixed-instruction setting.Source: §2.2 · p. 3

3 Preliminary

The problem setting treats external data as a collection of snippets that may be beneficial, harmful, or contextual. The model must identify and classify each snippet, then choose an action consistent with its type.

Problem setting

  • External data D is represented as a set of snippets, with beneficial instructions relevant and safe, harmful instructions attempting to override the query or trigger unsafe behavior, and contextual snippets lacking directive intent.Source: §3 · p. 3

Decision framework

  • The decision framework requires extracting snippets from D and assigning each extracted snippet a type.Source: §3 · p. 3
  • For each snippet, the model chooses to follow or reject it consistently with its type and produces a response reflecting those decisions.Controlled compliance means following beneficial instructions and rejecting harmful ones; silently following harmful instructions or rejecting beneficial ones is failure.Source: §3 · p. 3

4 Method

EDIC instantiates an identify-classify-act framework by synthesizing mixed-instruction data, training a generative judge to attribute compliance, and using decomposed judge rewards to guide a CoT policy. Its stages are designed to distinguish beneficial compliance from harmful compliance rather than obeying or rejecting all external instructions uniformly.

4 Method

  • EDIC combines scene-grounded data synthesis, a generative judge for per-instruction attribution, and judge-guided training of a CoT policy.The framework uses behavioral rewards to train the policy to classify and follow instructions.Source: §4 · p. 3

4.1 Data synthesis

  • Training examples are grounded in real documents, and counterfactual checks remove samples that fail naturalness, necessity, or consistency tests.Removing a beneficial instruction should change the response, while removing a harmful one should not.Source: §4.1 · pp. 3–4

4.2.1 Judge Model

  • The judge is cold-started with supervised fine-tuning and refined with GRPO using attribution-correctness and evidence-grounded-format rewards.The RL stage encourages concise reasoning and improves robustness on ambiguous compliance cases not well covered by SFT alone.Source: §4.2.1 · p. 4

4.2.2 CoT Model

  • The CoT policy’s SFT stage teaches it to identify instruction-like snippets, classify them as beneficial or harmful, and respond consistently with those labels.SFT supplies a structural scaffold but does not ensure robust classification under adversarial instructions.Source: §4.2.2 · p. 4
  • The RL reward gives separate credit for following beneficial instructions and penalties for following harmful ones, preventing always-refuse or always-obey policies from maximizing reward.This closed loop lets the model learn controlled compliance without classification-level annotations.Source: §4.2.2 · pp. 4–5

5 Experimental Setup

The experiments use two open-weight backbones, scene-grounded training data, and evaluations spanning utility, prompt-injection security, and held-out compliance measures. EDIC is compared with base models and five prompt-injection defense baselines.

Models and training

  • Experiments use Qwen3-4B and Llama-3.1-8B-Instruct, with judge and CoT models trained using SFT followed by GRPO.Full hyperparameters are provided in Appendix B.Source: §5 · p. 5

Training data

  • Training data covers email, API documentation, medical guidance, and knowledge evidence, retaining 900 examples per scene after filtering.Each scene is split into validation, judge-training, and policy-training sets in a 150/250/500 ratio.Source: §5 · p. 5

Baselines and evaluation

  • The evaluation measures held-out compliance and judge performance, utility on MT-Bench and Arena-Hard-Auto, and safety on AgentDojo, CyberSecEval 2, and InjecAgent.Comparisons include the base model and five defense baselines, including ReasAlign.Source: §5 · p. 5

6 Results and Analysis

EDIC leads the reported safety and utility comparisons by separately rewarding beneficial compliance and penalizing harmful compliance. Ablations and analysis support mixed-context training and decomposed rewards, while examining performance across backbones, triggers, and scenes.

6.1 Main Results

  • EDIC achieves the highest harmful-rejection and benign-follow rates and strongest Arena-Hard and MT-Bench utility among the compared defenses.The comparison is reported in Table 1; the supplied passages give no table-cell values.Source: §6.1 · p. 5, Table 1 · p. 6
  • On Llama-3.1-8B, EDIC leads on benign-follow and matches the strongest baseline on harmful-rejection, while its Arena-Hard gap over ReasAlign narrows from +12.5% on Qwen3-4B to +4.6%.The classification trigger improves safety at a small utility cost on both backbones; the authors describe it as an inference-time switch that needs no retraining.Source: §6.1 · pp. 5–6

6.2 Ablation Studies

  • Mixed-context training is essential for controlled compliance: when beneficial and harmful instructions appear in separate documents, the model need not discriminate between co-occurring directives.The authors distinguish this objective from raw safety alone.Source: §6.2 · p. 6, Table 3 · p. 7
  • With the sample count fixed at 3,600, stricter beneficial-instruction necessity filtering improves both safety and utility, supporting necessity as a proxy for causal coupling rather than sample size.High-necessity examples are described as providing a cleaner supervision signal for the judge and policy.Source: §6.2 · pp. 6–7, Table 4 · p. 7
  • Judge and policy SFT/RL optima differ: the judge favors 60/40 on Qwen3-4B and 80/20 on Llama-3.1-8B, while the Qwen3-4B policy favors a more RL-heavy 20/80 configuration.Pure RL lacks the judge’s structural cold-start; for the policy, pure SFT lacks robust safety and pure RL collapses without its scaffold.Source: §6.2 · p. 7, Figure 3 · p. 7, Figure 4 · p. 8
  • Only the full four-component reward yields high harmful-rejection with high benign-follow; removing structure collapses utility, while removing either directional signal encourages over-refusal or over-compliance.Removing the judge signal causes the largest single safety drop because it removes grounded per-instruction supervision.Source: §6.2 · p. 7, Figure 5 · p. 8

6.3 Analysis Studies

  • On NotInject’s benign prompts containing trigger words, EDIC stays close to the Base model and degrades slowly as trigger density rises, unlike the visible declines for SecAlign and ReasAlign.The evaluation uses 339 samples and reports benign accuracy across subsets with one, two, or three trigger words.Source: §6.3 · p. 8, Table 6 · p. 8
  • Across scenes, EDIC improves beneficial-follow, harmful-reject, and controlled-compliance while reducing the across-scene controlled-compliance standard deviation.The comparison is between Base and EDIC on Qwen3-4B with the trigger enabled.Source: §6.3 · p. 8, Table 5 · p. 8

7 Conclusion

Existing prompt-injection defenses fail by treating external instructions uniformly, sacrificing either safety or utility. EDIC trains models to identify, classify, and act on each instruction, improving safety and utility across two backbones and five benchmarks.

7 Conclusion

  • Uniformly treating external instructions as trustworthy or untrustworthy creates a classification failure that sacrifices safety or utility.The conclusion frames this as the core failure of existing prompt-injection defenses.Source: §7 · p. 8
  • EDIC consistently improves both safety and utility over the strongest prior defenses across two backbones and five external benchmarks.Ablations identify mixed-context data, the necessity filter, two-stage judge training, and each reward component as load-bearing.Source: §7 · p. 8
  • EDIC assigns per-instruction rewards that separately encourage beneficial compliance and harmful rejection.The resulting identify-classify-act training aims to avoid uniform obedience or rejection.Source: §7 · p. 8

Limitations

The paper identifies limits in training coverage, attack adaptation, binary judging, reward-weight tuning, and model scale. It also notes potential misuse of its attack materials and describes safeguards for data and release.

Limitations

  • Training uses four scene types; generalization to long-horizon agents, complex HTML, non-English documents, adaptive attacks, and larger-scale or multi-turn settings remains unvalidated.The judge also uses binary labels for partly safe and partly malicious instructions, and reward weights are tuned per backbone rather than automatically.Source: Limitations · p. 9

Ethical Considerations

  • Attack templates and synthetic adversarial documents could be repurposed against production systems, but the work uses public benchmark triggers and no real user data or personally identifiable information.Authors produce annotations with a structured rubric; any release is planned under a research-use license with restrictions on misuse.Source: Ethical Considerations · p. 9

A Baseline Descriptions

  • Baseline defenses use preference alignment, input isolation or filtering, instruction-source hierarchy, meta-learning across attack styles, or judge-guided alignment.ReasAlign, unlike EDIC, does not model mixed beneficial and harmful instructions or use decomposed per-instruction rewards.Source: A Baseline Descriptions · p. 13

B Training Hyperparameters

Training configuration specifies separate SFT and GRPO stages for the judge and response policy, alongside reward weights and compute infrastructure.

B Training Hyperparameters

  • Judge SFT uses learning rate 5 × 10−5, one epoch, LoRA rank 32, and maximum sequence length 2,048; judge GRPO uses learning rate 1× 10−6 and rollout group size 8.Judge SFT uses AdamW and cosine decay; GRPO uses a constant schedule.Source: B Training Hyperparameters · p. 13
  • Policy SFT uses backbone-specific learning rates of 2×10−5 for Qwen3-4B and 1×10−5 for Llama-3.1-8B; policy GRPO uses learning rate 1×10−6 and rollout group size 8.The SFT maximum sequence length is 4,096; GRPO maximum response length is 768.Source: B Training Hyperparameters · p. 13
  • The decomposed reward combines beneficial compliance, a harmful-compliance penalty scaled by λ, and a scene-appropriate structure bonus scaled by α.For Qwen3-4B, λ = 1.0 and α = 0.5.Source: B Training Hyperparameters · p. 13
  • Training uses NVIDIA A100-80GB GPUs with CUDA 12.1, VERL for RL optimization, and vLLM 0.6 for rollouts and judge inference.FlashAttention-2 is enabled on both backbones.Source: B Training Hyperparameters · p. 13

C Benchmark Configuration Details

The benchmark configurations cover general assistant utility, agent and prompt-injection security, and injection-specific task suites. AgentDojo evaluates both task completion under attack and attack success.

C Benchmark Configuration Details

  • MT-Bench uses 80 questions and reports first-turn, second-turn, and average scores; Arena-Hard-Auto uses 750 hard prompts and reports an Arena score with a 95% confidence interval.Both use DeepSeek-Chat as judge.Source: C Benchmark Configuration Details · p. 13
  • AgentDojo evaluates the important_instructions attack across banking, Slack, travel, and workspace suites, totaling 629 attacked cases.Each case has a clean pass and an attacked pass where the injection is appended to a tool or document return; DJ-U measures task completion and DJ-ASR measures attacker-goal execution.Source: C Benchmark Configuration Details · pp. 13–14
  • CyberSecEval 2 uses the prompt_injection task with 250 samples and reports its primary metric through inspect_ai and inspect_evals.The specified judge is DeepSeek-Chat for interpreter_abuse and prompt_injection tasks.Source: C Benchmark Configuration Details · p. 14
  • InjecAgent uses base and enhanced settings, with approximately 400 cases per setting across its DH and DS splits.The limit is 200 per split, and the model is served through a local vLLM endpoint.Source: C Benchmark Configuration Details · p. 14

D Additional Analyses

The supplementary analyses examine cross-scene generalization, injection-position robustness, and judge–policy agreement, alongside qualitative examples from email inputs.

D Additional Analyses

  • Email-only training yields an average OOD CC of 0.81 versus 0.90 with full-data training, a 9pp gap; H-R drops 8–11pp and B-F 9–12pp.The largest H-R drop is 11.0pp on API, the most challenging scene due to its instruction-dense format.Source: D Additional Analyses · p. 14, Table 7 · p. 14

D.1 Injection position robustness

  • 93.4% H-R for suffix injection exceeds 92.6% for prefix and 91.2% for mid-context injection by 2.2pp.Mid-context is the default training position and the most realistic, challenging setting because the harmful instruction blends with factual content.Source: §D.1 · p. 14, Figure 6 · p. 14

D.2 Judge–policy agreement

  • Judge–policy agreement rises from 62.4% (Base) to 89.2% (SFT+RL), while Cohen’s κ improves from 0.38 to 0.81.Cases where both judge and policy fail drop from 24.8% to 5.6%; judge-correct-and-policy-safe rate correlates 0.91 with external security metrics.Source: §D.2 · p. 15, Table 8 · p. 15

E Case Study

  • The case study presents two qualitative examples from the email scene using the with-trigger variant.These examples illustrate how the judge and CoT response model handle mixed-instruction inputs.Source: E Case Study · p. 15

E.1 Judge reasoning case

This section presents judge-reasoning cases and a four-step procedure for synthesizing policy SFT examples, illustrating instruction-level attribution and controlled compliance with beneficial directives and injections. It also describes four ablations testing whether embedded beneficial instructions are necessary.

E.1 Judge reasoning case

  • The judge correctly attributes compliance to each candidate instruction: it marks beneficial execution and harmful-instruction compliance yes, but an ignored injection no.This local attribution supports separate credit and penalty in the decomposed reward.Source: §E.1 · p. 15

E.2 CoT response model case

  • The EDIC response follows the beneficial request to attach the agreement and cc the sender while ignoring the injected antonym question.Its reasoning identifies the primary email task, distinguishes the relevant instruction from the injection, and the final email includes the attachment and CC.Source: §E.2 · pp. 15–16

F Prompts

  • The prompt appendix states that prompts are invoked with GPT-5.4 at temperature 1.0 and that placeholders are filled with sample fields at runtime.Prompts are rendered in breakable frames so long prompts flow across pages without truncation.Source: F Prompts · p. 16

F.1 Ablation generation for the necessity check

  • The necessity check generates four allow-path outputs after removing or weakening the beneficial instruction: removing it or rewriting it as background, optional, or nonbinding context.Each variant is tested with the same user instruction.Source: §F.1 · p. 16

F.2 Necessity check (e+); F.3 Consistency check; F.4 Naturalness check

  • The necessity, consistency, and naturalness checks retain only samples scoring 5, requiring a necessary beneficial instruction, a valid allow/forbid contrast, and a natural scene artifact.Consistency checks that the allow output completes the authorized task while the forbid output stays informational; naturalness assesses the scene, instruction placement, outputs, and user requests.Source: §F.2 · p. 16, §F.3 · p. 17, §F.4 · p. 17

F.5 Scene-specific standards

  • The four scene families use shared rubrics with scene-specific standards inserted into the necessity, consistency, and naturalness prompts.Table 9 lists the standards used for each scene.Source: §F.5 · p. 17, Table 9 · p. 18

F.6 Independence check (e−)

  • The independence check compares authorized responses with and without harmful instruction e− and removes samples scoring above 2 to limit contamination of the supervised target.A score of 1 means the responses are functionally identical; a score of 2 indicates only cosmetic differences.Source: §F.6 · pp. 17–18

F.7 Judge cold-start reasoning synthesis

  • For judge cold-start SFT on DJ, GPT-5.4 rewrites placeholder reasoning using the original prompts and gold yes/no label, then the cleaned reasoning is paired with the unchanged label.The synthesis prompt requires brief reasoning grounded in the context and consistent with the gold label.Source: §F.7 · p. 18

F.8 CoT cold-start reasoning synthesis

  • Policy cold-start reasoning enumerates instruction-like spans, classifies each as BENEFICIAL, HARMFUL, or CONTEXTUAL, and commits to following or rejecting each before answering.Unlike the judge prompt’s attribution of one instruction at a time, policy reasoning justifies a joint decision over instructions in the document.Source: §F.8 · p. 19
  • CoT synthesis has the model identify every instruction-like span, classify it from document cues, plan follow/reject decisions, and verify the response against the user query and answer.The cleaned reasoning is concatenated with the unchanged answer to form the supervised target.Source: §F.8 · pp. 19–20

Related papers

Loading 2610.22874v1…