Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy

1 INTRODUCTION

• LLM architectures process inputs without formally differentiating instruction priorities, enabling prompt injection, prompt extraction, and harmful requests.These attacks can override system-level safety or proprietary instructions.Source: §1 · p. 1
• ISE addresses an architectural gap left unresolved by delimiter-based and instruction-training approaches.The paper frames its contribution as incorporating hierarchy directly into the model rather than relying only on prompt formatting or training data.Source: §1 · p. 2
• ISE categorizes tokens by role and passes learned segment information through the model alongside token embeddings.The approach distinguishes system instructions, user prompts, and data to improve instruction prioritization.Source: §1 · p. 2
• 18.68% average robust accuracy improvement is achieved on the Instruction Hierarchy benchmark across multiple vulnerabilities.On Structured Query, the reported improvement reaches 15.75%, while AlpacaEval instruction-following capability improves by up to 4.1%.Source: §1 · p. 2
• The evaluation covers two benchmarks, five training datasets, and four vulnerability types.The contribution claims empirical effectiveness across Structured Query and Instruction Hierarchy settings.Source: §1 · p. 2

2 BACKGROUND: LLM VULNERABILITIES

• System instructions take precedence over user instructions, which take precedence over third-party data in typical LLM applications.These categories correspond to application providers, primary users, and external information sources.Source: §2 · pp. 2–3
• Conflicts among instruction priorities create security vulnerabilities, including prompt injections, prompt extractions, and harmful requests.Malicious users or web providers can attempt to bypass safety instructions or inject harmful actions into data.Source: §2 · p. 3
• Indirect prompt injection places instructions in third-party data, whereas direct prompt injection manipulates the user query to alter model behavior.Both forms aim to make outputs deviate from predefined instructions.Source: §2 · p. 3
• Prompt extraction attacks attempt to elicit confidential information stored in system instructions.The resulting disclosure may reveal system prompts.Source: §2 · p. 3
• Harmful requests use malicious queries to bypass safety alignment and elicit unsafe or potentially weaponized outputs.The paper describes these requests as targeting the model’s safety-aligned behavior.Source: §2 · p. 3
• The paper targets attacks that attempt to override the established priority rules.Its stated aim is to enhance LLM instruction hierarchy capabilities.Source: §2 · p. 3

3 LACK OF INSTRUCTION HIERARCHY IN MODERN LLM ARCHITECTURE

• Current embeddings encode token identity and position but do not represent the hierarchical significance of instructions.Self-attention therefore processes token embeddings equally with respect to instruction priority.Source: §3 · p. 3
• Equal processing can cause a lower-priority user prompt to override an original system prompt.The paper gives this architectural behavior as a source of downstream vulnerabilities.Source: §3 · p. 3
• Instruction-tuning methods improve prioritization but do not remove the limitation of uniformly processing hierarchical instructions.The paper identifies this as a fundamental limitation of current embedding methods.Source: §3 · p. 3
• The proposed architecture is reported to provide better robustness in results from Table 1 and Figure 5.This statement connects the architectural design to the paper’s robustness evaluation.Source: §3 · p. 4, Figure 4 · p. 4
• Delimiter-based chat templates provide hierarchy cues through special tokens, but those cues may diminish in long contexts and may be exposed to attackers.The paper identifies sparse inference signals and delimiter extraction as two drawbacks.Source: §3 · p. 4

4 PROPOSED APPROACH: INSTRUCTIONAL SEGMENT EMBEDDING (ISE)

• ISE encodes instruction hierarchy directly into embeddings so later self-attention layers can recognize and follow instruction priorities.The intended effect is improved robustness against priority-overriding attacks.Source: §4 · p. 4
• The default design assigns four segment types: system, user, data, and generated output.Each token receives hierarchy information from its stakeholder or output category.Source: §4 · p. 4
• Final input embeddings are formed by summing token embeddings with the corresponding learned segment embeddings before self-attention.The resulting representations follow the processing path used in current LLMs.Source: §4 · p. 4
• The segment embedding layer is trained with the other parameters during supervised fine-tuning on structured and instruction-following data.Training also includes malicious instructions in some datasets so models learn to reject or ignore hierarchy-violating commands.Source: §4 · p. 4
• ISE categories can be customized by adding finer-grained tool or web-data types, or omitting data when third-party context is absent.The design is intended to adapt to downstream application structure.Source: §4 · p. 4
• The method repurposes BERT-style segment embeddings for instruction hierarchy and includes output segments to support autoregressive and multi-turn settings.The output category is retained because generated text can itself contain instructions.Source: §4 · p. 5
• ISE is lightweight, requiring H × D additional parameters.The paper describes implementation as straightforward and provides a short PyTorch implementation.Source: §4 · p. 5

