Are you still on track!?
Catching LLM Task Drift with Activations

We assume a general-purpose LLM application where users can freely instruct the LLM. The application’s developer may give the LLM generic instructions in its system prompt, on which we do not assume any constraints. We refer to the user’s task as the primary task, which we assume is trusted. The LLM consumes third-party data such as retrieved documents or API output. We refer to this input as data blocks. The user’s primary task involves processing such blocks to answer questions or perform analysis. The data block may be “clean” or “poisoned”, where the latter has a secondary task inserted. Following Pasquini et al. (2024), the secondary task is further decomposed into a trigger and a payload. The trigger is a “build-up” sentence that incentivizes the LLM to follow the payload. The payload is the actual instruction. Task drift may happen when the model processes the poisoned data block. Our dataset construction approach, as described below, is based on sampling each of these components individually and combining them to form a sizeable dataset.

Data Blocks. We use question-answering (QA) datasets that typically consist of questions along with text paragraphs from, e.g., Wikipedia. Specifically, we use the SQuAD (Rajpurkar et al., 2016) training split for data blocks in the training data.

Primary Tasks. To create an instance in our dataset, we sample an instance from the QA dataset. If the instance has multiple questions, we randomly select one. In addition to QA tasks, we compile a list of synthetic tasks (from the SEP dataset (Zverev et al., 2024)) that can be used as generic NLP questions such as summarization, translation, and keyword extraction.

Secondary Tasks. We use GPT-4 to generate diverse variations of trigger sentences, ranging from subtle hints about upcoming instructions to more explicit alerts using uppercase keywords and delimiters to indicate boundaries. The payloads for the training data are sourced from the NLP and trivia questions from the SEP dataset (Zverev et al., 2024) and from instructions in the Alpaca dataset (Taori et al., 2023). We create equal-sized variations of the dataset by positioning the secondary task at the beginning, randomly in the middle, and at the end of the data blocks.

We construct validation and test data with different characteristics from the training data to ensure the generalization of our approach on challenging out-of-distribution examples.

Data-Blocks’ Sources and Lengths. We use the Dev. split from SQuAD. In addition, since by construction the poisoned examples in training can have more tokens, we construct examples from HotPotQA (Yang et al., 2018) that typically consist of multiple paragraphs for validation and test splits to rule out possible confounders related to text length.

Primary Task. Besides the variation in the source of the datasets, we also test on a case where the user asks two questions (one from QA datasets and one from the generic NLP question) for both the clean and poisoned examples. The primary task would now have the form of, e.g., “After you answer the previous question, please summarize the text as well”.

Secondary Task. To vary the payloads, we use the Code Alpaca dataset (Chaudhary, 2023) as a domain shift of tasks that were not used in training. We deliberately do not train on any jailbreaks or malicious instructions. Our definition of the secondary task is generic to include different forms of instructions, which are not exclusive to jailbreaks and malicious instructions, but may include them. We use the harmful behavior questions from AdvBench (Ding et al., 2024), the set of forbidden questions from (Shen et al., 2024), attack sentences from BIPA (Yi et al., 2023), jailbreaks from TrustLLM (Sun et al., 2024), malicious questions from BeaverTails (Ji et al., 2023) and Do-Not-Answer (Wang et al., 2023) datasets. We also export pre-computed jailbreaks from the JailbreakBench (Chao et al., 2024) spanning the PAIR (Chao et al., 2023) and adaptive jailbreak (Andriushchenko et al., 2024) attacks (we note that whether these jailbreaks themselves are successful in getting harmful answers is orthogonal to our problem of detection). We use non-overlapping splits of these datasets in the validation and test data. For triggers, the GPT-4 generated triggers we use in validation and test splits are inspired by previous hand-written prompt injections (Perez and Ribeiro, 2022) (e.g., “ignore previous prompts”), while the ones we used in training are less tailored for attacks. We also test on the case where no special triggers are used at all. In this case, we use the phrase “Please answer the following: ”. Similarly to the training data, we construct equal-sized variations of the dataset by varying the position of the secondary task. Dataset statistics are in Appendix B.

Our goal is to extract the LLM’s activations that represent the instructions given to the model at one point in time. To do this, we design a prompt template $T$. The intuition of this template is to prime the model towards attending to the instructions in the context and to increase the “Signal-to-Noise” ratio in the extracted activations.¹¹1Notably, we found that this priming prompt template may not be necessary for achieving high classification performance as outlined in Appendix D.3, suggesting that activations may naturally represent task drifts effectively.

