Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy

Wu, Tong, Zhang, Shujian, Song, Kaiqiang, Xu, Silei, Zhao, Sanqiang, Agrawal, Ravi, Indurthi, Sathish Reddy, Xiang, Chong, Mittal, Prateek, Zhou, Wenxuan

arXiv.org Artificial Intelligence 

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 instructionfollowing 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. Large Language Models (LLMs) have shown significant potential in enabling sophisticated agentic applications and facilitating autonomous decision-making across various domains, such as web agents, educational tools, medical assistance, and more (Yao et al., 2022; Gan et al., 2023; Abbasian et al., 2024). To optimize the use of AI applications, a structured approach to implementation is widely adopted. This involves clear distinctions among system instructions, user prompts, and data inputs, as illustrated in Figure 1. These instructions contain specific priorities that help the model execute functionalities correctly and better assist users. Modern LLMs process text without formal mechanisms to differentiate and prioritize instructions. Consequently, Figure 1: A demonstration of the hierarchy malicious attackers can easily exploit this limitation to of instructions, including system override priority roles, leading to various vulnerabilities. Prompt extraction (Zhang et al., 2024) aim to extract system messages, revealing proprietary prompts.