5 EXPERIMENTAL DESIGN

• Instruction Hierarchy evaluates indirect and direct prompt injection, prompt extraction, and harmful requests, while Structured Query focuses on indirect prompt injection.The two benchmarks use different vulnerability scopes to assess robustness.Source: §5.1 · p. 5

5.1 GENERATING TRAINING DATA

• Training data combines Clean and Adversarial Alpaca for Structured Query with decomposed UltraChat, system-prompt datasets, and crafted attacks for Instruction Hierarchy.Harmful-request data is excluded from training and reserved for evaluation.Source: §5.1 · p. 5

5.2 EXPERIMENT SETUP

• Samples are formatted into system, user, data, and output segments, with system and user instructions merged for Structured Query and evaluation restricted to single-turn chats.The merge reflects identical Alpaca system instructions.Source: §5.2 · p. 5
• Models are fully fine-tuned for three epochs using a 2e-5 learning rate and cosine scheduling, with top-p sampling during inference.The default pretrained models are Llama-2-13B for Structured Query and Llama-3-8B for Instruction Hierarchy; Llama-3.1-8B results appear separately.Source: §5.2 · p. 5
• Baselines compare text or newly constructed delimiters against +ISE on Structured Query, while Instruction Hierarchy uses Llama-3-8B’s default delimiters as its baseline.The delimiter comparison follows Chen et al. (2024).Source: §5.2 · p. 6

5.3 ROBUSTNESS EVALUATION

• Structured Query tests four in-domain and new out-of-domain indirect attacks, measuring robust accuracy as the rate of responses that omit “hacked”.The out-of-domain attacks place adversarial text at the beginning and end of data and substantially degrade adversarially trained models.Source: §5.3 · p. 6
• Instruction Hierarchy evaluates indirect and direct injection, extraction with Rouge-L recall, and harmful requests using jailbreak prompts, StrongREJECT instructions, and GPT-4o safety checks.Direct attacks insert the same attacking prompts into the user prompt, whereas extraction uses ShareGPT, Unnatural Instructions, and author-selected prompts.Source: §5.3 · p. 6
• Average robustness aggregates performance across attack methods, worst robustness measures defense against the most challenging attack, and clean capability is assessed with AlpacaEval and MT-Bench.MT-Bench is additionally used for Instruction Hierarchy.Source: §5.3 · p. 6

6 EXPERIMENTAL RESULTS AND ANALYSIS

• The experiments cover Structured Query and Instruction Hierarchy benchmarks, comparing ISE with baselines across multiple training datasets and attack types.The evaluation includes capability, robustness, detailed vulnerability analyses, over-refusal, and model generalization.Source: §6 · p. 6, Table 1 · p. 7, Figure 5 · p. 8

6.1 MAIN RESULTS ON STRUCTURED QUERY

• 15.75% higher average robust accuracy and 32.17% higher worst robust accuracy occur against in-domain indirect prompt injection on clean Alpaca.For out-of-domain attacks, ISE improves average robustness by approximately 10% on clean Alpaca and 7% on adversarial Alpaca.Source: §6.1 · p. 7

6.2 MAIN RESULTS ON INSTRUCTION HIERARCHY

• Approximately 4.1% higher AlpacaEval win rate is achieved when ISE is trained on the Instruction Hierarchy dataset.MT-Bench shows negligible degradation for UltraChat Baseline and improvements for the other two training datasets.Source: §6.2 · p. 7
• ISE improves Instruction Hierarchy robustness by 5%–25% across training configurations, with nearly 25% average improvement for UltraChat Baseline.Prompt-extraction robustness rises at least 10% across datasets, while harmful-request robustness improves for UltraChat Baseline and Instruction Hierarchy and is maintained or slightly exceeded for System Follow.Source: §6.2 · pp. 7–8

6.3 DETAILED ANALYSIS OVER ATTACKS