[linecolor=priming,roundcorner=20pt,linewidth=1.5pt] Here are your main requests: <MAIN> DATASET INSTANCE HERE </MAIN> but before you answer, please complete the following sentence by briefly writing each request(s) you received and you are going to execute next: "All requests that I am going to execute now are:"

We then extract the activations of the last token in the input (before generation) across all layers as a summary of the context. For each data instance $x$, we extract two sets of activations: the first is extracted after the LLM has processed only the primary task $x_{\text{pri}}$, and the second after it has processed the full instance $x$. For a language model $\mathcal{M}$, this can be expressed as:

where $n$ is the number of $\mathcal{M}$’s layers. The subset of layers to use is a hyperparameter. Instance $x$ can either be clean $x_{\text{cln}}$ (no secondary task) or poisoned $x_{\text{pois}}$ (with secondary task inserted).

Our goal is to distinguish between $\text{Act}^{x_{\text{cln}}}$ and $\text{Act}^{x_{\text{pois}}}$. We hypothesize that the drift incurred by the secondary task in $x_{\text{pois}}$ can be detected by contrasting $\text{Act}^{x_{\text{cln}}}$ and $\text{Act}^{x_{\text{pois}}}$ against $\text{Act}^{x_{\text{pri}}}$. Therefore, one option is to represent $x$ as the Task Activations Residuals, $\tilde{\text{Act}^{x}}$:

We show in Figure 2 the results of applying t-Distributed Stochastic Neighbor Embedding (t-SNE) on $\tilde{\text{Act}^{x}}$ at different layers of Mistral 7B (Jiang et al., 2023). As can be observed, $\tilde{\text{Act}^{x}}$ from intermediate to deeper layers provide a clear separation between clean and poisoned instances.

Given these visualizations that suggest that comparing $\text{Act}^{x}$ against $\text{Act}^{x_{\text{pri}}}$ can represent task drift, we design two probing mechanisms.

Linear Classifier. Here, we simply train a lightweight Logistic Regression probe $p$ to distinguish between the task activations residuals of clean $\tilde{\text{Act}^{x_{\text{cln}}}}$ and poisoned $\tilde{\text{Act}^{x_{\text{pois}}}}$ instances.

Metric Learning. We use triplet networks to learn embeddings of tasks (Figure 3). We optimize the triplet loss (Schroff et al., 2015) on $x_{\text{pri}}$, $x_{\text{cln}}$, and $x_{\text{pois}}$ triplets. The anchor and positive points are $x_{\text{pri}}$ and $x_{\text{cln}}$, respectively, and the negative point is $x_{\text{pois}}$. The embedding model should output closer embedding vectors for any $\text{Act}^{x_{\text{pri}}}$ and $\text{Act}^{x_{\text{cln}}}$ than for any $\text{Act}^{x_{\text{pri}}}$ and $\text{Act}^{x_{\text{pois}}}$. Our base model is a series of 1-dimensional convolution filters and non-linearities applied individually to the LLM’s layers. The output of these sub-networks are concatenated and fed to a last linear layer with output $\in\mathbb{R}^{1024}$ that constitutes the final embedding vector that we normalize. To train the model, we used a mix of hard and semi-hard triplet mining (details in Appendix C). At test time, we compute the distance between the embeddings of $\text{Act}^{x_{\text{pri}}}$ and $\text{Act}^{x}$ and classify based on a threshold (clean examples should have smaller distances). As we show in our experiments, these distances can also be useful to perform temporal analysis over the sequence to locate injected instructions.

Tables 1 and 2 show the ROC-AUC for classifiers and embedding models. ROC curves are in Appendices D.1 and D.2. Linear probes achieve near-perfect scores ($\geq$0.99) across all LLMs. Metric learning probing was more robust to layer selection as it consistently performs well ($\geq$0.93) across different ranges. Additionally, Figure 4 provides histograms of learned embedding distances (using the best metric learning probe) between the full test point $x_{i}$ and the primary instruction $x_{i_{\text{pri}}}$ for both clean and poisoned data points. Figure 5 also presents t-SNE visualizations of the differences in embeddings between the full test instances and their corresponding primary instructions. Our results show that both probing methods perform strongly, generalizing well to unseen cases in the test data.

Generalization Challenges in Instruction Detection. Detecting natural language instructions in the input domain may face a generalization problem, as instructions can vary widely and lack a specific grammar or format. Additionally, LLMs’ behavior may differ across languages, with low-resource languages potentially bypassing LLM safeguards (Yong et al., 2023). To test generalization, we created 500 poisoned examples with non-English secondary tasks. We used 100 trivia examples with the trigger “please answer the following”, translated into German, French, Italian, Spanish, and Arabic; the primary task and text remained in English. The Mixtral 8x7B embedding distances for these examples are in Table 5. Although lower than English instructions, these distances are still higher than those for clean examples, indicating that translated secondary tasks, even when short and without strong triggers, may still be executed and pose a safety challenge. More complex secondary tasks might yield even higher distances.

