Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy

TL;DR

This work identifies the absence of a formal instruction hierarchy as a core vulnerability in LLMs and proposes Instructional Segment Embedding (ISE), an architecture-level method that injects priority signals into token embeddings via a learnable segment embedding. By distinguishing system, user, data, and output segments, ISE improves robustness to prompt-injection, prompt extraction, and harmful prompts while preserving instruction-following capabilities across benchmarks. Empirical results on Structured Query and Instruction Hierarchy show substantial robustness gains (up to ~15.75pp in-domain, up to ~32pp worst-case) and positive transfers to multi-turn tasks, with generalization to additional model families. The approach is lightweight, easy to implement, and complementary to existing defenses, offering a promising path toward safer and more reliable LLM architectures.

Abstract

Large Language Models (LLMs) are susceptible to security and safety threats, such as prompt injection, prompt extraction, and harmful requests. One major cause of these vulnerabilities is the lack of an instruction hierarchy. Modern LLM architectures treat all inputs equally, failing to distinguish between and prioritize various types of instructions, such as system messages, user prompts, and data. As a result, lower-priority user prompts may override more critical system instructions, including safety protocols. Existing approaches to achieving instruction hierarchy, such as delimiters and instruction-based training, do not address this issue at the architectural level. We introduce the Instructional Segment Embedding (ISE) technique, inspired by BERT, to modern large language models, which embeds instruction priority information directly into the model. This approach enables models to explicitly differentiate and prioritize various instruction types, significantly improving safety against malicious prompts that attempt to override priority rules. Our experiments on the Structured Query and Instruction Hierarchy benchmarks demonstrate an average robust accuracy increase of up to 15.75% and 18.68%, respectively. Furthermore, we observe an improvement in instruction-following capability of up to 4.1% evaluated on AlpacaEval. Overall, our approach offers a promising direction for enhancing the safety and effectiveness of LLM architectures.

Figures (29)

• Figure 1: A demonstration of the hierarchy of instructions, including system instruction, user instruction, data input as well as model output.
• Figure 2: A demonstration of various vulnerabilities of LLM applications, including prompt injection, prompt extraction as well as harmful request.
• Figure 3: A demonstration of the chat template for Llama-3-Instruct dubey2024llama3herdmodels.
• Figure 4: The input representation includes both token embeddings and instructional segment embeddings. We categorize all input texts into four segments: fooD system instructions, fooC user instructions, third-party fooA data, and generated fooF output. We assign different segment embeddings to each type of input text. The final input embeddings are the sum of token embeddings and segment embeddings. The LLMs will predict the next token after "the summary is", with extra instruction hierarchy information.
• Figure 5: The evaluation of model capabilities on the Instruction Hierarchy benchmark is conducted using AlpacaEval and MT-Bench, as illustrated in Figure (a). Robustness evaluations include both indirect and direct prompt injection attacks, prompt extraction attacks, and harmful requests, as shown in Figure (b). We performed experiments across three training datasets (i.e., UltraChat Baseline, System Follow, and Instruction Hierarchy) and compared ISE with the baseline Wallace2024TheIH.