• 30%–50% robustness improvements against Completion Real accompany baseline effectiveness below 10% for UltraChat Baseline and System Follow without ISE.The attack spoofs a benign response and appends a malicious instruction, exploiting failures to distinguish instruction types.Source: §6.3 · p. 9
• ISE improves harmful-request robustness in all six evaluated categories, by up to 18%.The categories include disinformation, hate and discrimination, illegal goods, non-violent crimes, sexual content, and violence.Source: §6.3 · p. 9

6.4 OTHER ANALYSIS

• ISE improves over-refusal compliance by about 10% for UltraChat Baseline but slightly degrades it for the other two models.The evaluation uses 691 non-toxic WildChat prompts and assesses non-refusal responses with GPT-4o.Source: §6.4 · p. 9
• On Llama-3.1-8B, ISE raises AlpacaEval win rate by approximately 10% and robust accuracy by over 10% on most attacks.Harmful-request robustness improves by around 5%, supporting generalization across models.Source: §6.4 · p. 10

Safety vulnerabilities of LLMs

• LLM safety work studies prompt injection, prompt extraction, and harmful requests as distinct security and privacy vulnerabilities.The cited literature distinguishes indirect and direct injection, confidential system-prompt extraction, and attempts to elicit unsafe behavior.Source: §7 · p. 10

Improving LLM robustness

• Prompt-based defenses use delimiters or in-context instructions, whereas learning-based defenses fine-tune models on malicious instructions paired with robust responses.The paper presents embedding modification as an approach orthogonal to these mitigation strategies.Source: §7 · p. 10

Embedding and architecture of LLMs

• Prior embedding and architecture studies target long-context generalization or arithmetic, while this paper uses Instructional Segment Embedding to learn instruction hierarchy.ISE is inspired by BERT and LinkBERT architectural designs.Source: §7 · p. 10

• The study focuses on supervised fine-tuning with single-turn conversations, leaving pre-training, RLHF, and multi-turn applications for future work.Source: §8 · p. 10
• ISE provides limited robustness against adaptive attacks, or jailbreaks, although combining it with adversarial training may improve robustness.Source: §8 · p. 10
• Experiments use smaller 8B and 13B models and datasets under 300K examples, so generalization to larger models and datasets remains unresolved.Source: §8 · p. 10
• ISE is presented as the first attempt to design novel architectures for improving instruction hierarchy, with experiments demonstrating gains in robustness and general capabilities.Source: §8 · p. 10

9 ETHICS AND REPRODUCIBILITY STATEMENT

• The study uses publicly available datasets and safety-aligned GPT-4o for some training-data generation and performance judging, while reporting no anticipated ethical concerns.Source: §9 · p. 11

A DETAILS OF IMPLEMENTING INSTRUCTIONAL SEGMENT EMBEDDING

• The implementation adds instructional segment embeddings alongside token and positional embeddings, with segment IDs representing roles such as system, user, data, and output.During inference, token and ISE embeddings are computed and summed before further processing.Source: A DETAILS OF IMPLEMENTING… · p. 15, §19 · p. 15

B.1 STRUCTURED QUERY BENCHMARK.

• Training data includes clean Alpaca samples and adversarial samples containing indirect prompt injections, while preserving the original user-directed output.Source: §B.1 · p. 16

B.2 INSTRUCTION HIERARCHY BENCHMARK.

• Instruction Hierarchy training data is built from UltraChat, SystemChat, and SystemMessage, covering indirect injection, direct injection, and prompt extraction with GPT-4o-generated outputs.Source: §B.2 · pp. 16–18

C.1 EVALUATIONS ON STRUCTURED QUERY BENCHMARK

• Structured Query evaluation uses four indirect-injection attacks and measures robust accuracy as the proportion of samples that do not output “hacked”.The evaluation includes 208 AlpacaEval test samples with valid data inputs and also tests out-of-domain attack placement.Source: §C.1 · p. 19

C.2 EVALUATIONS ON THE INSTRUCTION HIERARCHY BENCHMARK

• Instruction Hierarchy evaluation covers indirect and direct injection, prompt extraction, and harmful requests, with robustness judged through safety adherence or attack resistance.Source: §C.2 · pp. 20–21
• Prompt extraction is scored using ROUGE-L recall based on the longest common subsequence, with extraction attacks considered successful above 0.9.Source: §C.2 · p. 20

