PaperlayerView exact revision on arXiv

Source-linked AI summary

Get my drift? Catching LLM Task Drift with Activation Deltas

Sahar Abdelnabi, Aideen Fay, Giovanni Cherubin, Ahmed Salem, Mario Fritz, Andrew Paverd

arXiv:2406.00799v6cs.CRcs.CLcs.CY

TL;DR

External data can inject instructions that shift an LLM away from the user’s task, and the paper studies whether activation changes can detect that drift. Its probes achieve over 0.99 ROC AUC on an out-of-distribution test set spanning unseen domains and attacks, and the authors release TaskTracker for further research.

  • Problem

    LLMs may treat instructions embedded in external data as executable, allowing task drift that can compromise security boundaries.

    Source: §I · p. 1

  • Method

    The paper detects task drift by contrasting model activations before and after processing external data and trains probes on activation deltas.

    Source: §I · p. 2, C. Catching the Drift · p. 5

  • Results

    Both probing methods exceed 0.99 ROC AUC on six language models and an out-of-distribution test spanning jailbreaks, malicious instructions, and unseen task domains and styles.

    Source: §I · p. 2

  • Takeaways & Limitations

    The defense does not require fine-tuning or sacrifice conversational utility, and TaskTracker provides resources for further task-inspection and interpretability research.

    Source: §I · p. 2, §VII · p. 13

    The authors note that future attacks may evade activation deltas, although their two adaptive attacks produced larger embedding distances than clean instances.

    Source: B. Limitations and Future Directions · p. 12

Abstract

from arXiv · show

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.

I. INTRODUCTION

External instructions can divert an LLM from a user’s task because models do not reliably distinguish data from instructions, creating a source-based security risk. The paper proposes detecting this task drift through activation changes and reports strong out-of-distribution performance, alongside the TaskTracker release.

I. INTRODUCTION

  • Task drift is deviation induced by external data; its security significance lies in instruction source because LLMs do not reliably distinguish data from instructions.The paper contrasts this source-based concern with traditional jailbreaks, which focus on malicious intent.Source: §I · p. 1

II. RELATED WORK

  • Prior mitigations include prompt-based solutions, task-specific models, and fine-tuning; some proposals require altering models, while task-specific models reduce general-purpose scope.The paper presents its activation-based defense as not requiring model alteration.Source: §I · p. 1, §II · p. 3

I. INTRODUCTION

  • The method compares last-token activations before and after external text to detect drift without relying on model output or modifying the LLM.The authors describe activation changes as leaving a trace of inserted instructions in the model’s activations.Source: §I · p. 2
  • Both probing methods exceed 0.99 ROC AUC on six language models and an out-of-distribution test spanning jailbreaks, malicious instructions, and unseen task domains and styles, outperforming Prompt Guard and Prompt Shields.The probes were not trained on the attacks in that test set.Source: §I · p. 2
  • TaskTracker releases a dataset, activations from six SoTA models, and probes to support further research on task inspection and interpretability.The evaluation dataset contains more than 500K examples.Source: §I · p. 2

II. RELATED WORK

  • Unlike fixed injection templates or trigger-specific backdoor detection, this approach examines task shifts in individual interactions and adapts to diverse injection forms.Its comparative probing starts from an arbitrary initial task at inference time as an anchor.Source: §II · p. 3

III. PRELIMINARIES: THREAT MODEL AND DATASET CONSTRUCTION

The threat model assumes a trusted user task followed by third-party data in a general-purpose retrieval-augmented application. The dataset construction varies tasks, data sources, injected instructions, and attack styles to test generalization beyond training conditions.

A. Threat Model and Problem Setup

  • The threat model assumes a general-purpose retrieval-augmented application, a trusted user task, and no closed set of user tasks or constraints on system prompts.Tasks may include questions, NLP tasks, combinations, or appended meta-prompts.Source: A. Threat Model and Problem Setup · p. 3
  • Poisoned data blocks contain an injected task decomposed into a trigger and payload; following any such instruction can constitute drift, whether benign or malicious.The trigger builds toward the payload, which contains the actual instruction.Source: A. Threat Model and Problem Setup · p. 3
  • The assumed input format is a primary task followed by data blocks; a classifier call may be skipped for follow-up questions that retrieve no new data.The authors note that multi-turn conversations can be adapted by formatting each turn in this sequence.Source: A. Threat Model and Problem Setup · pp. 3–4

