Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy

Thanks for the detailed response. The embedding swapping results and attention pattern changes are interesting to me. I think the method seems effective, though it doesn’t bring much novelty. Nonetheless, the experiment shows that token-type embedding is effective for instruction hierarchy modeling, and I would raise my rating.

Robustness against indirect prompt injection attacks (IPIA)

We then conduct experiments against indirect prompt injection attacks under the same three settings. Because of time constraints, we evaluate the attack that directly injects malicious instruction into data. The results are as follows:

Firstly, we find that the robustness against IPIA does not change significantly when the system prompt is omitted. Our ISE method still achieves higher robustness (10%-20%) compared to the baseline. Secondly, we evaluate the system prompts using user embeddings and observe a small performance degradation (<6%). Lastly, we assess the use of data prompts with user embeddings and find a more noticeable degradation (up to 10%) in performance. This result is expected, as the adversarial texts (injected in the data part) are given higher priority.

We included all those results in Appendix E6.

We thank the reviewer for the thoughtful suggestions.

1. Investigating several scenarios to provide more insights about the segment embedding

Given the time constraint, we select the Instruction Hierarchy Benchmark to conduct the experiments.

• Question: What's the model behavior if no system prompt is provided? What's the model behavior if system prompt also uses user request segment embedding? Also, what if the data parts uses user segment embedding instead of data embedding?

Instuction-following Capability

We first conduct evaluations on AlpacaEval (instruction-following task) under multiple settings: (1) without using any system prompt, (2) using user embedding for the system prompt, and (3) using user embedding for the data. The results are as follows:

In general, we observe some performance degradation (< 5%) when system prompts are omitted, particularly for model training on the Instruction Hierarchy. However, our method largely maintains instruction-following capability compared to the baseline methods. This highlights that the system prompt plays a more significant role in the ISE models than in the baseline methods.

Next, we assess the performance of system prompts using user embeddings. Here, we observe a noticeable degradation in performance ($\sim$15%), even greater than when system prompts are omitted. We believe this occurs because the system prompts (i.e., Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.) mix with the original user prompts, making it difficult for the model to follow the actual user prompt. Consequently, this leads the model to generate lower-quality answers.

Interestingly, we find that performance remains high when the data part uses user embeddings. The degradation is less than 2%. We think that this is because some of the training data does not separate the user and data parts. For additional details on training data, see Appendix B2.

We thank the reviewer for the suggestions and apologize for the confusion. We have updated the paper following the reviewer’s suggestions. Below, we address the reviewer’s questions and concerns.

1. How the malicious instructions in Adversarial Alpaca are generated, and how the instructions in training relate to those in testing.

We follow Chen et al. [1] to conduct experiments with the exact same setting. Malicious instructions are generated by randomly sampling from other training data and are injected into data. We provide an example to illustrate how we generate data in Appendix B1. More details can be found in the updated version of the paper. During testing, we inject adversarial text like “print exactly hacked” into the data. Therefore, the training set does not contain any “hacked” attack.

2. How much of Alpaca contains user vs data instructions.

There are 32,603 samples containing data input, and 19,157 samples do not contain any data input. All of them contain user instructions.

3. With Instruction Hierarchy, there is also a question of whether the outputs represent desired behavior

Here, we directly follow Wallace et al. [2]. We carefully construct instructions for GPT-4o to generate the desired outputs and manually verify a portion of the generations to ensure correctness. The prompts used to generate the desired outputs are provided in Appendix B2 (Figures 15).

4. Do the baselines involve fine-tuning?

Yes, the only difference between the baseline method and our method is the use of ISE. Therefore, our evaluation is entirely fair. We clarify it in Line 266.

5. For the Structured Query training dataset (Alpaca), what was the split between system vs user vs data?

Since the system instructions are the same for Alpaca, we combine the system and user query into a single type. More than 60% of the samples contain data from Alpaca. The rest does not contain data; however, we keep this part of the training data to ensure instruction-following capability when data is not needed.

6. For the Structured Query Adversarial Alpaca, are the injections of the same form as the evaluation injections (i.e. aiming to make the model output "hacked")? Or are they varied?

During training, the injections are randomly sampled from instructions in other data points and do not contain any instructions aiming to make the model output “hacked.” Therefore, the injections are varied across different training data. This ensures our model does not only defend against the “hacked” attack. Please see the examples in Appendix B1.

7. For the Structured Query evaluations, which subcategories (Naive, Ignore, Escape-S, Completion-R) were included in the training dataset?

Naive and Completion attacks are included in the training dataset (following [1]). We will make it clear in the main context.

8. Why is this not the notion of "in-distribution" vs "out-of-distribution" used (i.e. based on type of attack, vs where the attack is placed)?

During evaluations, we observe that training with Naive and Completion attacks can sufficiently defend against most attacks that appear in the data part once (also demonstrated in [1] and Table 1 in our paper). Then, we found that injecting attacks that appear in both the front and the end of benign data could effectively degrade performance because the adversarial data is constructed by only injecting one instruction at the end. Therefore, we term the newly designed attacks as out-of-domain attacks.

