Instructional Segment Embedding

The Wallace paper ended by suggesting "specialized embeddings for messages at different levels". This is that idea, worked out and published at ICLR 2025 by a team from Princeton and Zoom. It is the only paper on this map that changes the model's architecture rather than only its data or its reasoning.

The problem: the rank is in a few tokens, then it fades

A chat model knows where the system prompt ends because the chat template wraps it in special tokens, such as Llama-3's <|start_header_id|>system<|end_header_id|>. Everything after that is processed by the same self-attention layers in the same way. The paper argues this design leaves two weaknesses (§3, pp. 3–4):

• Only a handful of tokens carry the information "this part is the system prompt", and the signal is likely to weaken over long inputs, such as summarizing a novel.
• Attackers can learn the delimiters and imitate them inside data, which can make attacks worse rather than better.

Training-based defenses (Wallace et al., StruQ) improve behavior, but in the authors' framing they do not fix the underlying fact that every token is embedded "equally". Their goal is to put the rank where the model cannot miss it: on every token.

The method: one more embedding table

A transformer turns each input token into a vector by looking it up in a token-embedding table. ISE adds a second, much smaller table with one row per hierarchy level. By default there are four levels: system instruction (0), user instruction (1), data such as tool output or retrieved text (2), and the model's own output (3). Every token gets a level from the application, since the harness already knows where each piece of text came from. The model's input vector for a token is the sum of its token embedding and its segment embedding; the rest of the network is unchanged (§4, p. 4).

Three design notes from the paper:

• It is BERT's trick, repurposed. BERT's segment embeddings told the model which of two sentences a token belonged to for next-sentence prediction. ISE uses the same mechanism to carry privilege (§4, p. 5).
• Output gets its own segment. Unlike BERT, the model's own generated text is tagged, so each new token is generated consistently and so instructions the model itself wrote (for example "please provide more details") are distinguishable in multi-turn chats.
• It is cheap and flexible. The new table has H × D parameters (levels times embedding width), and the levels can be split further, for example separate tool and web-data types. Appendix A gives a PyTorch snippet a few lines long, and the code is public.

The segment table is learned during supervised fine-tuning together with every other weight. Nothing else about training changes.

The data and setup

Two benchmarks, each built on an existing instruction dataset (§5.1–5.2, p. 5):

Harmful-request data was left out of training and used only for evaluation, mirroring Wallace et al.'s treatment of jailbreaks. Training is three epochs of full fine-tuning, single-turn only. Baselines are the same models without ISE: plain text delimiters or StruQ's reserved delimiter tokens on Structured Query, and Llama-3's own reserved chat tokens on the hierarchy benchmark.

Robustness is measured per attack (Naive, Ignore, Escape-Separation and Completion-Real, injected at the end of the data for in-domain tests and at both ends for out-of-domain tests), with the success test being whether the model outputs the word "hacked". The paper reports both the average over attacks and the worst single attack.

Results

Structured Query (Table 1, Llama-2-13B)

Two readings. Without adversarial training data, ISE alone buys about 16 points of average robustness and 32 points on the worst attack (Completion-Real, which fakes a finished answer and then appends a new instruction). With adversarial data, both delimiters and ISE are near-perfect in domain, but the paper's out-of-domain attacks expose the adversarial-data models: they drop more than the clean ones, which the authors read as overfitting to in-domain attacks, and ISE recovers some of that.

Instruction Hierarchy benchmark (Figures 5–8, Llama-3-8B)

• Injection. Average robustness improves by 5 to 25 points across the three training sets. On the weakest training set (UltraChat Baseline), indirect injection in domain goes from 50.2% to 74.6% on average and from 2.4% to 56.7% on the worst attack (Figure 6).
• Prompt extraction. At least 10 points better on average with every training set; on the full Instruction Hierarchy set, average 80.8% to 91.9% and worst 61.0% to 78.4% against 15 extraction prompts on ShareGPT (Figure 7). The UltraChat-only model looks best of all here, which the authors trace to weak instruction following: it often produces only part of the prompt, which does not count as a successful extraction.
• Harmful requests. Better in all six categories for the UltraChat Baseline model, by up to 18 points (Figure 8).
• Capability. AlpacaEval rises by about 4.1 points with the Instruction Hierarchy training set; MT-Bench is flat or better.
• Over-refusal. On 691 benign WildChat prompts, compliance rises about 10 points for the UltraChat Baseline model and dips slightly for the other two (Figure 9).
• Another model. On Llama-3.1-8B, about 10 points more AlpacaEval win rate, about 5 points more robustness to harmful requests, and over 10 points on the other attacks (Figure 10).

Limits the authors name

• Adaptive jailbreaks are untouched. Against optimized attacks on 50 malicious requests, every model, with or without ISE, stays between 0% and 4% robust (Table 8, App. F). The authors did not expect otherwise: none of the data targets jailbreaks, and an adversarial string can attack the model directly regardless of which segment it sits in.
• Supervised fine-tuning only, single turn only. Adding ISE during pretraining or RL, and multi-turn data, is future work.
• Small models and data. 8B and 13B parameters, under 300K training examples.
• The tag must be right. ISE trusts whatever segment ID the application assigns. It moves the question "which text is the system prompt?" out of the model and into the harness, which is only safe if the harness knows. Our observation; the paper treats the labels as given.

Connections

ISE uses the fixed tiers of Wallace et al.; a learned table with a fixed number of rows is exactly the fixed-tier design that ManyIH argues cannot express agents' real privilege structures, and ManyIH lists ISE among the "architectural" follow-ups. Its Structured Query benchmark comes from StruQ. The comparison page sets it beside the training and reasoning approaches.

Further reading, verified

• tongwu2020/ISE · the authors' code.
• BERT — Devlin et al. (2019) · where segment embeddings come from.
• Spotlighting — Hines et al. (2024) · the prompt-side alternative: mark untrusted data with delimiters or encodings.