B. Training Dataset Construction

  • Training examples pair SQuAD text blocks with sampled questions or synthetic NLP tasks, while generated triggers and payloads vary in wording and position within the block.Payloads draw on SEP questions and Alpaca instructions; injected tasks are placed at the beginning, middle, or end.Source: B. Training Dataset Construction · p. 4

C. Generalization Checks

  • Validation and test data vary block sources and lengths, primary tasks, payload domains, and trigger styles; jailbreaks and malicious instructions are deliberately excluded from training.Test payloads include Code Alpaca and several jailbreak and malicious-instruction datasets, with some tests omitting special triggers.Source: C. Generalization Checks · p. 4

IV. TASK SCANNING AND TRACKING VIA ACTIVATIONS

The section collects last-token activations to represent the task before and after external data is processed. Its eliciting template primes attention to instructions, though high classification performance may not require that template.

A. Activations Collection

  • The eliciting template prompts the model to list the requests it received and will execute next, aiming to focus activations on instructions.The template’s stated purpose is to increase the signal-to-noise ratio in the extracted activations.Source: A. Activations Collection · p. 4
  • For each instance, the method extracts last-token activations across layers after the primary task alone and after the full instance.The full instance can be clean, with no injected task, or poisoned, with an injected task.Source: A. Activations Collection · p. 4
  • The eliciting template may not be necessary for high classification performance, suggesting task drift can be represented without this priming step.The paper reports this finding in Appendix G.Source: A. Activations Collection · p. 4
  • Table I varies primary-task type and number, data blocks, and injection tasks, with each row representing one combination.The table covers training datasets and validation and test generalization checks.Source: Table I · p. 5
  • Fig. 2 presents t-SNE visualizations of Mistral 7B task activation deltas across layers.Source: Figure 2 · p. 5

B. Activations Enable Task Drift Representation

The paper represents drift through changes in activations and probes those changes with linear classification or metric learning. The probes achieve strong results across models and generalize to unseen test cases.

B. Activations Enable Task Drift Representation

  • Task activation deltas subtract primary-task activations from full-instance activations to represent changes associated with processing external data.The paper hypothesizes that contrasting clean and poisoned instances against the primary task can capture drift.Source: B. Activations Enable Task Drift… · p. 5
  • Intermediate-to-deeper Mistral 7B layers show clear separation between clean and poisoned activation deltas in t-SNE visualizations.Source: B. Activations Enable Task Drift… · p. 5

C. Catching the Drift

  • Linear classification separates clean from poisoned deltas, while metric learning trains triplet embeddings so primary and clean instances are closer than primary and poisoned instances.At test time, metric learning classifies by thresholding the distance between primary-task and full-instance embeddings.Source: C. Catching the Drift · p. 5, Figure 3 · p. 6

A. Probes Can Effectively Detect Task Drift

  • Linear probes achieved ROC-AUC ≥0.99 across all models, while metric-learning probing performed consistently high (≥0.98), was robust to layer ranges in many models, and generalized well to unseen test cases.For Llama-3 70B, later layers gave weaker signals for both probes.Source: A. Probes Can Effectively Detect Task… · p. 6

V. EXPERIMENTAL EVALUATION

  • The evaluation covers six instruction-tuned models spanning different model families and sizes.The probes are trained per model and evaluated with ROC-AUC on balanced test data.Source: §V · p. 6

1) Performance Breakdown w.r.t. Dataset Properties:

The probes distinguish clean from poisoned examples across several dataset properties, although distances vary with injection placement, task composition, and instruction length. They also retain performance on contextual and naturally phrased examples.

