Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy
Tong Wu, Shujian Zhang, Kaiqiang Song, Silei Xu, Sanqiang Zhao, Ravi Agrawal, Sathish Reddy Indurthi, Chong Xiang, Prateek Mittal, Wenxuan Zhou
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 the instruction-following capability of up to 4.1% on AlpacaEval. Overall, our approach offers a promising direction for enhancing the safety and effectiveness of LLM architectures.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 6c87dd0d-af5c-4232-b5e4-06d937e39db9Cited by top-tier papers21
- PromptLocate: Localizing Prompt Injection AttacksYuqi Jia, Yupei Liu, Zedian Shao, Jinyuan Jia et al.S&P 2026 · 35 citations
- ASIDE: Architectural Separation of Instructions and Data in Language ModelsEgor Zverev, Evgenii Kortukov, Alexander Panfilov, Alexandra Volkova et al.ICLR 2026 · 28 citations
- Adaptive Attacks on Trusted Monitors Subvert AI Control ProtocolsMikhail Terekhov, Alexander Panfilov, Daniil Dzenhaliou, Caglar Gulcehre et al.ICLR 2026 · 26 citations
- ObliInjection: Order-Oblivious Prompt Injection Attack to LLM Agents with Multi-source DataReachal Wang, Yuqi Jia, Neil Zhenqiang GongNDSS 2026 · 24 citations
- MUZZLE: Adaptive Agentic Red-Teaming of Web Agents Against Indirect Prompt Injection AttacksGeorgios Syros, Evan Rose, Brian Grinstead, Christoph Kerschbaumer et al.USENIX Security 2026 · 18 citations
Builds on16
- ALBERT: A Lite BERT for Self-supervised Learning of Language RepresentationsZhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel et al.ICLR 2020 · 7,418 citations
- WebShop: Towards Scalable Real-World Web Interaction with Grounded Language AgentsShunyu Yao, Howard Chen, John Yang, Karthik NarasimhanNeurIPS 2022 · 1,477 citations
- WildChat: 1M ChatGPT Interaction Logs in the WildWenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie et al.ICLR 2024 · 504 citations
- LinkBERT: Pretraining Language Models with Document LinksMichihiro Yasunaga, Jure Leskovec, Percy LiangACL 2022 · 463 citations
- Improving Alignment and Robustness with Circuit BreakersAndy Zou, Long Phan, Justin Wang, Derek Duenas et al.NeurIPS 2024 · 362 citations
Related papers
- Evaluating the Instruction-Following Robustness of Large Language Models to Prompt InjectionZekun Li, Baolin Peng, Pengcheng He, Xifeng YanEMNLP 2024 · 15 citations
- StruQ: Defending Against Prompt Injection with Structured QueriesSizhe Chen, Julien Piet, Chawin Sitawarin, David A. WagnerUSENIX Security 2025
- Can LLMs Separate Instructions From Data? And What Do We Even Mean By That?Egor Zverev, Sahar Abdelnabi, Soroush Tabesh, Mario Fritz et al.ICLR 2025
- DRIP: Defending Prompt Injection via Token-wise Representation Editing and Residual FusionRuofan Liu, Yun Lin, Zhiyong Huang, Jin Song DongCCS 2026 · 3 citations
- Control Illusion: The Failure of Instruction Hierarchies in Large Language ModelsYilin Geng, Haonan Li, Honglin Mu, Xudong Han et al.AAAI 2026 · 21 citations
