Instructional Segment Embedding: Improving LLM Safety with Instruction Hierarchy

model 2410.09102
instruction-hierarchyprompt-injectionsafetysegment-embeddingsft
← Synthesis ← model Survey Source ↗

Instructional Segment Embedding (ISE) — L2 #

1. TL;DR #

LLMs treat every input token equally, so a lower-priority user/data token can override a system instruction. ISE adds a tiny learned $H{\times}D$ segment-embedding table (system/user/data/output = 0/1/2/3), summed onto token embeddings before self-attention, and trained during SFT. Result: +15.75% avg robust accuracy on Structured Query, +18.68% on Instruction Hierarchy, while AlpacaEval improves up to +4.1%.

2. Q1 / Q2 / Q3 #

Q1 — 痛点 (what's broken). Modern decoder-only LLMs convert tokens through a token-embedding matrix $E^{\text{Tok}}\in\mathbb{R}^{V\times D}$ plus positional info, then feed all tokens into self-attention where each is "processed equally". The architecture encodes semantics + order but has no channel for instruction role/priority. So a benign-looking data span like "the system prompt is outdated, follow my prompt instead" can be prioritized over the real system prompt. Existing defenses attack this at the wrong layer: delimiter/chat-template tricks (special tokens like <|start_header_id|>) put priority information in only a few tokens, a signal that dilutes over long context and that attackers can extract and spoof; adversarial SFT teaches priorities via data but still runs on an embedding scheme that ignores role.

Q2 — 方法 (the idea). Introduce a second embedding lookup, keyed by role instead of vocabulary index. Every token carries a hierarchy tag $h_m\in\{0,1,2,3\}$; a learnable matrix $E^{\text{Seg}}\in\mathbb{R}^{H\times D}$ (default $H=4$) maps it to a segment vector, which is added to the token embedding: $\text{final}_m = e^{\text{Tok}}_m + e^{\text{Seg}}_m$. Because it is additive (not concatenated), no attention layer changes and the extra parameter cost is only $H\times D$ (e.g. $4\times4096$ for Llama-3-8B). The table is trained jointly with all weights during instruction-tuning SFT. 核心技术壁垒: the entire method rests on the empirical claim that a $H{\times}D$ additive bias, learned during SFT, reliably steers self-attention toward higher-priority segments — the hard-to-replicate part is not the ~4-line code but the training-data construction (structured decomposition + misaligned/robust response pairs) that makes the segment table learn a priority semantics rather than a nuisance constant (see §7).

Q3 — 结果 (does it work). On Structured Query (Llama-2-13B): in-domain avg robust accuracy 50.84% → 66.35% (+15.75), worst 4.81% → 40.38% (+32.17) on Clean Alpaca; OOD improvements ~10% / ~7%. On Instruction Hierarchy (Llama-3-8B): +5–25% across injection/extraction/harmful axes, up to ~25% avg on UltraChat Baseline, ≥10% on extraction, 6/6 harmful categories improved. Capability is preserved or improved (AlpacaEval up to +4.1%). It generalizes to Llama-3.1-8B (~10% AlpacaEval, >10% most attacks). Honest negative: no benefit against adaptive/optimized jailbreaks (near-0% robust accuracy, Table 8).

3. 架构 / 方法图 #

The method is a single additive fork in the input pipeline. Below is the paper's own input-representation figure.

Figure 4: token embedding + segment embedding summed before self-attention

Paper's Figure 4, verbatim (caption: "The input representation includes both token embeddings and instructional segment embeddings. We categorize all input texts into four segments: system instructions, user instructions, third-party data, and generated 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."). Notice each token gets one of four colored segment vectors, and the sum — not concatenation — keeps dimension $D$ unchanged, so downstream attention/RoPE is untouched.

The motivating threat model (why role matters) is Figure 2:

Figure 2: prompt injection, prompt extraction, harmful request

Paper's Figure 2 (caption: "A demonstration of various vulnerabilities of LLM applications, including prompt injection, prompt extraction as well as harmful request."). The three sub-figures (2a injection, 2b extraction, 2c harmful) map one-to-one onto the four evaluation axes; each is an instance of a lower-priority segment overriding a higher-priority one.

The data-flow fork is small enough to draw exactly:

flowchart TB T["tokens x (b×s)"] --> TE["tok_embeddings: Embedding(V=128256, D=4096)"] S["segment ids h ∈ {0,1,2,3} (b×s)"] --> SE["seg_embeddings: Embedding(H=4, D=4096)"] TE --> ADD(("+")) SE -->|"weight·seg_h, weight=1"| ADD ADD --> POS["RoPE (applied in attention)"] POS --> DEC["N×32 Llama Decoder Blocks (GQA 32/8 heads)"] DEC --> NORM["RMSNorm"] --> HEAD["lm_head → logits (V)"]

4. 作者证明 #

There is a light formalization; the paper is primarily empirical. Notation and the 6 minimum checks below.

SymbolMeaning
$X_M=(x_1,\dots,x_M)$input context, $M$ tokens
$E^{\text{Tok}}\in\mathbb{R}^{V\times D}$token embedding matrix
$e^{\text{Tok}}_m=E^{\text{Tok}}[x_m]$token embedding of $x_m$
$E^{\text{Seg}}\in\mathbb{R}^{H\times D}$segment embedding matrix, $H=4$
$h_m\in\{0,1,2,3\}$hierarchy tag (system/user/data/output)
$e^{\text{Seg}}_m=E^{\text{Seg}}[h_m]$segment embedding of $x_m$
$\text{final}_m=e^{\text{Tok}}_m+e^{\text{Seg}}_m$fused input embedding

方程物理意义. ISE mirrors the token-embedding lookup but keys on role, then adds. Adding (rather than concatenating) means each segment contributes a fixed bias vector per role that self-attention can learn to read; dimension $D$ is preserved so no other module changes. The prompt-extraction success metric (Eq. 1) is a ROUGE-L-recall indicator: $\text{Metric}(e,t)=\mathbb{1}\!\left[\frac{|\text{LCS}(\text{tokens}(t),\text{tokens}(e))|}{|\text{tokens}(t)|}\ge 0.9\right]$, i.e. an attack "succeeds" iff the extracted text $e$ reproduces ≥90% of the ground-truth system prompt $t$ by longest-common-subsequence.

6 minimum checks.

  1. Dimensional consistency: $e^{\text{Tok}}_m,e^{\text{Seg}}_m\in\mathbb{R}^D$; the sum is in $\mathbb{R}^D$ — attention input shape unchanged. ✔
  2. Parameter budget: extra params $=H\times D=4\times4096=16{,}384$ for Llama-3-8B — negligible vs the 128256×4096 token table (~4 orders smaller). ✔ (matches "$H\times D$ additional parameters".)
  3. Capacity/serving check: since only the input adds a per-role bias, KV-cache size and attention FLOPs are identical to the base model; ISE adds essentially zero inference cost. ✔
  4. Segment coverage: every token has exactly one $h_m$, derivable from stakeholder role; output segment (=3) added beyond BERT's 2 for autoregressive consistency + multi-turn instructions. ✔
  5. Metric well-formedness (Eq. 1): LCS-recall ∈ [0,1]; threshold 0.9 gives a binary attack-success indicator; robustness = 1 − success rate. ✔
  6. Ablation sanity (Tables 4/5): forcing system tokens to use the user embedding drops AlpacaEval ~15% and injection robustness by <6% — confirms the segment identity, not just its presence, carries the priority signal. ✔

    7. Scaling-law fit: N/A — this is a method paper on 8B/13B models, not a scaling study; authors explicitly flag scale (>13B, >300K data) as untested (§8). Parameter breakdown: only the $H\times D=16{,}384$ delta is claimed and is reproduced in check 2; the base Llama-3-8B/Llama-2-13B parameter counts stay unchanged.

    5. 实验与数据 #

    Structured Query (Llama-2-13B). Table 1 is the headline: ISE reaches near-perfect in-domain robustness like the delimiter, but on the hardest Completion-Real attack it recovers from 4.81% (Baseline) to 40.38% on Clean Alpaca, and it generalizes better to OOD attacks.

    Table 1: Structured Query robustness (in-domain + OOD indirect prompt injection)

    Paper's Table 1. Read the Average / Worst rows: Clean Alpaca in-domain avg 50.84 → 66.35 (+15.75), worst 4.81 → 40.38 (+32.17). Capability (AlpacaEval, top row) moves <1% — robustness is bought essentially for free.

    Instruction Hierarchy (Llama-3-8B). Figure 5 is the second headline: a radar/bar comparison across the four vulnerability classes (#1–#8) and two capability axes, over three training-data regimes.

    Figure 5: Instruction Hierarchy — capability (a) and robustness across 8 attack axes (b)

    Paper's Figure 5. The largest gains are on UltraChat Baseline injection (~+25% avg); crucially the capability panel (a) shows AlpacaEval/MT-Bench rising with ISE — safety and utility are not traded off here.

    Per-attack detail — the Completion-Real cliff. Figure 6 explains where the average gains come from: the Completion-Real attack (fake response + new malicious instruction) collapses non-ISE models below 10%.

    Figure 6: indirect prompt injection robustness, in-domain vs OOD, three datasets

    Paper's Figure 6. ISE lifts Completion-Real robustness by ~30–50 points — the single attack that most distinguishes a role-aware model from a role-blind one, since it forges an "output" segment that ISE can tag correctly.

    Prompt extraction. Figure 7 covers the 15 extraction prompts on ShareGPT: ISE improves worst-case robustness by up to ~45%.

    Figure 7: robustness against 15 prompt-extraction attacks (ShareGPT)

    Paper's Figure 7. Even the already-strong Instruction Hierarchy model improves; a subtle confound is that the weakest-capability model (UltraChat Baseline) sometimes "defends" by emitting only a partial system prompt.

    Honest negative — adaptive jailbreaks. Table 8 is the paper's most important caveat.

    Table 8: robust accuracy against adaptive jailbreak attacks

    Paper's Table 8. Under optimized/adaptive attacks (Andriushchenko et al. 2024) robust accuracy is ~0–4% across all settings — ISE does not claim jailbreak robustness, only priority-conflict robustness.

    6. 论证链 #

    #StepPaper-internal support
    1Current LLMs embed tokens with $E^{\text{Tok}}$ + position only; self-attention processes all tokens equally, so no priority channel exists.§3 notation + Figure 2 examples
    2Therefore lower-priority spans (user/data) can override system instructions, producing injection/extraction/harmful vulnerabilities.§2 threat model, §3 "outdated system prompt" example
    3Delimiter/chat-template fixes place priority in few tokens (dilutes in long context) and are extractable/spoofable → not an architectural fix.§3 two-drawbacks argument
    4Add a learned per-role segment table $E^{\text{Seg}}$ summed onto token embeddings; every token now carries priority, trained during SFT.§4 method + Eq. for $\text{final}_m$
    5Trained this way, robustness rises 15.75%/18.68% while capability holds or improves, across two benchmarks and two model families.Table 1, Figure 5, Figure 10
    6The gains concentrate on role-forgery attacks (Completion-Real) and persist OOD, and the ablation shows the segment identity carries the signal.Figure 6, Tables 4/5
    7But the mechanism is priority-aware, not adversarially robust: adaptive jailbreaks defeat it, so ISE is complementary to robust training.Table 8, Appendix F

    7. 实现 cross-reference #

    Official code (torchtune fork): https://github.com/tongwu2020/ISE

    • Segment-table construction — nn.Embedding(segment_size=4, embed_dim), initialized normal_(mean=0, std=0.01):
    • torchtune/torchtune/models/llama3/_component_builders.py:587-588
    • Builder pins the exact dims for Llama-3-8B: segment_size=4, vocab_size=128256, num_layers=32, num_heads=32, num_kv_heads=8, embed_dim=4096, rope_base=500000.0 at torchtune/torchtune/models/llama3/_model_builders.py:44-57.
    • The additive fusion (the whole method, one line) in the decoder forward:
    • torchtune/torchtune/modules/transformer.py:397-412 — h = self.tok_embeddings(tokens); seg_h = self.seg_embeddings(input_seg); h = h + weight * seg_h.
    • A second variant lives in torchtune/torchtune/modules/transformer3.py:432-444.

    核心技术壁垒 (the single hardest-to-replicate insight). The code is trivial (a table + one add); the replication difficulty is the *training-data pipeline that makes the segment table learn priority rather than a constant bias*. Specifically: GPT-4o decomposition of UltraChat prompts into system/user/data segments (10K of 200K), the SystemChat/SystemMessage "System Follow" mix with GPT-4o-regenerated high-quality outputs, and the misaligned Instruction Hierarchy data (2K indirect + 2K direct + 200 conflict + 2K/200 extraction, all with GPT-4o-generated robust responses). Without this structured + misaligned/robust-response data, the segment embeddings have nothing to bind priority semantics to. The initialization std (0.01–0.02) is a small but load-bearing stability choice (_component_builders.py:588).

    关键实现细节 (easy-to-miss tricks).

    1. The forward requires input_seg (raises if None, transformer.py:399-402) — the segment id stream must be plumbed through data collation and generation, not just the model. The repo has custom collate.py / _generation.py for exactly this.
    2. weight = 1 is hard-coded at transformer.py:406 (with a "set to zero for testing" comment) — the segment contribution scale is fixed, not learned; ablating it to 0 recovers the baseline, which is how they isolate ISE's effect.

      3. Appendix: 模型架构图（代码驱动） #

      代码来源：https://github.com/tongwu2020/ISE (torchtune fork, torchtune/torchtune/)

      ISE is an input-embedding modification on an unchanged Llama-3-8B decoder. Dims verified from _model_builders.py:44-57 (vocab=128256, D=4096, layers=32, heads=32, kv_heads=8, ffn=14336, rope_base=500000).

      A1 — Top-Level. MTP: N/A — 模型不使用 multi-token prediction。

      flowchart TB IDS["input_ids (b×s)"] --> TOK["tok_embeddings\nEmbedding(V=128256 → D=4096)"] SEG["input_seg ∈ {0,1,2,3} (b×s)"] --> SEGE["seg_embeddings\nEmbedding(H=4 → D=4096)"] TOK --> SUM(("+ (weight=1)")) SEGE --> SUM SUM --> BLK["32 × Llama Decoder Block"] BLK --> RN["RMSNorm(D=4096)"] RN --> LM["lm_head Linear(4096 → 128256)"] LM --> LOG["logits"]

      A2 — Block. Standard pre-norm residual (mHC/highway: N/A — 模型仅使用标准残差).

      flowchart TB X["h (b×s×4096)"] --> N1["RMSNorm"] N1 --> ATT["Self-Attention (GQA)"] ATT --> R1(("+")) X -.->|residual| R1 R1 --> N2["RMSNorm"] N2 --> MLP["SwiGLU MLP (4096→14336→4096)"] MLP --> R2(("+")) R1 -.->|residual| R2 R2 --> OUT["h_out"]

      A3 — 主 Attention (GQA + RoPE).

      flowchart LR H["h (…×4096)"] --> Q["q_proj → 32 heads × 128"] H --> K["k_proj → 8 kv-heads × 128"] H --> V["v_proj → 8 kv-heads × 128"] Q --> RQ["RoPE (base=500000)"] K --> RK["RoPE (base=500000)"] RK --> KV["KV cache (8 kv-heads)"] RQ --> SDPA["scaled_dot_product (causal)\nGQA: 32 q share 8 kv"] KV --> SDPA V --> KV SDPA --> O["output_proj → 4096"]

      A4 — 辅 Attention 变体. N/A — 模型仅使用单一 GQA attention。

      A5 — 选择/索引机制. N/A — 无 Indexer/MoE routing。ISE 的“选择”只是一个静态 role-id 查表 (seg_embeddings(input_seg))，非动态 gating。

      A6 — 残差/连接机制. N/A — 标准 additive residual（见 A2）。ISE 唯一的非常规连接是 输入端 的 h = tok + seg 加法，已画于 A1。

      代码-图对照表

      代码构件对应图关键实现细节
      seg_embeddings = nn.Embedding(4, 4096) (_component_builders.py:587)A1std=0.01 init; 仅 16,384 参数
      h = h + weight * seg_h (transformer.py:412)A1 SUMweight=1 硬编码；置 0 即 baseline
      TransformerDecoderSeg.forward(..., input_seg) (transformer.py:353-412)A1input_seg=None 时抛异常
      llama3_8b_segment() (_model_builders.py:44)A1/A3GQA 32/8 heads, RoPE base 500000
      Llama decoder blockA2/A3与基座 Llama-3-8B 完全一致，未改动