1) Performance Breakdown w.r.t. Dataset Properties:

  • Poisoned examples retain higher distances than clean ones across dataset properties, with end-placed injections producing higher distances than other locations.Mistral 7B shows relatively smaller distances for mixed-task poisoned examples than QA-poisoned examples; other models are more consistent.Source: 1) Performance Breakdown w.r.t.… · p. 6
  • The clean–poisoned separation persists for one-task and mixed primary tasks, despite those task types forming distinct embedding neighborhoods.The training data contained only one task per example; Mistral 7B has relatively smaller distances for mixed-task poisoned examples than QA-poisoned examples.Source: 1) Performance Breakdown w.r.t.… · p. 6
  • Jailbreak injections have lower distances for some models, which the authors hypothesize may reflect their longer instruction spans diluting the contrastive signal.They note that retrieval-ranking constraints may limit the size of injected instructions in practice.Source: 1) Performance Breakdown w.r.t.… · pp. 6–7
  • Distances are comparable whether injected tasks were executed or not; without a strong trigger, they can decrease but remain above clean-example distances.The authors frame detection as identifying drift in perceived instructions rather than determining attack success.Source: 1) Performance Breakdown w.r.t.… · pp. 7–8, Table VI · p. 7
  • Probes still perform highly when primary and injected questions concern the same paragraph, and higher layers remain strong when clean questions are phrased naturally.The latter comparison uses poisoned examples phrased as instructions and clean examples containing the same questions as natural text.Source: 1) Performance Breakdown w.r.t.… · pp. 8–9, Table VIII · p. 9, Table VII · p. 9

2) Tracking the Distances To Enable Locating Injections:

The probes track growing distances as injected instructions are read, need relatively little training data, and compare favorably with classifiers on the reported benchmarks. On a challenging non-public set, their comparison with baselines depends on retraining.

3) Size of Training Data:

  • With 5% of the original training data (40K instances), the linear probe achieves AUC 0.9986, with slight variations as training data increases to 75%.The authors note this may mean only small amounts of application data are needed for deployment.Source: 3) Size of Training Data · p. 9

5) Generalization to “even more” Challenging Datasets:

  • The challenging non-public test set includes long documents, multilingual text, varied encodings, subtle injection phrasing, tool use, and application-specific data.The dataset contains approximately 13K instances and cannot be released through TaskTracker.Source: 5) Generalization to “even more”… · pp. 9–10

4) Comparison to Classifiers:

  • On the collected test set, the probe trained on benign data outperforms Prompt Guard and Prompt Shields in a zero-shot attack setting.The compared classifiers were trained on attacks, according to the authors’ account of Prompt Guard’s model card and the Prompt Shields comparison.Source: 4) Comparison to Classifiers · p. 9, Table X · p. 10

2) Tracking the Distances To Enable Locating Injections:

  • Distances gradually increase as more tokens of an injected task are ingested, both in an example and averaged across 2,300 poisoned examples.Clean-example distances fluctuate within a lower range; the trend suggests last-token activations are useful in the normal operation mode.Source: 2) Tracking the Distances To Enable… · p. 9, Figure 7 · p. 10

5) Generalization to “even more” Challenging Datasets:

  • On the challenging non-public test set, the probe is close to Prompt Shields without retraining and outperforms it after retraining on the same data.Without retraining, it also significantly outperforms Prompt Guard; the comparison is reported for Phi-3 14B.Source: Table XI · p. 10, 5) Generalization to “even more”… · p. 10

VI. DISCUSSION

The discussion argues that instruction detection from text alone faces generalization challenges because instructions vary and models respond differently across languages and contexts. Activation-based detection is presented as sensitive to instructions a model is likely to attend to and as compatible with meta-prompts.

A. Why Use Activations?

  • Instruction detectors face generalization challenges because instructions vary widely, and behavior may differ across languages; translated injected tasks still have higher distances than clean examples.In the reported non-English test, the injected tasks were translated into German, French, Italian, Spanish, and Arabic while the primary task and text remained English.Source: A. Why Use Activations? · p. 10, Table XII · p. 11
  • Which instructions are relevant is model-specific and contextual, making relevance difficult to detect from natural language alone.The paper gives differing model ability to decode base64 and its in-context injection experiment as examples.Source: A. Why Use Activations? · p. 11
  • Embedding distances decrease as meta-prompts more strongly mark data and instruct the model to ignore conversation instructions, reaching the clean-example baseline.The authors interpret this as evidence that detection reflects instructions the model is likely to attend to, and report compatibility with meta-prompt defenses.Source: Table XIII · p. 11, A. Why Use Activations? · p. 11
  • The embedding model still detects delimiter-based evasion attempts, while the paper notes that text-output task elicitation can yield generic or hallucinated instructions.The authors say their method uses no text generation and can be used alongside other mitigations.Source: Table XIV · p. 11, A. Why Use Activations? · p. 11

