Get my drift? Catching LLM Task Drift with Activation Deltas

TL;DR

The paper tackles prompt-injection-induced task drift in retrieval-augmented LLM applications by introducing activation-delta representations derived from last-token activations. It develops two lightweight probes (linear classifier and triplet-based metric learning) to detect drift without modifying the LLM, and validates them on a large, diverse synthetic dataset spanning six SoTA models, achieving near-perfect ROC-AUC on out-of-distribution attacks. Key contributions include a scalable data synthesis pipeline, cross-model generalization, and a public TaskTracker toolkit with activation data and inspection tools to advance post-hoc task decoding and interpretability. Practically, the approach offers a cost-efficient, deployment-friendly defense that complements existing prompt-based defenses and supports robust, monitorable safeguards for LLM-based RAG systems across domains and languages.

Abstract

LLMs are commonly used in retrieval-augmented applications to execute user instructions based on data from external sources. For example, modern search engines use LLMs to answer queries based on relevant search results; email plugins summarize emails by processing their content through an LLM. However, the potentially untrusted provenance of these data sources can lead to prompt injection attacks, where the LLM is manipulated by natural language instructions embedded in the external data, causing it to deviate from the user's original instruction(s). We define this deviation as task drift. Task drift is a significant concern as it allows attackers to exfiltrate data or influence the LLM's output for other users. We study LLM activations as a solution to detect task drift, showing that activation deltas - the difference in activations before and after processing external data - are strongly correlated with this phenomenon. Through two probing methods, we demonstrate that a simple linear classifier can detect drift with near-perfect ROC AUC on an out-of-distribution test set. We evaluate these methods by making minimal assumptions about how users' tasks, system prompts, and attacks can be phrased. We observe that this approach generalizes surprisingly well to unseen task domains, such as prompt injections, jailbreaks, and malicious instructions, without being trained on any of these attacks. Interestingly, the fact that this solution does not require any modifications to the LLM (e.g., fine-tuning), as well as its compatibility with existing meta-prompting solutions, makes it cost-efficient and easy to deploy. To encourage further research on activation-based task inspection, decoding, and interpretability, we release our large-scale TaskTracker toolkit, featuring a dataset of over 500K instances, representations from six SoTA language models, and a suite of inspection tools.

Figures (16)

• Figure 1: In LLM applications, instructions can flow from poisoned (ideally "data-only") sources, enabling attacks. In this example, a user asks a question (step 1) upon which external search results are retrieved (step 2), which contain an injected task (i.e., prompt injection). We propose to catch (and potentially locate) the LLM's drift from the initially given user's task via contrasting the LLM's activations before (step 1) and after (step 2) feeding the external data, computing activation deltas. The plot shows the activation deltas before and after the onset of the injected task.
• Figure 2: t-SNE visualizations of the task activation deltas of Mistral 7B across different layers.
• Figure 3: Details of the metric learning probing method. We first extract the activations of the primary, clean, and poisoned instances in the dataset when wrapped within our eliciting prompt$T$. We read the activations of the last token in the input context sequence across the LLM's layers. We next train a triplet embedding model with repeated triplet mining. After training, the embedding distances between $\text{Act}^{x_{i_\text{pri}}}$ and $\text{Act}^{x_i}$ should be closer when $x_i$ is a clean point.
• Figure 4: Histogram of embedding distances between $x_i$ and $x_{i_\text{pri}}$ in the case of clean and poisoned data points for the best embedding model trained on the activations of different LLMs.
• Figure 5: t-SNE visualization of the embeddings from metric learning probes, showing they have learned meaningful representations from the activations. Each point in the visualization represents the difference in the embeddings of the full test instance $x_i$ and its corresponding primary $x_{i_\text{pri}}$.