TL;DR:

• We fine-tuned Llama3-8B to enhance its resistance to common indirect prompt injection attacks on Question/Answer tasks on emails and articles. 
• We released scripts that can be used to reproduce the fine-tuning process on specific datasets and use cases. 
  
• The fine-tuned model and a quantized version are available on Hugging Face and Ollama for testing and experimentation. 
  

Organizations have been increasingly integrating generative AI (GenAI) features into their platforms and applications. For example, Microsoft Copilot and Google Gemini allow users to summarize emails or documents seamlessly within their UIs. NotebookLM goes a step further, and even allows the generation of a full podcast discussing any article or link provided by the user. To implement these GenAI features, applications will write prompts to an underlying LLM: such prompts typically include an instruction (such as “summarize this email”) and the data on which the instruction needs to operate (e.g. email body). When the data originates from untrusted sources, it can contain prompt injection attacks designed to alter the LLM's intended behavior, leading to unexpected or malicious outputs. 

Prompt injection attacks are similar to SQL injection in databases, where untrusted input manipulates queries to produce unintended results. The consequences vary depending on the application. For instance, an injection attack might cause the LLM to generate output that, when rendered in a user's browser, enables data exfiltration or client-side attacks (see example). In more severe cases, especially with autonomous agents, prompt injections can lead the LLM to perform unauthorized and dangerous actions on behalf of an attacker (see example). 

Jailbreaks and prompt injection attacks are ongoing challenges for LLMs. Solutions to mitigate these vulnerabilities generally fall into two categories: implementing external checks on the LLM's inputs and outputs (as we describe in our Security Canvas for Prompt Injection) and training LLMs to be inherently more resistant to such attacks. Neither approach is perfect on its own, but they work well together and are complementary. In this blog, we focus on the latter strategy—fine-tuning models to enhance their resistance to common indirect prompt injection attacks. 

We demonstrate the process to fine-tune Llama3-8B to enhance its resistance to common indirect prompt injection attacks, and evaluate the results. Our approach is inspired by the methods presented in Microsoft's BIPIA paper and OpenAI's Instruction Hierarchy paper but does not require modifications to the embedding layer and uses a limited, task-specific dataset. While the resulting model may not be universally applicable, it offers increased resistance to prompt injection for specific use cases. 

The terms jailbreaks and prompt injection are often used interchangeably, and indeed they are related attacks used to disalign language models (LLMs) from their intended behavior. 

Jailbreaks are techniques that exploit language patterns to override the constraints set by the creators during model training and Reinforcement Learning with Human Feedback (RLHF). These patterns manipulate the model into producing outputs that bypass built-in restrictions or system instructions provided by the user. 

Prompt injection occurs in applications where user data is used to instruct an LLM to perform tasks (e.g., summarizing or translating text). If the user’s input contains hidden instructions or jailbreak patterns, the model’s behavior can be manipulated by an attacker, causing unintended or malicious outputs. 

The most serious form of this is indirect prompt injection, where the harmful data fed to the LLM comes from a third-party attacker (such as the body of an email). In this scenario, the actual user of the GenAI application becomes the victim. This attack vector was first explored in the paper "Not what you've signed up for: Compromising Real-World LLM-Integrated Applications with Indirect Prompt Injection" and can result in severe consequences like social engineering, data exfiltration, or unauthorized actions—especially if the LLM operates as an agent with access to external tools. 

Our approach is inspired by two key research papers: Benchmarking and Defending Against Indirect Prompt Injection Attacks on Large Language Models (BIPIA) by Microsoft and The Instruction Hierarchy: Training LLMs to Prioritize Privileged Instructions by OpenAI. 

The BIPIA paper describes a white-box defense against indirect prompt injection attacks. The authors introduce special <data> and </data> tokens into the embedding layer of the LLM. By fine-tuning the model to ignore any instructions contained within these tokens, the LLM learns to distinguish between trusted instructions and untrusted user content. This approach effectively targets indirect prompt injection but requires modifying the embedding layer, which can be complex to implement. 

The Instruction Hierarchy paper addresses the broader issue of LLMs lacking an inherent instruction hierarchy. Without it, models struggle with conflicting instructions from different sources. The authors define a hierarchy (e.g., system messages, user messages, tool outputs, ...) and train the model using a large dataset to prioritize higher-level instructions over lower-level ones. They use Reinforcement Learning from Human Feedback (RLHF) to teach the model to follow only lower-level instructions that align with higher-level directives. This method does not require modifying the embedding layer but requires a larger dataset and more complex training process than BIPIA’s.

We aimed to enhance Llama3-8B's resistance to indirect prompt injection attacks without modifying its embedding layer. Drawing inspiration from both BIPIA's white-box method and OpenAI's Instruction Hierarchy, we developed a practical solution tailored to specific use cases like summarizing emails or documents. 

Instead of introducing new tokens into the embedding layer as BIPIA did, we leveraged the model's existing architecture by defining data delimiters within the system prompt. We used <<<data>>> and <<</data>>> markers to enclose user-provided content. These delimiters are not new tokens; they are processed by the standard tokenizer, so no changes to the embedding layer are needed. 