B. Limitations and Future Directions

The authors examine activation-based drift detection’s robustness, limits, and open questions, including how activations represent tasks and how attacks might evade detection. They conclude that activation deltas detect drift across models and new tasks without fine-tuning.

B. Limitations and Future Directions

  • Without a task-eliciting prompt, individual-layer probes on Mistral retained strong performance, but the best activation-collection strategy across models remains unstudied.The authors also identify adversarial robustness across collection strategies as an open question.Source: B. Limitations and Future Directions · p. 11
  • The activations’ underlying signal remains unclear: they may reflect semantic task content or syntactic and positional features, with early-layer detection suggesting the latter.The authors call for more systematic study of why activation-based detection works.Source: B. Limitations and Future Directions · p. 12
  • The tested lie prompts did not elicit deceptive responses from open-source models, leaving open whether genuinely deceptive answers could reduce activation deltas.The authors note that prior work detecting untruthfulness in internal states may suggest such manipulation is unlikely.Source: B. Limitations and Future Directions · p. 12
  • Both tested adaptive attacks produced larger embedding distances than clean instances, so neither evaded the defense.The attacks framed injected instructions as steps toward the primary task or told the model to lie; Table XV summarizes the comparisons.Source: Table XV · p. 12, B. Limitations and Future Directions · p. 12
  • Attackers may need attacks that work across users’ unknown primary tasks; very long injected instructions may lower distances, though query relevance may constrain them.The authors argue that attacker constraints in retrieval-augmented applications merit further study.Source: B. Limitations and Future Directions · p. 12
  • Probe transfer across fine-tuned model versions was not evaluated; the authors suggest using a stable, potentially smaller proxy model for detection.The proxy would be separate from the model used for the downstream application.Source: B. Limitations and Future Directions · p. 13

VII. CONCLUSION

  • The conclusion presents activation deltas as a way to detect drift across six model families and sizes and new tasks and attacks, without fine-tuning.The authors report releasing TaskTracker with a dataset, raw activations, and inspection tools.Source: §VII · p. 13

APPENDIX

The appendix describes the data used to train and evaluate the probes, metric-learning design and sampling, and compute requirements. It also reports the resources used for activation collection.

APPENDIX

  • Activation generation took around 36 to 96 hours per model, while metric-learning probe training took around 12 hours on one A100 80GB GPU.The authors plan to release activations to support reproducibility and follow-up research.Source: APPENDIX · p. 15, Table XVI · p. 15

A. Training Dataset

  • Training paired SQuAD text blocks and tasks with injected instructions yielded 418,110 clean and 418,110 poisoned samples.Injected tasks came from Alpaca, GPT-4 triggers, handwritten trivia, and generic NLP prompts; each task combination was repeated at the beginning, middle, and end of the text.Source: A. Training Dataset · p. 15

B. Validation and Test Dataset

  • Validation and test sets combine clean and poisoned examples with varied primary tasks and injection locations; the test set contains 31,134 instances of each class.Test injections draw on Code Alpaca, BeaverTails, Do-Not-Answer, and JailbreakBench; HotPotQA was not used to train the probes.Source: B. Validation and Test Dataset · p. 15

C. Training Details of Metric Learning Probes

  • Metric learning places primary-task and clean activations closer than primary-task and poisoned activations, using triplet loss and online hard or semi-hard negative mining.Warm-up with semi-hard sampling followed by hard sampling improved generalization over naive triplet sampling, which plateaued and failed to generalize.Source: C. Training Details of Metric Learning… · pp. 15–16
  • The metric-learning model processes each layer with lightweight convolutional subnetworks, concatenates their outputs, and normalizes a 1024-dimensional embedding.The authors choose convolutions instead of linear layers to keep the model lightweight.Source: C. Training Details of Metric Learning… · p. 15

D. Hyperparameters

This section gives the metric-learning training settings, explains test-time classification by embedding distance, and summarizes probe ROC curves across models and layers.

D. Hyperparameters

  • The metric-learning probe used margin α=0.3, batch sizes 2500 for mining and 1024 for training, and semi-hard sampling before hard mining after 3000 steps.Training ran for five epochs with Adam at learning rate 0.0005, decayed by 0.05 every 800 steps.Source: D. Hyperparameters · p. 16

