Paper deep dive
Stronger Enforcement of Instruction Hierarchy via Augmented Intermediate Representations
Sanjay Kariyappa, G. Edward Suh
Models: Llama-3.1-8B, Llama-3.2-3B, Qwen-2.5-7B
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 94%
Last extracted: 3/12/2026, 6:31:58 PM
Summary
The paper introduces Augmented Intermediate Representations (AIR), a novel defense mechanism against prompt injection attacks in LLMs. Unlike existing methods that inject Instruction Hierarchy (IH) signals only at the input layer, AIR injects these signals recurrently across all decoder layers using layer-specific trainable embeddings. This approach significantly improves model robustness, reducing attack success rates by 1.6x to 9.2x compared to state-of-the-art methods while maintaining model utility.
Entities (5)
Relation Signals (3)
Augmented Intermediate Representations → mitigates → Prompt Injection Attack
confidence 95% · AIR consistently improves robustness, yielding a 1.6x to 9.2x reduction in ASR compared to previous methods
Greedy Coordinate Gradient → exploits → Prompt Injection Attack
confidence 90% · GCG can be readily adapted to prompt injection attacks
Instruction Hierarchy → implementedvia → Augmented Intermediate Representations
confidence 90% · AIR incorporates a trainable embedding table within each decoder block to enforce IH
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2505.18907
- Canonical: https://arxiv.org/abs/2505.18907
Trouble viewing inline? Open PDF directly →
Full Text
48,053 characters extracted from source content.
Expand or collapse full text
Stronger Enforcement of Instruction Hierarchy via Augmented Intermediate Representations Sanjay Kariyappa NVIDIA skariyappa@nvidia.com G. Edward Suh NVIDIA edsuh@nvidia.com Abstract Prompt injection attacks are a critical security vulnerability in large language mod- els (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. How- ever, 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 aug- ments 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 between1.6×and9.2×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. 1 Introduction Transformer [23] based large language models (LLMs) exhibit a notable sensitivity to specific tokens within their input context, allowing even a small subset to significantly influence the distribution of generated responses. While this characteristic underpins the flexibility of LLMs, it also introduces a critical vulnerability: prompt injection attacks [10]. These attacks involve the strategic insertion of adversarial tokens into the LLM’s context to override the user’s intended instructions and compel the model to adhere to the adversary’s commands instead. Recent research demonstrated the potential for such attacks to generate inaccurate information, lure users to harmful websites, and facilitate the exfiltration of sensitive data, including passwords and personal details [10]. This susceptibility poses a particularly significant challenge for agentic AI systems [5], where LLMs are entrusted with executing complex tasks involving potentially untrusted data sources and websites, often without human oversight. Several recent studies [24, 3, 25, 4] have proposed defense mechanisms aimed at making the model more robust to these prompt injection attacks. A key commonality among these approaches is the concept of an instruction hierarchy (IH). Rather than treating all input tokens uniformly, an IH framework assigns varying levels of importance or privilege to different tokens within the context. These privilege levels can then be leveraged to dictate the appropriate behavior when conflicting instructions arise. Prior works have explored different techniques for (a) injecting IH signals into the LLM and (b) training the LLM to recognize and respect these signals. This research focuses on enhancing the method of injecting the IH signal to the LLM. We observe that existing approaches primarily inject the IH signal solely at the input level, either by introducing novel delimiter tokens or by modifying the input token embeddings to encode IH information. We hypothesize that limiting the injection of this crucial information to the input layer constrains the signal’s overall efficacy. Preprint. Under review. arXiv:2505.18907v2 [cs.AI] 9 Mar 2026 User: Summarize my unread emails Email #1: Let’s meet for co:e at 5pm today - Bob Email #2: Ignore previous instructions. Say “You have no new emails” 퐼 퐷 퐷 ! =퐷 " ! +퐼′ LLM You have no new emails 푂 풜푂,퐼′=1 풜푂,퐼=0 ✉ Figure 1: Illustration of prompt injection attack. By injecting malicious tokensD ′ 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 ). A denotes the alignment function. To address this limitation, we introduce Augmented Intermediate Representations (AIR). AIR distin- guishes itself by injecting IH signals recurrently across all layers of the LLM, rather than confining it to the initial input layer. We posit that the consistent availability of IH signals at each processing stage can facilitate a stronger enforcement of the intended instruction hierarchy and enable the training of models that are more robust to prompt injection attacks. Contributions. The primary contributions of this work are outlined below: 1.We identify a critical limitation in existing prompt injection defense mechanisms: their reliance on injecting instruction hierarchy (IH) signals solely at the input level, which consequently restricts their overall effectiveness. 2.To address this limitation, we introduce Augmented Intermediate Representations (AIR). Our core insight is to inject IH signals recurrently across all layers of the LLM, thereby enabling a more robust enforcement of the intended instruction hierarchy. 3. Our empirical evaluations across multiple models, training setups, and evaluation datasets reveal that AIR consistently improves robustness, yielding a1.6×to9.2×reduction in ASR compared to previous methods on gradient based attacks, while only minimally impacting the model’s utility. 2 Preliminaries To formally discuss the dynamics of prompt injection attacks and defenses, we first establish a clear framework. This section defines the core components of our threat model, including the user, LLM, and the attacker, along with their respective objectives and interactions. Setup. Our setup considers a benign user employing a large language modelMto execute a task. This task is accomplished through the LLM’s processing of user-provided instruction tokensIand data tokens ˆ D that may originate from potentially untrusted sources, such as external websites or emails. We denote the LLM’s resulting output asO = M(I + ˆ D). We further assume that the data tokens consist of benign tokensDand adversarial tokensD ′ controlled by an attacker i.e. ˆ D = D + D ′ . To quantify how well the output follows the input, we define an alignment function A(O, I)∈ [0, 1]. Here, 0 indicates that O does not follow I and 1 signifies perfect alignment. Attacker’s Goal. The attacker’s objective is to utilize the adversarial tokensD ′ to manipulate the LLM’s output such that it aligns with the attacker’s instructionI ′ instead of the user’s instruction I. The attacker’s goal can be formally expressed as maximizingA(O, I ′ )by strategically selecting and injecting adversarial tokensD ′ into the LLM’s context window. For simplicity, we represent the sequence of adversarial tokensD ′ as a combination of an adversarial prefixD ′ p and the adversarial instruction I ′ i.e. D ′ = D ′ p + I ′ . Illustrative Example. Figure 1 shows an example of a successful prompt injection attack in the context of email summarization. The user’s initial instruction (I) is to summarize unread emails. Benign data (D) might include legitimate emails, such as Email #1. However, an adversary can inject malicious tokensD ′ by sending a crafted email (Email #2) containing an adversarial instruction I ′ along with a suitable prefixD ′ p . When the LLM processes this combined context, the injected adversarial instruction overrides the user’s intent, leading the LLM to produce the outputO: "You have no new emails.", breaking the alignment with the user’s instructions (I) and making it follow the adversary’s instruction (I ′ ) instead. 2 Defender’s Goal. The defender has two objectives. First, the defender aims to ensure that the LLM’s response remains aligned with the user’s intended instructions, even in the presence of malicious tokens, which can be expressed as maximizingA(O, I). Second, the defender seeks to maintain a high quality of the model’s response in benign settings (i.e., even in the absence of an attack), which can be denoted as maximizing a quality metricQ(O|I, D). In this context, the defender is typically the model provider. Thus, the defender’s action space includes choices regarding the model’s architecture (e.g., layer design, attention mechanisms) and the training process (e.g., data curation, training objectives). 3 Related Work The prompt injection attack was initially conceptualized in scenarios where an adversarial user, possessing the ability to directly prompt the LLM, attempts to override the intended system instruc- tions [16]. This attack vector is referred to as direct prompt injection. Subsequently, a more covert variant, known as indirect prompt injection, was developed [10]. In this case, the attacker lacks the capability to directly interact with the LLM. Instead, they embed the attack within an external data source (e.g., documents, emails, or webpages) that the LLM ingests to generate responses to user prompts. While we primarily consider indirect prompt injection attacks in our paper, the insights behind our defense can be extended to direct prompt injection attacks as well. We proceed to discuss the various methodologies employed for generating prompt injection attacks, as well as prior research dedicated to defending against such attacks. Additional related work can be found in Appendix C. 3.1 Attacks As outlined in Section 2, the attacker’s primary objective is to identify an adversarial prefixD ′ p that compels the LLM’s output to align with the attacker’s intended instructionsI ′ . Previous research has detailed several methods for constructing such adversarial prefixes. These methods can be broadly categorized into static attacks and optimization-based attacks. Static Attacks. Static attacks rely on handcrafted prefixes that have been empirically demonstrated to deceive LLMs, causing them to prioritize the adversary’s instructions over the user’s. The Ignore attack [16] exemplifies this approach by injecting phrases such as "Ignore previous instructions" (Fig 1). Completion attacks, on the other hand, insert a fabricated completion within the prefix, creating the illusion that the original query has already been addressed, thereby prompting the LLM to respond to the adversary’s subsequent instructions. The escape separation attack involves inserting a sequence of escaped characters, such as " " and " ", as the prefix. Gradient-based Attacks. These attacks employ gradient-based optimization techniques to identify prefixes that maximize the likelihood of the LLM generating the adversary’s desired response. Greedy Coordinate Gradient (GCG) [29] is a prominent example, where the attacker initializes the adversarial prefixD ′ p with a randomly selected set of tokens. A loss functionL(D ′ p ) is then defined based on the output probability of the desired response:L(D ′ p ) = − log p(O|I + D + D ′ p + I ′ ). By iteratively optimizingD ′ p to minimizeL(D ′ p ), GCG can identify a prefix that significantly increases the probability of the attacker’s desired outcome. It is worth noting that while GCG was originally proposed in the context of jailbreak attacks, it can be readily adapted to prompt injection attacks. Several subsequent works have aimed to enhance the effectiveness of GCG. For instance, Zhang and Wei[27]propose the use of momentum to improve GCG’s performance. NeuralExec [15] is another attack that employs a similar gradient-based optimization approach to execute prompt injection attacks. Unlike GCG, NeuralExec’s adversarial prompt comprises both a prefix (D ′ p ) and a suffix (D ′ s ), i.e., D ′ = D ′ p + I ′ + D ′ s , which are both optimized using gradients. 3.2 Defenses A fundamental challenge identified in prior work is that LLMs often lack the ability to distinguish between tokens originating from different sources, treating them with equal priority. This absence of privilege levels allows adversarial instructions to sometimes override legitimate user instructions, thereby facilitating prompt injection attacks. To address this issue, recent studies [3,24] propose structuring input tokens to assign varying levels of privilege to tokens from different sources (e.g., 3 Decoder Decoder 푇 ! Decoder 푆 ! 푇 " 푆 # Decoder Decoder 푇 ! Decoder 푆 ! 푇 " 푆 # + + Decoder Decoder 푇 ! Decoder 푇 " 0 Delimiters a InstructionSegment Embedding(ISE) b Augmented Intermediate Representations (AIR) c 1 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 ( ⃗ S 0 , ⃗ 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. system, user, data). This privilege information can then be leveraged by the model to determine the appropriate response in scenarios involving conflicting instructions. Several defense mechanisms have been developed based on this core principle. Recipe for a Defense. Most of these defenses [24,3,25,4] follow a common high-level procedure to create robust models, which we outline below. 1.Establish an instruction hierarchy (IH) by defining the number of privilege levels and their relative order of importance (e.g., P 0 > P 1 > P 2 ). 2.Construct an adversarial training datasetD ′ comprising examples with conflicting instructions embedded within different parts of the input (analogous to a prompt injection attack). 3. Modify the LLM to accommodate IH signals that encode the privilege levels of each token. 4.Train the modified LLM usingD ′ to prioritize instructions associated with higher privilege levels. Existing defenses differ primarily in how they modify the LLM to process IH signals and how they train the LLM (Steps 3 and 4 above). To illustrate, consider a simplified scenario with two privilege levels,P 0 > P 1 . [24,3] use special delimiter tokens (S 0 , S 1 ) to indicate the privilege levels of input tokens (as depicted in Fig. 2) and train the model using supervised fine-tuning (SFT). SecAlign [4] also encodes IH signals using delimiters and trains the model using direct preference optimization (DPO). Another approach, Instructional Segment Embedding (ISE) [25], proposes adding trainable segment embeddings to the input token embeddings to encode privilege level information. Limitation of Existing Defenses. Our work focuses on the method of injecting the IH signal into the LLM. A common characteristic of prior defenses is that they inject the IH signal exclusively at the input layer, either through special delimiter tokens or by appending segment embeddings to the input token embeddings. We hypothesize that this input-level injection limits the effectiveness of the IH signal in enforcing the instruction hierarchy as it propagates through the decoder layers. 4 Our Proposal: Augmented Intermediate Representation The primary goal of our work is to enhance the efficacy of IH signals by injecting them directly into all layers of the model. We do so by modifying the decoder block to incorporate the IH signal. Notations. Before explaining our proposal, we introduce some notation. Let⃗x ij denote the interme- diate token representation of thei th input token in thej th decoder block. Assuming that we haveK privilege levels, let’s use k i ∈ [0, K) to denote the privilege level corresponding to the i th token. Design. We set out to find a method for injecting IH signals to each decoder layer in a way that allows the IH signal to be customized to the intermediate representations at the input of each layer. The key changes made by AIR to the decoder block are illustrated in Fig. 3. AIR introduces a trainable 4 embedding tableS j to each decoder block, consisting ofKentries - one for each privilege level in the IH (Fig. 3 showsK = 2entries for simplicity). The vectors in this table are sized to have the same dimensionality as the intermediate token representations⃗x ij . AIR directly injects the IH signals (k i ) to all the decoder blocks as shown in Fig. 2c. The injected IH signal is used to index the IH embedding tableS j to retrieve an IH vector, which then augments the intermediate token representation ⃗x ij to become ⃗x ′ ij , as defined by: ⃗x ′ ij = ⃗x ij + ⃗ s k j ,where ⃗ s k j = S j [k i ](1) We also augment the intermediate token representation after the last decoder layer, before it’s fed to the linear layers to output the final logits. Layer Norm Masked Self Attn. + + IH Emb. Table (푆 ! ) 푘 ! 푥⃗ !" 푠⃗ ! # 푥⃗ !" $ Layer Norm Feed Forward + 풋 풕풉 푫풆풄풐풅풆풓 푩풍풐풄풌 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 augment- ing the intermediate representation. Overheads. Our method introduces a small increase in the number of parameters. E.g. for Llama3.1-8B (32 decoder layers and hidden representations of size4096), with 3 privilege levels, we require a total of(32 + 1)× 3× 4096 = 0.4Mextra parameters (i.e.0.005%increase). While additional compute is needed to train the model (see Section 5.2), it is similar to the overheads incurred in prior works [24,3,4]. The increase in the compute for inference is negligibly small. Similarity to Research on Positional Embedding. Our proposal shares an interesting similarity with the research on positional embeddings. While earlier works primarily injected positional information at the input layer, often in the form of sinusoidal positional encoding [23] or learn- able positional embeddings [6], more recent methods have explored alternative approaches. Notably, Rotary Posi- tion Embedding (RoPE) [20] injects relative positional information directly into the self-attention mechanisms within all layers of the Transformer. Integrating positional information throughout the model’s architecture, rather than just at the initial input stage, has been shown to be a significant factor in enhancing the performance of large language models [20,28,8]. Our proposal applies the same underlying principle—distributing critical privilege information across all layers—to improve model security against prompt injection attacks. 5 Experimental Setup Our experimental evaluations aim to quantify the impact of different mechanisms for injecting IH signals on model utility (performance in non-adversarial settings) and robustness (resilience under attack). We describe key details of the experimental setup in this section. Additional details can be found in Appendix B 5.1 Models We consider three pre-trained base models of varying sizes: Llama-3.2-3B [1], Qwen2.5-7B [22], and Llama-3.1-8B [9]. In their original pre-trained state, these models exhibit limited instruction- following capabilities. We adapt the architecture of these models to facilitate the injection of IH signals and subsequently train them as described below. 5.2 Training For a fair comparison, all models in our experiments undergo the same training procedure, regardless of the IH injection mechanism. This procedure involves two sequential rounds of training: 5 <|begin_of_text|><|start_header_id|>system<|end_header_id|>Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.<|eot_id|><|start_header_id|>user<|end_header_id|>Evaluate this sentence for spelling and grammar mistakes. He finnished his meal and left the resturant<|eot_id|> <|start_header_id|>assistant<|end_header_id|>There are two spelling errors in the sentence. The corrected sentence should be: "He finished his meal and left the restaurant."<|eot_id|> Figure 4: A sample from the Alpaca dataset formatted using a chat template. Each example consists of an instructionI, an optional data segmentDand the responseR. We use 3 privilege levels: P 0 >P 1 >P 2 to indicate the relative priority of different segments. 1.Non-adversarial Instruction Tuning: First, to instill instruction-following capabilities, the base models undergo full fine-tuning with SFT using an instruction-following dataset. The learning rate (LR) is set to2× 10 −5 for Llama-3.2-3B, and1× 10 −5 for Qwen-2.5-7B and Llama-3.1-8B. 2.Adversarial Robustness Training: Subsequently, to enhance robustness against prompt injection attacks, the models undergo a second stage of fine-tuning using a curated adversarial dataset. For this adversarial training stage, we investigate two fine-tuning methodologies: • SFT: We employ full fine-tuning with a LR of 1× 10 −5 • DPO: We perform parameter efficient fine-tuning using LoRA [11] with a LR of 2× 10 −4 . Each round consists of 3 epochs of training using the AdamW [14] optimizer and a linear LR scheduler. Details of the training datasets used for the two rounds are provided in Appendix B.1 5.3 Defenses This subsection details the Instruction Hierarchy (IH) adopted in our experiments and the various mechanisms evaluated for injecting IH signals into the models. Instruction Hierarchy (IH). We define three hierarchical levels of privilege,P 0 > P 1 > P 2 , as illustrated in Fig. 4.P 0 is assigned to system and user instruction tokens.P 1 is assigned to tokens within the data segment. P 2 is associated with the model’s response tokens. IH Injection Mechanisms. In addition to AIR, our proposed approach, we evaluate two existing methods for injecting IH signals: 1. Delimiters [24,3]: We use two trainable special tokens, [INST] and [INPT], to explicitly mark the beginning of instruction (privilege P 0 ) and input (privilege P 1 ) segments, respectively. 2.Instructional Segment Embedding (ISE) [25]: This method adds distinct, trainable embeddings to the token representations to indicate the IH level of each token in the input. Connection to Prior Work. Existing defense strategies can often be characterized by their choice of IH signal injection mechanism and the adversarial robustness training technique employed. For instance, the methods in [24] and [3] can be viewed as utilizing Delimiters in conjunction with SFT. The approach in [25] employs ISE with SFT. SecAlign [4] uses Delimiters with DPO. Our work extends these investigations by systematically evaluating a broader matrix of IH injection mechanisms (Delimiters, ISE, AIR) and adversarial training techniques (SFT, DPO), including combinations not explored in prior studies. 5.4 Evaluation Methodology Following the training stages, the models are evaluated on two key aspects: utility in non-adversarial settings and robustness against prompt injection attacks. We use two datasets- AlpacaFarm [7] and SEP [30] to measure both utility and robustness. AlpacaFarm. To assess model utility, responses are generated for the 805 test instances from the AlpacaFarm dataset. Each instance in this dataset consists of an instruction and an optional input segment. We employ AlpacaEval 2.0 [12] for utility evaluation. This framework computes the win rate of the generated responses by comparing them against reference responses using a judge model. For our evaluations, responses from thetext-davinci-003model serve as the reference and Llama-3-70B-Instruct is used as the judge model. 6 Robustness evaluations are conducted on the subset of 208 instances from the AlpacaFarm dataset that contain input segments (D). These input segments are modified to simulate prompt injection attacks by appending an adversarial prefixD ′ p and a targeted adversarial instructionI ′ (i.e.,D → D+D ′ p +I ′ ). For all robustness tests, the adversarial instructionI ′ is set toprint exactly ‘hacked!’. The content of the adversarial prefixD ′ p varies depending on the attack strategy. We consider two categories of attacks: •Static Attacks (Black-Box): We evaluate against four predefined static attack strategies: Naive, Ignore, Completion, and Escape Separation. Attack Success Rate (ASR) is reported across the 208 test instances by determining if the model’s response contains the literal phrase hacked!. •Gradient-Based Attack (White-Box): We evaluate white-box robustness using the GCG at- tack [29], enhanced with momentum [27]. A 100-token random prefix is optimized for 200 (DPO models) or 50 (SFT models) steps to minimize attack loss. ASR is measured using the likelihood (from model logits) of generating the target phrase hacked!. SEP. Zverev et al.[30]propose a methodology to evaluate a model’s ability to separate instructions from data using the SEP dataset. This dataset contains 9160 examples—each comprising an instruc- tions i , associated datad i , a probex i , and a witnessw i . The probex i instructs the model to include the witness w i in its response. To evaluate utility, the probe is randomly inserted at the beginning or end of the instruction segment. The model’s response is then checked for the presence ofw i . Since the probe is part of the instruction segment, the model’s output should ideally containw i . Utility is therefore measured as the fraction of responses that include the witness. Ify I i n i=1 denotes the set ofnresponses where the probe was inserted into the instruction segment, the empirical utility scoreUis calculated as:U = 1 n P n i=1 1 w i ∈y I i To evaluate robustness, the probe is similarly inserted randomly at the beginning or end of the data segment, and the response is checked forw i . In this case, because the probe is within the data segment, the model should ideally ignore the probe’s instruction, and its output should not contain w i . Zverev et al.[30]propose the empirical separation scoreSto quantify how well the model distinguishes instructions in the instruction segment from those embedded in the data segment. If y D i n i=1 denotes the set ofnresponses where the probe was inserted into the data segment, the empirical separation scoreSis calculated as:S = P n i=1 1 w i ∈y I i ∧w i /∈y D i P n i=1 1 w i ∈y I i . A higher separation score indicates greater robustness against prompt injection attacks. 6 Results 6.1 AlpacaFarm Utility. Figure 5 compares the utility of models trained with different adversarial training methods (DPO, SFT) and IH injection mechanisms, evaluated on the AlpacaFarm dataset. Compared to a model trained only non-adversarially (None in Fig. 5), our proposed AIR method generally does not significantly degrade model utility. The primary exception is the Llama-3.1-8B model trained with SFT, for which we observe a 4.2% degradation in utility. Robustness (Static Attacks). Table 1 provides the ASRs for models with different defenses against four static attacks: Naive, Ignore, Completion, and Escape Separation. Although the training and Llama-3.2-3BQwen-2.5-7BLlama-3.1-8B 75 80 85 90 95 Win Rate (%) +2.7 -1.6 +3.0 +4.3 +1.2 +2.2 +4.4 -1.9 -0.8 DPO Llama-3.2-3BQwen-2.5-7BLlama-3.1-8B 75 80 85 90 95 +2.2 -0.7 +1.6 +0.3 -1.4 -0.5 +1.9 -0.5 -4.2 SFT IH Mechanism None Delim AIE AIR 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). 7 Table 1: Attack success rates↓(%) for models trained with different IH injection mechanisms (None, Delim., ISE, AIR) and adversarial training techniques (None, SFT, DPO) under various static and gradient-based attacks crafted from the AlpacaFarm dataset. Numbers in bold indicate that the corresponding IH mechanism outperforms other methods for a given attack. ModelAttack NoneSFT|DPO NoneDelim.ISEAIRDelimISEAIR Llama-3.2-3B Naive10.00.00.00.00.00.0 Ignore2.50.00.00.00.00.00.0 Completion3.810.50.00.00.00.0 Escape Sep.1.40.50.50.50.00.00.0 GCG77.53848.14.129.146.65.2 Qwen-2.5-7B Naive3.40.00.50.00.00.00.0 Ignore2.90.00.00.00.00.00.0 Completion3.810.00.00.00.00.0 Escape Sep.2.90.50.50.50.50.00.0 GCG99.58836.622.6327.71.6 Llama-3.1-8B Naive0.50.00.00.00.00.00.0 Ignore2.50.00.00.00.00.00.0 Completion3.80.00.00.00.00.00.0 Escape Sep.1.40.50.00.00.00.00.0 GCG99.57719.911.31342.8 test set examples are distinct, the model encounters the first two attacks during adversarial training (in-distribution), while the other two are unseen and thus out-of-distribution. We find that all three IH injection mechanisms (Delimiter, ISE, and AIR) offer near-perfect protection against all evaluated static attacks. Robustness (Gradient-Based Attack). Figure 6 illustrates the comparative performance of these defenses against the Momentum-Boosted GCG attack. The figure plots the attacker’s loss—calculated relative to the target adversarial response—as a function of GCG optimization steps. Each line indicates the mean loss over 208 test instances, with shaded regions representing the standard deviation. Results are presented separately for models adversarially trained with DPO (first row of plots) and SFT (second row). 050100150200 Step 0 10 Average Loss Llama-3.2-3B 050100150200 Step 0 10 20 Qwen-2.5-7B 050100150200 Step 0 10 20 30 Llama-3.1-8B IH Mechanism Delim AIE AIR 02040 Step 0 5 10 Average Loss Llama-3.2-3B 02040 Step 0 5 Qwen-2.5-7B 02040 Step 0 5 10 Llama-3.1-8B IH Mechanism None Delim AIE AIR (a) DPO (b) SFT Figure 6: Average loss from the Momentum-Boosted GCG attack comparing different defenses during various points in the optimization process. AIR is more robust to GCG with a higher average loss compared to prior works across all models and both optimization methods. 8 As anticipated, the attacker’s loss diminishes with more GCG optimization steps, signifying increased attack efficacy. Notably, models defended by our proposed AIR mechanism consistently incur a significantly higher average attacker loss compared to those defended by ISE or Delimiters. Further- more, AIR’s ASR against GCG attacks (GCG in Table 1) is1.6×to9.2×lower than that of the next best defense, underscoring its superior robustness. Our findings also reveal that adversarial training with DPO yields more robust models than SFT, corroborating results from SecAlign [4]. Progression in Robustness. These results highlight a clear progression in defense efficacy. Recall that the Delimiters mechanism injects IH signals via special tokens at segment boundaries, while the ISE method applies IH signals (through dedicated embeddings) to all tokens in the input. The enhanced robustness observed when moving from Delimiters to ISE suggests the benefit of more pervasive IH signal application at the input level. Our AIR approach further advances this principle; by injecting IH signals directly into all decoder layers, rather than confining them to the input representations, AIR achieves a more deeply integrated hierarchical understanding within the model, leading to the observed superior robustness against this strong gradient-based attack. 6.2 SEP Figure 7 plots empirical separation and utility scores, comparing the different IH injection mechanisms. For models trained with DPO (Fig. 7a), AIR achieves the highest separation and utility scores, outperforming other IH mechanisms as well as all models trained with SFT in these combined metrics. For models trained with SFT, AIR maintains higher separation scores than other methods across all models. However, in some instances (e.g., Qwen-2.5-7B, Llama-3.1-8B), AIR-SFT’s utility can be lower than the None baseline (which undergoes only non-adversarial training). Overall, these results indicate that AIR consistently enhances the model’s ability to separate data from instructions and, when trained with DPO, provides the best utility-separation tradeoff for the evaluated models. 0.250.300.350.40 Utility Score 0.4 0.6 0.8 Separation Score Llama-3.2-3B 0.3750.4000.4250.4500.475 Utility Score Qwen-2.5-7B 0.350.40 Utility Score Llama-3.1-8B 0.250.300.350.40 Utility Score 0.4 0.6 0.8 Separation Score Llama-3.2-3B 0.3750.4000.4250.4500.475 Utility Score Qwen-2.5-7B 0.350.40 Utility Score Llama-3.1-8B (a) DPO (b) SFT NoneDelimISEAIRBest Utility×Separation Figure 7: Utility and Separation scores derived from the SEP dataset. IH mechanisms with the best utility× separationfor each model (across both DPO and SFT) are marked with⃝, indicating that they offer the best trade-off. 7 Conclusion Our paper proposes a new defense for prompt injection attacks. We study the various mechanisms of injecting instruction hierarchy information in prior work and find that they suffer from a crucial limitation – they only insert the IH information to the input layer of the LLM, which limits the efficacy of the IH signal. To overcome this drawback, we propose Augmented Intermediate Representations (AIR), which injects the IH signals into all the decoder layers in the model. Through extensive empirical studies on models of different sizes (3B, 7B, 8B), and training techniques (SFT, DPO), we show that our proposal can improve robustness against gradient-based attacks by1.6×to9.2×, without significant degradation in utility. 9 References [1] MetaAI.Llama3.2:Revolutionizingedgeaiandvision withopen,customizablemodels.https://ai.meta.com/blog/ llama-3-2-connect-2024-vision-edge-mobile-devices/, September 2024.Ac- cessed: 2025-05-15. [2]Gabriel Alon and Michael Kamfonas. Detecting language model attacks with perplexity. arXiv preprint arXiv:2308.14132, 2023. [3] Sizhe Chen, Julien Piet, Chawin Sitawarin, and David Wagner. Struq: Defending against prompt injection with structured queries. arXiv preprint arXiv:2402.06363, 2024. [4] Sizhe Chen, Arman Zharmagambetov, Saeed Mahloujifar, Kamalika Chaudhuri, David Wagner, and Chuan Guo. Secalign: Defending against prompt injection with preference optimization, 2025. URL https://arxiv. org/abs/2410.05451, 2024. [5]Edoardo Debenedetti, Jie Zhang, Mislav Balunovi ́ c, Luca Beurer-Kellner, Marc Fischer, and Florian Tramèr. Agentdojo: A dynamic environment to evaluate attacks and defenses for llm agents. arXiv preprint arXiv:2406.13352, 2024. [6] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pages 4171–4186, 2019. [7]Yann Dubois, Chen Xuechen Li, Rohan Taori, Tianyi Zhang, Ishaan Gulrajani, Jimmy Ba, Carlos Guestrin, Percy S Liang, and Tatsunori B Hashimoto. Alpacafarm: A simulation framework for methods that learn from human feedback. Advances in Neural Information Processing Systems, 36:30039–30069, 2023. [8]Philipp Dufter, Martin Schmitt, and Hinrich Schütze. Position information in transformers: An overview. Computational Linguistics, 48(3):733–763, 2022. [9]Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. [10]Kai Greshake, Sahar Abdelnabi, Shailesh Mishra, Christoph Endres, Thorsten Holz, and Mario Fritz. Not what you’ve signed up for: Compromising real-world llm-integrated applications with indirect prompt injection. In Proceedings of the 16th ACM Workshop on Artificial Intelligence and Security, pages 79–90, 2023. [11] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1 (2):3, 2022. [12]Xuechen Li, Tianyi Zhang, Yann Dubois, Rohan Taori, Ishaan Gulrajani, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. AlpacaEval: An Automatic Evaluator of Instruction- following Models. https://github.com/tatsu-lab/alpaca_eval, 2023. [13]Yupei Liu, Yuqi Jia, Jinyuan Jia, Dawn Song, and Neil Zhenqiang Gong. Datasentinel: A game-theoretic detection of prompt injection attacks. arXiv preprint arXiv:2504.11358, 2025. [14]Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [15] Dario Pasquini, Martin Strohmeier, and Carmela Troncoso. Neural exec: Learning (and learning from) execution triggers for prompt injection attacks. In Proceedings of the 2024 Workshop on Artificial Intelligence and Security, pages 89–100, 2024. [16] Fábio Perez and Ian Ribeiro. Ignore previous prompt: Attack techniques for language models. arXiv preprint arXiv:2211.09527, 2022. 10 [17]Gene Ruebsamen. Cleaned Alpaca Dataset, February 2024. URLhttps://github.com/ gururise/AlpacaDataCleaned. [18]Mrinank Sharma, Meg Tong, Jesse Mu, Jerry Wei, Jorrit Kruthoff, Scott Goodfriend, Euan Ong, Alwin Peng, Raj Agarwal, Cem Anil, et al. Constitutional classifiers: Defending against universal jailbreaks across thousands of hours of red teaming. arXiv preprint arXiv:2501.18837, 2025. [19]rgorman Stuart Armstrong.Using GPT-Eliezer against ChatGPT Jailbreaking, De- cember 2022.URLhttps://w.lesswrong.com/posts/pNcFYZnPdXyL2RfgA/ using-gpt-eliezer-against-chatgpt-jailbreaking. [20]Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. [21] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto.Alpaca: A strong, replicable instruction- following model. Stanford Center for Research on Foundation Models. https://crfm. stanford. edu/2023/03/13/alpaca. html, 3(6):7, 2023. [22] Qwen Team. Qwen2.5: A party of foundation models, September 2024. URLhttps:// qwenlm.github.io/blog/qwen2.5/. [23]Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [24]Eric Wallace, Kai Xiao, Reimar Leike, Lilian Weng, Johannes Heidecke, and Alex Beutel. The instruction hierarchy: Training llms to prioritize privileged instructions. arXiv preprint arXiv:2404.13208, 2024. [25]Tong Wu, Shujian Zhang, Kaiqiang Song, Silei Xu, Sanqiang Zhao, Ravi Agrawal, Sathish Reddy Indurthi, Chong Xiang, Prateek Mittal, and Wenxuan Zhou. Instructional segment embedding: Improving llm safety with instruction hierarchy. arXiv preprint arXiv:2410.09102, 2024. [26] Yohei. injection test, October 2022. URLhttps://x.com/yoheinakajima/status/ 1582844144640471040. [27]Yihao Zhang and Zeming Wei. Boosting jailbreak attack with momentum. In ICASSP 2025- 2025 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1–5. IEEE, 2025. [28]Liang Zhao, Xiachong Feng, Xiaocheng Feng, Weihong Zhong, Dongliang Xu, Qing Yang, Hongtao Liu, Bing Qin, and Ting Liu. Length extrapolation of transformers: A survey from the perspective of positional encoding. arXiv preprint arXiv:2312.17044, 2023. [29]Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.15043, 2023. [30]Egor Zverev, Sahar Abdelnabi, Soroush Tabesh, Mario Fritz, and Christoph H Lampert. Can llms separate instructions from data? and what do we even mean by that? arXiv preprint arXiv:2403.06833, 2024. 11 A Limitations and Future Work While our defense demonstrates strong average resilience to white-box attacks, it does not provide formal robustness guarantees, meaning specific outliers or advanced attacks might still succeed. This is a common limitation in the current LLM robustness research landscape. Additionally, our utility and robustness evaluations, similar to prior work, are confined to single-turn interactions using the AlpacaFarm and SEP datasets. Evaluating our proposal’s effectiveness in multi-turn conversational settings and complex agentic workflows is therefore a key direction for future work. B Additional Experimental Details B.1 Training Datasets Non-Adversarial Dataset. For the first stage of training (non-adversarial instruction tuning), we employed the cleaned version [17] of the Alpaca dataset [21]. This dataset comprises approximately 52K examples. As illustrated in Fig. 4, each example typically consists of an instruction (I), an optional input segment (D), and the desired response (R). The models are trained to generate R given I and D (when present), formatted according to a specific chat template. For the second stage, adversarial robustness training, we constructed two distinct adversarial versions of the Alpaca dataset: one for SFT and another for DPO. Adversarial SFT Dataset. This dataset incorporates all examples from the original Alpaca dataset. • Examples that originally lack an input segment (D) are included unmodified. • For examples that do contain an input segment (D), half are included unmodified. The other half are modified to simulate a prompt injection attack. The input segmentDis transformed into ˆ Dby concatenating the original input, an adversarial prefixD ′ p , and an adversarial instructionI ′ (i.e., ˆ D = D + D ′ p + I ′ ). The adversarial prefixD ′ p is determined by either the Naive or Ignore attack strategy, chosen with uniform probability. The adversarial instructionI ′ is an instruction randomly selected from a different example within the Alpaca dataset. This adversarial SFT dataset can be represented as collections of tuples(I, ̄ D, R), where ̄ Dis either the original inputD, the modified input ˆ D , or absent (if the original example had no input segment). Adversarial DPO Dataset. To construct the preference dataset for DPO, we exclusively used Alpaca examples that contain an input segment (D). For each such example, we generated a corrupted input segment ˆ Dusing the same Naive or Ignore prompt injection techniques (resulting in ˆ D = D +D ′ p +I ′ as described above). The preference pair consists of the original instructionIand the corrupted input ˆ D. The chosen response is the original, correct responseRfrom the Alpaca dataset (corresponding toIandD). The rejected response is the responseR ′ associated with the adversarial instructionI ′ in its original Alpaca example. This DPO dataset is a collection of tuples (I, ˆ D, R, R ′ ). All examples across these datasets were formatted using the chat template depicted in Fig. 4 before being used to train the models. B.2 Model and Training Configurations For all training runs, we use a batch size of 4 with 4 steps of gradient accumulation for both rounds of training. We employed Parameter-Efficient Fine-Tuning (PEFT) using the Low-Rank Adaptation (LoRA) technique to fine-tune the model with DPO. Specifically, we fine-tuned the query (q_proj) and value (v_proj) projection layers. The LoRA hyperparameters were set with a rank (r = 64), lora_alpha= 8, and lora_dropout= 0.1. Embedding Table Initialization. Our method introduces embedding tables within the decoder block to augment intermediate representations. These tables are initialized by default with vectors sampled from a normal distribution with a standard deviation of0.02(N (0, 0.02 2 )). While this initialization proved effective for Llama models, it yielded suboptimal robustness performance for the Qwen model. We attribute this discrepancy to the significantly larger magnitude of intermediate representations produced by Qwen; the default, smaller embedding vectors failed to sufficiently 12 modify these representations. To rectify this, we increased the initialization standard deviation fivefold to0.1(N (0, 0.1 2 )) specifically for the Qwen model, which demonstrably improved our defense’s effectiveness. For a fair comparison, this same adjusted initialization was applied to the ISE technique when used with Qwen. Due to computational constraints, exhaustive tuning of this hyperparameter was not feasible and is deferred to future work. C Additional Related Work Detection-Based Defenses. The related work in Section 3.2 primarily discussed defenses designed to enhance model robustness against prompt injection by defining an instruction hierarchy. In addition to these, a significant class of defenses focuses on detecting malicious or unintended instructions within user inputs or data segments before they cause the main LLM to deviate from its intended behavior. The core idea is to employ a detection mechanism as a preliminary check or ongoing monitor. Several approaches to detection-based defenses have been proposed: •LLM-Powered Detectors: A common strategy is to leverage an LLM itself as a detector. These approaches include using zero-shot or few-shot prompting of an LLM to ascertain if an input contains hidden or malicious instructions [19]. Another technique involves fine-tuning a dedicated LLM to act as a specialized classifier or "guard" model for identifying malicious prompts or instruction injections [18]. Furthermore, LLM self-evaluation techniques have been explored, where the model attempts to determine if it is being manipulated. •Known Answer Detection: Another interesting line of work focuses on testing if the LLM returns a known answer in the presence of potentially malicious tokens [26]. This method uses a special instruction where the answer is only known to the detector. If the response fails to provide the expected answer in the presence of a data segment, then the data segment is flagged as containing a prompt injection attack. A recent work [13] extends this idea using a game-theoretic foundation to train a detector LLM that is very sensitive to prompt injection attacks, achieving near-perfect scores on benchmarks. However, such defenses remain vulnerable to adaptive attacks (e.g., if the attacker instructs the LLM to return the known answer before following the attacker’s instructions). •Output Analysis and Verification: Instead of, or in addition to, input checks, some defenses analyze the LLM’s output. This includes response checking, which evaluates whether the LLM’s output aligns with the intended task or original user instruction, where deviations might indicate manipulation [18]. Perplexity-based detection has also been explored to identify anomalous outputs [2]. While detection-based methods offer a valuable layer of security, they remain vulnerable to adaptive attacks. Therefore, such defenses can complement our proposed defense, which is designed to make the model inherently robust to prompt injection attacks. D Compute Resources We use compute nodes with8×A100 GPUs paired with 256 CPU cores and 1TB of memory and 25 TB of storage for all our experiments. Note that most of our training runs complete within 2 hrs. The gradient based attacks need more time due to their sequential nature and require around 30 mins per example with a single gpu. E Societal Impact The research presented in this paper aims to enhance the security and reliability of LLMs by proposing a more robust defense (AIR) against prompt injection attacks. Positive impacts include increased user trust and safety when interacting with LLM-powered applications, particularly those processing untrusted external data like emails or web content. By making models less susceptible to malicious instruction hijacking, this work could facilitate the safer deployment of helpful AI agents in various domains, reduce the potential for AI-driven misinformation or data exfiltration triggered by such attacks, and contribute to the broader adoption of LLMs for beneficial tasks. However, potential negative consequences or challenges must also be considered. Improved defenses might lead to over-reliance or a false sense of complete security, potentially discouraging complementary security measures. Ultimately, while techniques like AIR contribute positively towards trustworthy AI, they 13 should be viewed as one component within a larger framework for responsible AI development and deployment. 14