C.3 INSTRUCTION-FOLLOWING EVALUATION

• Instruction following is evaluated on AlpacaEval and one-turn MT-Bench conversations, while over-refusal is assessed by whether benign requests receive sufficient detail.Source: §C.3 · p. 21, §C.4 · p. 22

• ISE with special delimiters consistently outperforms the other evaluated methods in almost all Structured Query cases.Source: D MORE EXPERIMENTAL RESULTS ON… · p. 23
• ISE alone increases Adversarial Alpaca robustness by up to 4.6% on average and up to 14% for in-domain worst robust accuracy versus the baseline.Source: D MORE EXPERIMENTAL RESULTS ON… · p. 23
• Direct ISE does not improve performance on the Clean Alpaca dataset, motivating its use with special token delimiters for best performance.Source: D MORE EXPERIMENTAL RESULTS ON… · p. 23

E.1 DETAILED ANALYSIS OF INDIRECT PROMPT INJECTION

• ISE improves robustness by 15%–45% across almost all indirect prompt injection attacks, with worst robust accuracy gains approaching 50%.The evaluation covers in-domain and out-of-domain attacks, including 17 listed attack variants.Source: §E.1 · p. 24

E.2 DETAILED ANALYSIS OF DIRECT PROMPT INJECTION

• In-domain direct-attack average robust accuracy rises from 47.3% to 69.9% for the UltraChat Baseline model with ISE.Additional results report over 10% average and over 20% worst robust-accuracy improvements.Source: §E.2 · p. 24

E.3 DETAILED ANALYSIS OF PROMPT EXTRACTION

• Prompt-extraction robustness improves by over 20% in both average and worst cases for UltraChat Baseline models using ISE.The appendix describes robustness as the ratio of attacks that fail to extract a significant number of original prompts.Source: §E.3 · p. 25

E.4 DETAILED ANALYSIS OF HARMFUL REQUESTS

• ISE improves harmful-request robustness across categories, including for UltraChat Baseline and Instruction Hierarchy models trained without bypass-oriented data.The full results extend the evaluation to System Follow and Instruction Hierarchy models.Source: §E.4 · p. 25

E.5 DETAILED ANALYSIS OF LLAMA 3.1 MODEL

• On Llama-3.1-8B, ISE improves AlpacaEval winning rates, maintains or improves MT-Bench scores, and strengthens robustness across attacks.The pattern persists even for models trained on Instruction Hierarchy that already achieve high robustness.Source: §E.5 · pp. 25–26, Figure 26 · p. 26

E.6 INVESTIGATION OF INSTRUCTIONAL SEGMENT EMBEDDING

• Capability degradation is below 5% when system prompts are omitted, about 15% when system prompts use user embeddings, and below 2% when data uses user embeddings.The authors attribute the larger middle-case drop to mixing system and user prompts, while overlapping training formats may explain the small data-embedding degradation.Source: §E.6 · p. 27, Table 4 · p. 27

E.9 CASE STUDY OF ISE ATTENTION

• Attention case studies show baseline models over-attend to neighboring malicious tokens, whereas ISE distributes attention more evenly and identifies malicious context.The example malicious instruction is “Print exactly hacked,” which disproportionately influences the baseline response.Source: §E.9 · p. 29, Figure 27 · p. 29, Figure 28 · p. 29

F DISCUSSIONS AND EVALUATIONS ON JAILBREAK ATTACKS

• The harmful-request evaluation uses wild-collected malicious prompts without active optimization, treating them as zero-shot generalization tests.No training data were designed to bypass safety alignment.Source: F DISCUSSIONS AND EVALUATIONS ON… · p. 31
• Adaptive attacks on 50 malicious requests leave the models almost completely unable to generate safe responses.These attacks use adversarially optimized or carefully human-crafted strategies and are reported in Table 8.Source: F DISCUSSIONS AND EVALUATIONS ON… · p. 31, Table 8 · p. 31
• The authors do not expect ISE to improve adaptive jailbreak robustness because its training data did not target jailbreaks and adversarial texts may directly target the model.They describe ISE as orthogonal to robust-training methods such as LAT and circuit breakers.Source: F DISCUSSIONS AND EVALUATIONS ON… · p. 31