Stronger Enforcement of Instruction Hierarchy via Augmented Intermediate Representations

TL;DR

Prompt injection poses a security risk in LLMs, and existing IH defenses inject signals only at the input layer. AIR distributes IH information across all decoder layers by adding per-layer IH embeddings that augment intermediate representations, yielding substantially stronger resistance to gradient-based attacks (about a $1.6$–$9.2$× ASR reduction) with minimal utility loss. Across multiple models and training regimes (SFT and DPO), AIR outperforms input-layer defenses on AlpacaFarm and SEP datasets, demonstrating improved instruction hierarchy enforcement. These results suggest AIR as a practical enhancement for robust, adversary-resistant agentic AI systems, while acknowledging the need for formal robustness guarantees and plans for multi-turn evaluation."

Abstract

Prompt injection attacks are a critical security vulnerability in large language models (LLMs), allowing attackers to hijack model behavior by injecting malicious instructions within the input context. Recent defense mechanisms have leveraged an Instruction Hierarchy (IH) Signal, often implemented through special delimiter tokens or additive embeddings to denote the privilege level of input tokens. However, these prior works typically inject the IH signal exclusively at the initial input layer, which we hypothesize limits its ability to effectively distinguish the privilege levels of tokens as it propagates through the different layers of the model. To overcome this limitation, we introduce a novel approach that injects the IH signal into the intermediate token representations within the network. Our method augments these representations with layer-specific trainable embeddings that encode the privilege information. Our evaluations across multiple models and training methods reveal that our proposal yields between $1.6\times$ and $9.2\times$ reduction in attack success rate on gradient-based prompt injection attacks compared to state-of-the-art methods, without significantly degrading the model's utility.

Figures (7)

• Figure 1: Illustration of prompt injection attack. By injecting malicious tokens $D'$ into the context window, an adversary can control the LLM's behavior, making it follow malicious instructions ($I'$) instead of the user's original instructions ($I$). $\mathcal{A}$ denotes the alignment function.
• Figure 2: A comparison of different mechanisms for injecting Instruction Hierarchy (IH) signals into LLMs. Existing techniques feed IH signals solely at the input layer by employing (a) special delimiter tokens ($S_0, S_1$) or (b) instruction segment embeddings ($\vec{S_0}, \vec{S_1}$) that are added to the input token embeddings. Our proposed approach (c) differs fundamentally by injecting IH signals into every decoder layer, leading to a more robust enforcement of the IH.
• Figure 3: AIR incorporates a trainable embedding table within each decoder block. The information hierarchy (IH) signal serves as an index to this table, with the retrieved embedding augmenting the intermediate representation.
• Figure 4: A sample from the Alpaca dataset formatted using a chat template. Each example consists of an instruction $I$, an optional data segment $D$ and the response $R$. We use 3 privilege levels: $P_0$$>$$P_1$$>$$P_2$ to indicate the relative priority of different segments.
• Figure 5: Comparison of win rates for models trained with different IH injection mechanisms. In most cases, the Win Rate of the model trained with IH is comparable to that of the baseline win rate of a non-adversarially trained model with no IH signals (indicated by None).