E. Classification

  • At test time, the metric-learning probe compares primary-task and full-instance embeddings by L2 distance, then thresholds that distance to classify examples as clean or poisoned.Source: E. Classification · p. 16

F. ROC Curves of Probes

  • The linear classifier performs best from the halfway point of model depth onward and worse in earlier layers, especially layer 0; the metric-learning ROC curves span different LLMs.Figures 9 and 8 show the linear-probe layer curves and metric-learning curves across models, respectively.Source: F. ROC Curves of Probes · p. 16, Figure 8 · p. 16, Figure 9 · p. 17

G. Results of Training Probes on Activations Without a Prompt Template

The appendix evaluates probes using activations collected without the task-eliciting prompt, reporting performance patterns, computational implications, and an execution-status comparison. The evidence is primarily from Mistral 7B.

G. Results of Training Probes on Activations Without a Prompt Template

  • Without the task-eliciting prompt, the linear probe remained strong across most layers, while metric learning dropped with the last 15 layers but retained signal on individual layers 7 and 15.The no-template t-SNE visualization shows more segmentation and noise than the template-based visualization, which may explain why triplet-network training did not converge.Source: G. Results of Training Probes on… · p. 18, Table XVIII · p. 19, Table XVII · p. 19
  • For Mistral, both probes appeared to perform strongly without the prompt template, though metric learning may need more layer hyperparameter search.Which activation-collection strategy performs better across models and offers better adversarial robustness remains unstudied.Source: G. Results of Training Probes on… · p. 18
  • Collecting activations without a specific template avoids extra forward passes over the context at test time, and probe classification adds negligible cost.The model’s primary-task and full-instance activations can be collected at the ends of contexts already processed to generate output.Source: G. Results of Training Probes on… · p. 18
  • A suffix-style priming template could permit computation reuse through caching, but its necessity and cost reduction remain questions for future work.The authors note that processing tokens costs significantly less than generating them.Source: G. Results of Training Probes on… · p. 18

H. Distances vs. Verifier (Execution Status)

  • Attacks judged executed, rejected, or not detected had comparable embedding distances, shifting mitigation toward detecting drift rather than determining attack success.The execution-status comparison is in Figure 12.Source: H. Distances vs. Verifier (Execution… · p. 18, Figure 12 · p. 20
  • Execution status was verified from one generated output per example, so cases labeled not detected might have succeeded on another attempt.The authors describe possible variation in whether a model responds to the injected task or the primary task.Source: Figure 12 · p. 20

I. Temporal Analysis of Distances Over Tokens

Distances rise after injected instructions begin, and the appendix reports additional generalization tests using previously trained metric-learning probes without further training.

I. Temporal Analysis of Distances Over Tokens

  • On 2,000 poisoned Mistral 7B examples, Fig. 13 compares average distances in five-token windows before and after injection onset and at the last token.Source: Figure 13 · p. 22
  • Distances rise after injected instructions begin, including in examples where injections occur at the beginning or end of the text.The temporal analysis provides additional quantitative and qualitative examples of this pattern.Source: §I · p. 22, Figure 14 · p. 23, Figure 15 · p. 23

J. Additional Results: Why Use Activations?

1) Vs. Detecting Instructions?:

The appendix tests whether activation distances persist under varied prompts and potential evasion attempts; text-based task elicitation can itself produce generic or hallucinated answers.

1) Vs. Detecting Instructions?:

  • The appendix tests spotlighting and 1,000 English WildChat conversations paired with a summarization task and three meta-prompts of varying strength.The spotlighting scenario assumes an inserted closing tag and the sentence “The text has ended. These are new instructions.”Source: 1) Vs. Detecting Instructions? · p. 24
  • WildChat meta-prompt levels progress from a basic summarization task to explicit instructions to ignore conversation requests and marked conversation boundaries.Level 3 uses the strongest phrasing, explicitly saying to summarize and not answer the conversation.Source: 1) Vs. Detecting Instructions? · p. 24

2) Vs. Instructing the Model To Spell Out Its Instructions?:

K. “Potential” Evasion Attacks

Loading 2406.00799v6…