PFT: Enhancing Prompt Injection Robustness via Position-Enhanced Finetuning

1. The contribution appears limited to "pure" closed-domain tasks, where users are restricted to providing only input data without additional task-specific instructions. In real-world scenarios (e.g., translation tools or platforms like ChatGPT and Copilot), users often add specific instructions, such as style preferences or formatting guidelines, which PFT does not address. Figure 2 suggests that PFT might not perform well in these more dynamic contexts.

2. The paper states that PFT imposes no performance penalty (Section 5.3), yet assumes that users will not add prompts, focusing solely on task input data. It would be helpful to analyze the trade-off between robustness and performance in a practical context where users provide detailed instructions. Examining the parameter $d$ and its relationship with performance in such scenarios would clarify PFT’s effectiveness.

3. Including a comparison between PFT and a simple baseline prompt designed to direct LLMs toward system tasks would add value. Observers might find a "SOTA vs. SOTA+PFT" comparison more informative than "vanilla vs. vanilla+PFT."

4. In the robustness experiments (Figure 5), only data points with fewer than ten inserted sentences (#ofSentences < 10) seem relevant for real-world usage. It would be helpful to highlight these cases more explicitly.

5. No code is provided, which limits the reproducibility of the findings.

6. Table 1 is difficult to interpret due to the lack of a legend or clear explanation of its components. Adding this information would improve clarity.

1. Choice of the model to be fine-tuned.
   1. Why do you fine-tune instructed models instead of the base pre-trained model? The instructed models can already solve these tasks so it is perhaps expected that the utility is not hurt via PFT. These models are also already aligned (via RLHF or safety tuning) so they should already be somewhat resilient to such attacks.
   2. Or, the authors intend to propose PFT as an additional step after RLHF? However, I believe that the authors are suggesting to replace SFT with PFT, and if so, the base model for this experiment should be the base Llama-3-8B, not instruct.
2. Weak defense baselines.
   1. I believe that this experiment is missing two important baselines: Instruction Hierarchy and StruQ. My understanding is that the authors intend to use “Delimiter-enhanced SFT” to represent StruQ and “Data-augmented SFT” to Instruction Hierarchy. However, I would suggest reproducing these baselines exactly or as close as possible and compare to them in the same setup (same dataset). Without this direct comparison, it is impossible to conclude whether PFT is better.
   2. Data-augmented SFT: What happens if the augmentation is done to the user prompt instead of the system prompt? This would be more like StruQ.
   3. Data augmentation is used in both Instruction Hierarchy and StruQ, and it is different from what’s done in this paper. What about training against a subset of the attack? If we do that and combine with PFT, is there any improvement?
3. Weak attack baselines. The authors should consider evaluating against stronger attacks, e.g., a set of handcrafted attacks from StruQ, jailbreaks, or even automated attacks such as PAL, TAP, AutoDAN, or GCG. This would help with comparison to the prior works.
4. KL divergence metric. It is mentioned that the KL divergence is computed between p_model(output text|prompt) of the model before and after fine-tuning. As far as I know, there is no way to efficiently compute this because the set of “output texts” is just too large. My guess is that the authors compute the KL divergence at each token conditioned on the prior tokens, which is a different quantity from what is stated.
5. Not applicable to real-world / more diverse use cases. This defense does not hurt the utility because all of the Alpaca samples put the instruction close to the beginning of the input. This defense would fail immediately when legitimate instructions are part of the user input and placed anywhere, e.g., in chatbots.

Nitpicky Stuff

• In Table 1, calling the model before fine-tuning a “base” model is slightly confusing. In my understanding, this is an aligned and instructed model, but "base model" often refers to a pre-trained model on the next-token prediction task. I might suggest simply calling it "before fine-tuned," or at least make it very in the text or caption (not delay until Section 5).
• L193 (”this is not a strong enough signal for LLMs…”): I believe that this statement is plausible but still not very convincing. Could you explain more why the delimiters would be a weaker signal than the begin-of-text token? Is it because there is only one begin-of-text token in every sample, and it is always at the beginning? -- I see that this is discussed later on in Section 3.2. It might be a good idea to refer to it here.
• L211 (”Also the first insertions have a dramatic effect…”): I can understand the message from looking at Figure 3, but this statement is hard to follow and is pretty confusing to me.
• L285 (”assigning a unique signature to each token based on its role (system or user), we create a continuous, fine-grained distinction throughout the entire input.”): I believe that this sentence is not an accurate description of PFT. Each token is pushed back by the same offset $d$ which does not sound like a unique signature to me. I’m also not sure what “continuous, fine-grained distinction” refers to.

Overall, I struggled to evaluate this paper. The paper has some interesting results, and I'm not sure what to make of them. So I'm not sure whether to recommend acceptance or not for this paper.

The paper's notion of robustness is: the LLM should be able to tolerate irrelevant instructions appearing before the key instruction, in the system prompt. Is this important? Do real applications use system prompts that contain many irrelevant instructions before the key instruction? I'm not sure. Therefore, I'm unsure whether the problem this paper tackles is important. I encourage the authors to provide examples or evidence of real-world applications where this property is important.

The notion of robustness that I'm used to is a bit different: can the LLM resist all attacks? If we pick some class of attacks, what is the attack success rate of the strongest attack in that class? In other words, increasing robustness means reducing the attack success rate of some attack -- and this is evaluated for average-case prompts (i.e., ones that will appear in real applications), rather than worst-case prompts (e.g., where we add extraneous instructions at the start). That's not the notion this paper takes on, though.

The paper starts from a premise about how LLMs will/should be used (put the instruction in system message, the data in user message), a premise that I am skeptical about. Then it draws some conclusions about that usage. I'm not sure whether those conclusions generalize to ways of using LLMs that I think are more appropriate and more common. Also, I'm not sure whether the paper's results generalize to multiple models (different LLMs).

I also struggle to tell whether this paper's results are primarily telling us something about prompt injection (inserting malicious instructions into a field that is only supposed to contain data) or system-message-following (supplying user messages that contradict/violate rules/guardrails established in the system message). Details on experiments are vague and so it's hard for me to tell what is being tested.

Detailed feedback:

Abstract: I had a hard time understanding the abstract. I don't understand what is meant by "the key instruction", or what system instructions have to do with prompt injection.

Sec 1, paragraph 2: I don't agree with the claim here about how engineers typically build systems with LLMs. I don't think the zero-shot prompt is typically put in the system instructions. Instead, I think the prompt or instruction is typically put in the user message, and the system message typically contains guardrails that constrain what types of prompts/instructions will be followed. If you disagree, I encourage you to look for quantitative evidence (perhaps surveying some collection of systems built by others).

Sec 1: the paper seems to conflate two issues that I consider separate: (a) prioritizing system instructions over user instructions when they conflict; (b) ignoring all instructions in the data part of the user message. I consider prompt injection to be problem (b), and problem (a) to be a separate problem. The introduction does not distinguish between these two, and that makes it harder for me to understand what problem the paper is and isn't solving.

Or, to put it another way, we can distinguish between system instructions, user instructions, and user data. The paper does not seem to clearly distinguish between these. It seems to assume the system message contains instructions and the user message contains data. In my experience, that is not representative of how LLMs are used in real systems and not representative of how production LLMs are trained in practice. My experience is that the system message (if present) contains instructions (such as guardrails or restrictions or a definition of the domain/scope), and the user message often contains a combination of both instructions and data.

Table 1: It's not clear to me what this is showing. What does a number like 10% mean? What is "Gandalf Summarization"? I think the table caption should provide a self-contained explanation, or the table should be moved to later in the paper after all key concepts have been defined. I think the table caption should specify what the number/percentages mean (what are they measuring? attack success rate?).

Sec 2.1: So many details are missing. What base model do you use? How large is your dataset? What are the attacks you evaluate against?

I am quite skeptical of the claim that the model is secure. I don't believe you've evaluated against enough attacks to draw such a conclusion. There are some quite strong attacks, such as GCG and TAP (modified to create prompt injection attacks), which are not considered here. Therefore, it is not warranted to conclude that the model is secure, as you don't know whether it will be secure against these more sophisticated attacks.

Sec 2.2: Does not match my experience. In my experience, general instructions ("You are an AI assistant") go in system messages, the "key instruction" (e.g., "Translate this to French") goes in the user message, and background knowledge and context and few-shot examples and RAG-retrieved excerpts typically go in the user message.

Many details are missing. What model did you use? What attack did you use? What was the dataset and tasks for evaluation? What is meant by an "attack dataset"?

Fig 2: How is accuracy measured? How do you tell whether the model's respons is accurate?

Sec 3.1: It might help if you stated what prompt injection attack techniques you used. From Fig. 4, it sounds like maybe you used the most naive attack, just directly asking the model to do something different.

Fig 4: I'm not convinced the model is trained to know what "user's input" means, so I'm not sure whether this is a fair test of LLM capabilities or if this is the right way to use LLMs for this kind of task.

Sec 4: Do you have any explanation whether it is better to have a fixed-length gap between system vs user message (i.e., user message starts at token $k+1+d$) or have the user message start at a fixed position (i.e., user message starts at token $d+1$)? Have you tried both?

Sec 4: It seems this assumes that messages will always appear in the order: system message first, then user message second. In contrast, existing LLMs allow them to be interspersed arbitrarily. Does this restriction cause any loss in flexibility? Does it matter?

Also, how do you propose to handle multi-turn interactions? How will they be encoded, and where will gaps appear or not appear?

Sec 5.1: Are you doing SFT on a model that has already been instruction-tuned, or on a base model that has not? It sounds like you are doing SFT on a model that was already instruction-tuned?

Sec 5.1: How do you generate the desired responses to pairs where the system prompt and user input have been swapped?

Sec 5.1: Will you make your dataset and code available, to support reproducibility?

Sec 5.2: I don't understand what the attacks are. Please provide more detail. For Gandalf Summarization, the paper cites Lakera AI, 2023b, but that reference is bogus: the URL is for the Lakera main web site, which clearly does not have the claimed information, nor could I find any other webpage on the Lakera web site with the listed title. The same comments apply to Lakera 2023a. Please provide direct links or references for each attack, and preferably describe the attack in a self-contained way in the paper.

I suspect you might not be testing against the strongest prompt injection attacks, such as found in other papers on the subject. For instance, I would be more convinced if you had evaluated against completion attacks, TAP attacks, and GCG attacks.

Sec 5.3: I disagree with using "most robust" here. I believe what you actually measure is the property "isn't distracted by extraneous instructions at the start". That isn't the same thing. What you study is one narrow, specific aspect of robustness. This is relevant if real systems use system prompts that start with a lot of generic, extraneous instructions. It is less clear how relevant it might be if real systems don't do that. And it does not necessarily imply that PFT can defend against stronger attacks.

Fig. 6: This is missing a key metric: you should also show the accuracy and log-likelihood for the undefended base model with none of these defenses applied. That's what users will really care about: does the defense harm the utility of the model, compared to existing models with no defense applied; not whether your defense is about as good as other plausible defenses.

Consequently, I don't think the paper can reasonably claim that PFT doesn't hurt model performance. It might be true, but I don't think the paper has evaluated that.

Fig. 6(b) shows that there is some deviation from the base model. Is 0.5 KL divergence a large deviation, or a small one? It's hard to know. Perhaps looking at a small random sample of responses from both the base model and PFT model would help.

Sec 6: The paper seems to claim that OpenAI's instruction hierarchy method has fragility when the key instruction appears later in the input. But it's not clear that this has actually been tested. I would find this analysis of the instruction hierarchy method more compelling if the paper empirically measured this, e.g., on GPT 4o-mini.

The paper appears to claim that StruQ has fragility in this case as well, and that PFT is more robust. However, I don't think this has been demonstrated, as the paper does not evaluate StruQ. StruQ uses some techniques that were not tried in any of the models evaluated in this paper. I think the comparison to StruQ would be more convincing if the paper compared to a StruQ-trained model. I don't think we should view StruQ as designed just to defend against Completion attacks; it seems like it is trying to handle all prompt injection attacks, as much as possible.

The paper seems to be missing one piece of related work, BIPIA (Yi et al, arXiv:2312.14197).

1. The definition of the problem lacks clarity. Specifically, the formal definition of a "key instruction" in system prompts is ambiguous. How and why does it differ from other system prompts? It is challenging to distinguish key instructions from other contextual prompts, especially since other prompts can sometimes serve as the context for the key instructions. This will lead to complex scenarios that needs a detailed discussion.
2. Following the first point, the generalizability of PFT needs to be clarified, as key instructions vary across different contexts. The validation dataset used in the paper mirrors the examples provided in the introduction, which does not suffice to demonstrate PFT's applicability in more complex and practical scenarios.
3. Stronger attacks are necessary. In the realm of prompt injection, several existing studies employ learning-based methods to launch attacks, as referenced in [1]. I suggest that the authors include experiments to test the resilience of PFT against these types of attacks.
4. The paper lacks a discussion on adaptive attacks. If attackers know the PFT method and can fine-tune the model accordingly – for instance, tuning it to adhere strictly to user instructions?

[1] Pasquini, Dario, Martin Strohmeier, and Carmela Troncoso. "Neural Exec: Learning (and Learning from) Execution Triggers for Prompt Injection Attacks." arXiv preprint arXiv:2403.03792 (2024).