Research map / Large language models / Language, safety and interpretability
Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Instructional Segment Embedding (ISE), a learned segment-type embedding added to token embeddings that tags system, user, data, and output tokens so LLMs can distinguish instruction priority, trained via SFT on Llama models.
Problem
LLMs treat all input tokens equally and lack an architectural mechanism to distinguish and prioritize instruction types (system, user, data), making them vulnerable to prompt injection, prompt extraction, and harmful requests.
Approach
Add a BERT-style learned segment embedding, with one vector per instruction role (system, user, data, output), summed with token and positional embeddings, so instruction priority is encoded at the architecture/embedding level rather than only via delimiters or instruction-tuning data.
Key results
- On Structured Query (Clean Alpaca, Llama-2-13B), ISE improves average in-domain robust accuracy by 15.75% (50.84 to 66.35) and worst by 32.17%.
- On Instruction Hierarchy, ISE gives robustness gains from 5% to 25% across training configurations, with up to ~25% average gain for UltraChat Baseline.
- AlpacaEval win rate improves by ~4.1% when training on the Instruction Hierarchy dataset with ISE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add ISE | Llama-2-13B trained on Clean Alpaca · Structured Query, in-domain indirect prompt injection | average robust accuracy: 50.84 -> 66.35; 15.75% | Helped |
| Add ISE | Llama-2-13B trained on Clean Alpaca · Structured Query, out-of-domain indirect prompt injection | average robust accuracy: 47.12 -> 61.06; ∼10% | Helped |
| Add ISE | Llama-2-13B trained on Adversarial Alpaca · Structured Query, out-of-domain indirect prompt injection | average robust accuracy: 80.53 -> 90.26; ∼7% | Helped |
| Add ISE | Llama-3-8B trained on Instruction Hierarchy · AlpacaEval | win rate: ∼4.1% | Helped |
| Add ISE | Llama-3-8B across Instruction Hierarchy training datasets · Instruction Hierarchy, prompt extraction | average robust accuracy: at least 10% | Helped |
| Add ISE | Llama-3-8B trained on UltraChat Baseline · Instruction Hierarchy, direct prompt injection, in-domain | average robust accuracy: 47.3% -> 69.9% | Helped |
| Add ISE | Llama-3-8B trained on UltraChat Baseline · Instruction Hierarchy, Completion Real prompt injection | robust accuracy: approximately 30% to 50% | Helped |
| Add ISE | Llama-3-8B trained on UltraChat Baseline · Instruction Hierarchy, harmful requests | robust accuracy: up to 18% | Helped |
| Add ISE | Llama-3-8B trained on UltraChat Baseline · WildChat over-refusal evaluation | compliance rate: about 10% | Helped |
| Add ISE | Llama-3.1-8B trained on UltraChat Baseline · AlpacaEval | win rate: approximately 10% | Helped |
| Use ISE with special token delimiters instead of ISE alone (ablation) | Llama-2-13B on Structured Query · Structured Query, Adversarial Alpaca, indirect prompt injection | average robust accuracy: 94.60 -> 99.83 | Helped |
| Omit the system prompt (ablation) | ISE models on Instruction Hierarchy · AlpacaEval | win rate: < 5% | Hurt |
Surprising findings
- Using ISE alone without special delimiters does not improve performance on Clean Alpaca, so ISE is best combined with delimiters.
- Models trained on UltraChat Baseline with ISE showed highest prompt-extraction robustness, partly because weaker instruction following yields partial system prompts that count as failed attacks.
Limitations
- Limited robustness against adaptive attacks and jailbreaks.
- Only supervised fine-tuning studied with single-turn conversations; pretraining/RLHF not explored.
- Experiments limited to 8B and 13B models and datasets under 300K samples.
Experimental setup
Datasets: Alpaca-Cleaned-50K UltraChat-200K SystemChat SystemMessage ShareGPT Unnatural Instructions WildChat Jailbreakchat DAN StrongREJECT
Benchmarks: Structured Query Instruction Hierarchy AlpacaEval 1.0 MT-Bench WildChat over-refusal
Metrics: robust accuracy average robustness worst robustness win rate MT-Bench score compliance rate
Baselines: Baseline (no ISE) Delimiter (StruQ) Llama-3-Instruct default template
Builds on: BERT segment embeddings StruQ instruction hierarchy SFT
Closest papers by meaning
- Context-Aware Hierarchical Learning: A Two-Step Paradigm towards Safer LLMs (NeurIPS 2025 · similarity 0.90)
- Can LLMs Separate Instructions From Data? And What Do We Even Mean By That? (ICLR 2025 · similarity 0.90)
- HarmAug: Effective Data Augmentation for Knowledge Distillation of Safety Guard Models (ICLR 2025 · similarity 0.89)
- On Prompt-Driven Safeguarding for Large Language Models (ICML 2024 · similarity 0.89)
- SG-Bench: Evaluating LLM Safety Generalization Across Diverse Tasks and Prompt Types (NeurIPS 2024 · similarity 0.89)
- InstructZero: Efficient Instruction Optimization for Black-Box Large Language Models (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.