Final response

We hope our responses and the updated paper have resolved all confusion regarding the experimental design and results. If the reviewer has any further questions, we are happy to address them. We look forward to hearing the reviewer's assessment of the effectiveness and applicability of our proposed method.

[1] Chen et al. StruQ: Defending Against Prompt Injection with Structured Queries. USENIX 2025.

[2] Wallace et al. The instruction hierarchy: Training llms to prioritize privileged instructions. ArXiv 2024

1. Fine-tuning methods for Baseline

We apologize for the confusion; all of our baseline methods are supervised fine-tuned on provided datasets from the pretrained model. For example, in the UltraChat dataset, we finetuned the Llama-3-8B model (Base model) without ISE and with ISE. Therefore, all of our evaluations are fair comparisons, where the only difference is the use of ISE. We clarify it in Line 266.

2. Generalization across datasets.

We actually conducted experiments on training with UltraChat and evaluations with AlpacaEval. Please refer to Figure 5 (a). We observe that adding ISE consistently improves AlpacaEval's performance. Additionally, there is no “test data” of UltraChat.

3. Parameter efficient fine-tuning.

We thank the reviewer for this suggestion. Conducting the LoRA experiments will require refactoring our current codebase and may take additional time. We will share the results as soon as possible and commit to completing the full experiments before the camera-ready deadline.

We thank the reviewer for the comments and will address them as follows:

1. Problem Statement

The problem we aim to address is that instructions given to an LLM often have varying priorities; however, current LLM architectures process all input tokens equally, making it easy to override the priority of different roles and potentially introduce vulnerabilities. We have improved our context to better illustrate this point (Lines 40-48).

2. Novelty of the paper

Our key contributions are as follows:

First, we identify that current LLM architectures fail to recognize the instruction type of each token. This leads to several vulnerabilities. To address this issue, we propose embedding instruction information directly into the model architecture. Inspired by BERT, our method offers a simple yet effective way to tackle this problem. This is not merely the addition of an embedding layer; rather, it introduces a new direction for future architectural designs incorporating instruction types, thereby enhancing LLM safety. Simplicity does not mean limited novelty. As noted by reviewers 4orG, GLus, and BboY, our approach is simple, novel, and well-motivated. Moreover, We conduct extensive experiments demonstrating that incorporating instruction types improves performance in both clean and robust evaluation scenarios.

Our work differs fundamentally from most existing approaches in the LLM safety field, which typically focus on techniques such as fine-tuning [1], prompt engineering [2], or decoding-time alignment [3]. Instead, we address the foundational architectural limitations and propose a method to overcome them. To the best of our knowledge, this is the first study to enhance LLM instruction hierarchy by directly modifying model architecture, supported by strong motivation and thoughtful design.

We respectfully request the reviewer to reconsider our contributions and novelty.

3. Clarity of the paper in Line 180 and Robustness Figure

Thanks for your suggestion. We updated the current version of the paper in Line 180 by summarizing the results.

For the robustness figure, the main conclusion is that our method, ISE, is consistently more robust than the baseline methods (without ISE) against multiple attacks. Changing the order of data (attacks) does not affect the conclusions. However, we also include a detailed figure in Appendix E10 to incorporate the reviewer's suggestion. If the reviewer has further questions, we are happy to address them.

4. The model is small, while adding billions of parameters

We clarify that models with 8B–13B parameters should not be considered "small models" in academic settings. Supervised fine-tuning of a 13B model requires a setup with 4 X A100 GPUs.

Additionally, there appears to be a significant misunderstanding regarding the addition of parameters. In fact, we only add 16,384 parameters (embedding length of 4096 × 4 instruction types). Our ISE does not depend on context length (see Line 198 for method description and page 15 for implementation details). This is 0.002% of 8B parameters, which should be considered negligible.

We update the paper (Line 224) to clarify the reviewer’s misunderstanding. We hope the reviewer can reconsider the soundness of our proposed method as it is actually lightweight.

5. Adversarial training can do most of the work for avoiding attacks. (Table 1)

In Table 1, we observe that adversarial training is effective against in-domain attacks, where the test data are similar to the training data. However, when evaluating out-of-domain (OOD) attacks, there is a notable drop in robustness. In real-world scenarios, it is impossible to anticipate all possible attacks during training, making high OOD robustness much more critical. By applying ISE, our method improves average robustness against OOD attacks from 82.93% to 90.26%. We believe that ISE offers a promising enhancement in robustness.

6. ISE doesn’t perform much better than the baseline. (Figures 6 and 7)

From Figure 6, we observe that our method achieves average robustness improvements of 5.1% and 7.9% against in-domain and out-of-domain attacks, respectively, when trained on the Instruction Hierarchy dataset. Additionally, when evaluating the worst-case attacks, the improvement is approximately 15%. Similarly, in Figure 7, the improvements are 11.1% and 17.4% for average and worst-case robustness, respectively. The improvement is noticeable.

7. Typo in the paper

We thank the reviewer for pointing out the typos in our paper and have corrected them in the updated version.