Jatmo: Prompt Injection Defense by Task-Specific Finetuning
Julien Piet, Maha Alrashed, Chawin Sitawarin, Sizhe Chen, Zeming Wei, Elizabeth Sun, Basel Alomair, David Wagner
TL;DR
Prompt-injection attacks threaten LLM-integrated applications, and Jatmo offers a defense by building per-task, safety-focused capabilities through task-specific fine-tuning of a non-instruction-tuned base using outputs from a teacher instruction-tuned system. The method supports fully or semi-synthetic data generation via GPT-4 to bootstrap task datasets and employs a three-stage pipeline of dataset creation, labeling with a teacher, and fine-tuning, enabling security with limited labeled data. Empirically, Jatmo achieves near-teacher-quality performance on seven tasks while dramatically reducing prompt-injection success (from about 87% to 0.5%), with no runtime overhead and potential cost savings due to smaller task-specific models. Acknowledging limitations such as single-task specialization and upfront training costs, the work demonstrates a practical pathway to safer LLM deployments and provides open-source access to the approach.
Abstract
Large Language Models (LLMs) are attracting significant research attention due to their instruction-following abilities, allowing users and developers to leverage LLMs for a variety of tasks. However, LLMs are vulnerable to prompt-injection attacks: a class of attacks that hijack the model's instruction-following abilities, changing responses to prompts to undesired, possibly malicious ones. In this work, we introduce Jatmo, a method for generating task-specific models resilient to prompt-injection attacks. Jatmo leverages the fact that LLMs can only follow instructions once they have undergone instruction tuning. It harnesses a teacher instruction-tuned model to generate a task-specific dataset, which is then used to fine-tune a base model (i.e., a non-instruction-tuned model). Jatmo only needs a task prompt and a dataset of inputs for the task: it uses the teacher model to generate outputs. For situations with no pre-existing datasets, Jatmo can use a single example, or in some cases none at all, to produce a fully synthetic dataset. Our experiments on seven tasks show that Jatmo models provide similar quality of outputs on their specific task as standard LLMs, while being resilient to prompt injections. The best attacks succeeded in less than 0.5% of cases against our models, versus 87% success rate against GPT-3.5-Turbo. We release Jatmo at https://github.com/wagner-group/prompt-injection-defense.