The key idea was to fine-tune the model to recognize these delimiters and ignore any instructions found within them. We relied on the system message to establish this behavior, capitalizing on the model's inherent tendency to follow system prompts. The system message explicitly instructs the model to treat text within <<<data>>> markers as data and to disregard any embedded instructions. 

We created a custom dataset comprising benign content and examples of indirect prompt injection attacks. The benign content included emails and articles generated to simulate typical user data. We then crafted various attack payloads consisting of a jailbreak technique followed by a potentially malicious instruction that could be used for data exfiltration or social engineering attacks. We varied the placement of these malicious payloads within the content to expose the model to different attack scenarios. 

Our resulting dataset consisted of 10,000 training samples. This size provided sufficient diversity to cover a range of content and attack scenarios without overfitting the model. 

Each sample consisted of: 

• Context: The user-provided data, such as the content of an email. 
  
• Question: A user query related to the content, such as "Please summarize the email".
• Ideal: The desired response, accurately answering the query while ignoring any instructions within the data delimiters. 

Example entry: 

The final dataset to fine-tune the model on was then created by adding the following system prompt and surrounding the data/context with the “<<<data>>>” delimiters:

The resulting samples (FormattedFinal.jsonl) look like this:

The next section will cover recomemndations for users and developers of GenAI based assistants.We fine-tuned Llama3-8B using the TogetherAI platform, which provides tools for efficient training and deployment of LLMs. The training process lasted approximately 30 minutes. We trained the model for three epochs with a batch size of 8 and a learning rate of 1e-5. 

To assess the effectiveness of our fine-tuning approach, we created a separate test dataset (TESTselection.jsonl). This dataset consisted of prompt injection attacks designed to insert specific canary words into the model's output. The presence of these canary words indicated a successful attack.

We evaluated both the original Llama3-8B model and our fine-tuned model using this test set (Compare.py and Evaluate.py). The results were as follows:

    • Original Model: Achieved a pass rate of 64%, meaning it successfully resisted 64% of the attacks but succumbed to the remaining 36%.

    • Fine-Tuned Model: Achieved a pass rate of 100%, successfully resisting all prompt injection attacks in the test dataset.

These results indicate that our fine-tuning process was effective in enhancing the model's resistance to the specific prompt injection patterns included in our training data.

We also built an application to perform some manual anecdotal testing and comparisons. The application implements a simple email summarization use-case, where the body of an email is fed to the LLM, which is asked to summarize it:

The prompts used by the application are provided here:

We then injected a malicious payload into the body of the email, and observed the result on the original Llama model:

Simply switching to our fine-tuned model (keeping the system prompt and email body the same), we can see how the fine-tuned model ignored the added instructions:

It's important to note that while the fine-tuned model showed improved resistance to the tested attacks, it is not immune to all forms of prompt injection. The attack surface has been reduced concerning the injection patterns present in the training dataset, but the model is most definitely still to be vulnerable to any attacks not represented in the training data. Moreover, for this model to function properly, it is necessary to block any instance of the data markers (<<<data>>> and <<</data>>>) that appears in the input.

Additionally, due to time constraints during this project, we did not evaluate whether the fine-tuning process affected the model's general language capabilities. Future work should include running standard LLM benchmarks, such as the Massive Multitask Language Understanding (MMLU) benchmark, to ensure the model has retained its original abilities across a range of tasks.

A future improvement could be to fine-tune the model not only to ignore injected instructions, as we've already done, but also to alert the user when such instructions are detected. This would involve expanding the training dataset to include samples where the model adds warnings like the following to any response that involves an injection attack:

“Warning: the original document provided seems to contain potentially malicious instructions, which have been ignored.”

We have published the fine-tuned model on Hugging Face: withsecure/Llama3-8B-PromptInjectionHardened. Additionally, we created a quantized version (Q8_0) that can be run using Ollama (https://ollama.com/withsecure/llama3-8b-prompt-injection). Quantization, performed using llama.cpp, is a process that reduces the model's size by compressing its weights. This allows for more efficient inference on hardware with limited resources without significantly impacting performance.

The following example shows how to load and use the quantized version of our fine-tuned model locally in Ollama:

1. Benchmarking and Defending Against Indirect Prompt Injection Attacks on Large Language Models, https://arxiv.org/abs/2312.14197
2. The Instruction Hierarchy:Training LLMs to Prioritize Privileged Instructions, https://arxiv.org/html/2404.13208v1 
3. Should you let ChatGPT control your browser?, https://labs.withsecure.com/publications/browser-agents-llm-prompt-injection
4. When your AI Assistant has an evil twin (Google Gemini Prompt Injection), https://labs.withsecure.com/publications/gemini-prompt-injection 
5. LLM01:2023 - Prompt Injections. OWASP Top 10 for Large Language Model Applications, https://owasp.org/www-project-top-10-for-large-language-model-applications/Archive/0_1_vulns/Prompt_Injection.html
   
6. OWASP Top 10 for Large Language Model Applications, https://owasp.org/www-project-top-10-for-large-language-model-applications