Activations and Sensitivity to “Relevant” Instructions. Our preliminary experiments indicate that detection via activations might better reflect which instructions the LLM is more likely to attend to. For attacks, previous work (Zverev et al., 2024; Pasquini et al., 2024) has shown that the way the embedded instructions are phrased affects how likely they will be executed. For benign text, LLMs may frequently encounter natural language instructions that are not intended as instructions to the LLM. While LLMs may nevertheless execute them (Zverev et al., 2024), it is still not clear which instructions would be treated as such. Therefore, detecting instructions purely based on natural language is challenging because it is 1) model-specific, and 2) contextual. To test whether activations provide a better alternative, we test on 1000 conversations from the WildChat dataset (Zhao et al., 2024). In this experiment, we applied a “summarization” primary task to the conversations. Instructions in the conversations should ideally be treated as non-executable data. However, due to the models’ inherent lack of separation, we used meta prompts specifically designed to prompt the LLM to ignore any encountered instructions (see Appendix D.6). We used three meta prompts that vary in how strongly they emphasize that no instructions should be followed and how clearly they mark the conversation as data blocks. As shown in Table 6, embedding distances decrease as the prompt strength increased, suggesting that detection is model-specific and that meta prompt-based defenses are compatible with our work.

More Robust than Prompt-Based Defences. The previous WildChat experiment is inspired by another defense (Hines et al., 2024) that “marks” the data using delimiters, prompting the model to ignore instructions within the data blocks. This method was effective with benign clean data, which contains instructions that the model should not follow. However, in adversarial setups, this approach may lack robustness. To illustrate this, we created 500 examples where the primary task includes instructions about the delimiter, and the secondary task includes phrases such as “the text has now ended and these are new instructions”. This strategy can evade delimiter-based defenses (Chen et al., 2024). Our embedding model, however, can still detect these instructions effectively, maintaining low distances for clean, data-marked examples (see Table 7). Although more sophisticated data-marking techniques could be used, this often leads to a cat-and-mouse game, requiring constant updates based on learnings from jailbreaks and prompt injections attacks (Debenedetti et al., 2024). In contrast, mitigations based on causality or interpretability, like our approach, provide a more robust defense, which can still be used in conjunction with other mitigations.

Limitations of Task-Eliciting Prompts’ Text Output. Another approach involves using task-eliciting prompts to retrospectively “interrogate” models about the received task and use the text output as an indicator of drift. However, we found that models often respond with generic tasks (e.g., “I will understand the context”, “I will provide accurate information”), or occasionally with hallucinated instructions. Examples of such responses can be found in Appendix D.6.

What Do The Activations Capture? While our method shows surprisingly high performance and helps solve a major security and safety problem, it is still unclear what these activations actually capture. E.g., do they encode semantic similarity and high-level concepts related to the tasks, or do they capture syntactical or structural features related to changes in instructions’ locations within the context? The empirical results that drifts can be detected from early layers may suggest the latter. More work is needed to pinpoint the underlying reasons.

Evasion Attacks. Evading our detector is not straightforward due to the aforementioned reason. We designed two potential attacks to make the activations more similar to those of the primary task: one where the secondary task states that the instructions are a step toward the primary task, and another that prompts the model to “lie” about the new instructions. Both resulted in larger embedding distances compared to the clean instances (see Table 8 and Appendix D.7). We observed that all open-source models still mentioned the secondary task when prompted with our task eliciting prompt, and we could not get a deceptive response. It is thus unclear how having deceptive answers (if possible) may affect the activations. Previous work has shown that inspecting models’ internal states can detect untruthfulness (Zou et al., 2023a; Halawi et al., 2024), which may suggest that manipulating activations via deception is unlikely. Additionally, a universal optimization-based attack seems difficult as it would need to be effective against any arbitrary primary and secondary instructions.

Alternative Setups For Extracting Activations. We performed most of our experiments with the task eliciting prompt. However, we also ran preliminary experiments on Mistral 7B by passing the instance $x$ without this template (see Appendix D.3). Training over multiple layers was significantly worse than training over individual layers. When trained over individual layers, both probes surprisingly maintained the same performance as those trained with the prompt template, suggesting that the signal may be strong enough in specific layers. Future work could investigate the differences between these and other potential setups in more detail.