🤖 machine learning

When Benchmarks Lie: Evaluating Malicious Prompt Classifiers Under True Distribution Shift

This paper reveals that standard cross-validation significantly overestimates the generalization of malicious prompt classifiers by 8–16.5 AUC points compared to a rigorous Leave-One-Dataset-Out evaluation, exposing widespread dataset-dependent shortcuts in model features that current domain-generalization techniques fail to mitigate.

Original authors: Max Fomin

Published 2026-07-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Max Fomin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a security guard for a high-tech museum. Your goal is to teach the guard to spot a thief hiding a stolen artifact, no matter how they try to disguise it. In the world of artificial intelligence, these "thieves" are malicious instructions trying to trick a robot brain (called a Large Language Model) into doing something it shouldn't, like revealing secrets or writing harmful code. This is called "prompt injection." To train the guard, scientists usually show them a huge pile of photos of known thieves and known innocent visitors. If the guard gets 99% of the photos right, we usually cheer and say, "Great job! You're ready for the real world!"

But here is the catch: what if the guard isn't actually learning what a thief looks like? What if they are just learning what the photo album looks like? Maybe all the photos of thieves were taken in a red room, and all the photos of innocent people were taken in a blue room. The guard might just be checking the wall color instead of looking at the person's face. This paper is about a group of researchers who decided to test if their security guards were actually smart or just really good at memorizing the color of the training room. They built a new test where the guard has to spot a thief in a brand-new room they've never seen before, to see if they can really tell the difference between good and bad.

The Great "Room Color" Trick

The researchers, led by Max Fomin, discovered that the standard way we test AI security guards is lying to us. They found that when we test these guards using the usual method (mixing all the training photos and testing photos together), the guards look like geniuses, scoring incredibly high marks. But when the researchers used a stricter test called LODO (Leave-One-Dataset-Out), where the guard is forced to face a completely new "room" (a new dataset) they have never seen during training, the guards suddenly looked much less impressive.

In simple terms, the standard tests were like giving the guard a multiple-choice quiz where the answers were hidden in the question format. The paper shows that current security classifiers are often just memorizing "dataset shortcuts." Instead of learning the deep, complex patterns of a malicious attack, they are learning superficial clues like "this text looks like it came from the 'WildJailbreak' file, so it must be bad," or "this email looks like it came from the 'Enron' file, so it must be safe."

The Experiment: 18 Different Rooms

To prove this, the team built a massive training ground with 18 different datasets (collections of text), containing over 105,000 samples. These included everything from direct jailbreak attempts to harmless business emails. They trained security classifiers on these texts using four different AI models (including Llama, Gemma, and Qwen).

When they used the standard "mix-and-match" testing method, the classifiers scored near-perfectly, with accuracy rates often above 99%. It looked like a total victory. However, when they switched to the LODO method—training on 17 datasets and testing strictly on the 18th one they had never seen—the scores dropped dramatically.

The gap was huge. The standard tests overestimated the AI's ability to generalize by 8.0 to 16.5 percentage points. In some specific cases, the drop was even more shocking, with accuracy gaps ranging from 1% to 25% depending on the dataset. For example, a classifier that looked like a 99.5% genius on standard tests dropped to a 79.4% performer when tested on a new dataset called "Mosscap." This proved that the AI wasn't learning the concept of "bad behavior"; it was just recognizing the "style" of the training data.

The "Shortcut" Detective Work

The researchers didn't just stop at finding the gap; they wanted to know why it happened. They used a special tool called a Sparse Autoencoder (SAE), which acts like a microscope to look at the tiny internal switches (features) inside the AI's brain. They wanted to see which switches were firing when the AI made a decision.

They found that a significant chunk of the AI's "brain power" was being wasted on dataset shortcuts.

28% to 44% of the top features the AI used to make decisions were actually just shortcuts that identified the dataset, not the attack.
They even built a separate "dataset identity classifier" that could guess which dataset a piece of text came from with 96.6% accuracy just by looking at the AI's internal signals. This confirmed that the datasets were so distinct that the AI could easily tell them apart, and it was using that ability to cheat on the security test.

The paper also explored whether standard fixes could solve this. They tried common tricks like "adversarial training" (teaching the AI to fight back against tricky examples) and "reweighting" (making the AI pay more attention to hard examples). Unfortunately, none of these standard fixes closed the gap. The researchers found that the signals for "what dataset this is" and "is this a safety threat" were tangled together inside the AI's brain. Trying to remove the "dataset" signal accidentally hurt the "safety" signal, making the problem even harder to solve.

The Silver Lining: Better Explanations

While the news about the gap is serious, the paper offers a clever tool to help fix the problem in the future. The researchers developed a way to "weight" the AI's explanations. Instead of just asking, "Which features made you think this was bad?", they asked, "Which features made you think this was bad and stayed consistent across all the different rooms?"

By filtering out the features that were just "dataset shortcuts," they could generate much more reliable explanations for why an AI flagged a prompt. They found that this method changed the top 20 reasons for flagging a prompt in about 98% of cases, effectively cleaning up the AI's reasoning and removing the "cheating" logic.

The Bottom Line

This paper doesn't claim to have solved the problem of AI security. In fact, it explicitly states that they did not find a method to close the gap using current techniques. Instead, they provided a crucial reality check. They showed that the current "gold standard" for testing AI safety is flawed because it lets AI models cheat by memorizing the source of the data rather than understanding the danger.

The takeaway for anyone building or using AI agents is clear: if your security guard scores 99% on a test where the training and testing data come from the same pile, you can't trust them to handle a real-world attack from a new source. You need to test them on data they have never seen before to know if they are truly safe. The authors have released their tools so that others can run these stricter tests, ensuring that future AI guards are actually learning to spot the thieves, not just the color of the room.

Technical Summary: When Benchmarks Lie: Evaluating Malicious Prompt Classifiers Under True Distribution Shift

Problem Statement

The deployment of LLM-based agents in autonomous environments introduces critical security vulnerabilities, specifically prompt injection, jailbreak attacks, and harmful requests. While recent literature has developed classifiers to detect these threats, current evaluation practices significantly overestimate generalization capabilities. Standard protocols typically employ train-test splits where test samples originate from the same dataset sources as training data. This methodology allows classifiers to exploit "dataset-identity signals"—features indicating the provenance of the data (e.g., formatting specific to WildJailbreak or Enron) rather than learning generalizable attack semantics. Consequently, models report near-perfect performance (AUC > 0.99) that fails to translate to true out-of-distribution (OOD) scenarios.

Methodology

Benchmark Construction

The authors compiled a benchmark of 18 datasets totaling 105,000 samples, covering four families of malicious inputs (harmful requests, jailbreaks, prompt injections including indirect/agent-based, and extraction attacks) alongside benign sources. The datasets vary in composition, with some being 100% malicious, 100% benign, or mixed-class.

Evaluation Protocol: Leave-One-Dataset-Out (LODO)

To measure true generalization, the paper proposes Leave-One-Dataset-Out (LODO) evaluation. Unlike standard cross-validation (CV) or held-out test sets that mix data from the same sources, LODO trains a classifier on $N-1$ datasets and evaluates it strictly on the held-out $N$ -th dataset. This ensures the classifier has never seen any examples from the test distribution's source.

Model and Feature Analysis

Models: Experiments span four LLMs from three families: Llama-3.1-8B, Gemma-3-27B, Qwen-3.5-2B/4B, and partial results for Llama-3.3-70B.
Classifiers: The study utilizes activation-based classifiers (linear probes and 2-layer MLPs) trained on:
- Raw Activations: Extracted from the residual stream (primarily layer 31 for Llama).
- Sparse Autoencoder (SAE) Features: Decomposed activations from layer 27 (Llama) and layer 42 (Gemma) to enable interpretability.
Shortcut Analysis: The authors introduce a LODO coefficient retention metric ( $r_j$ ) to analyze SAE features. This metric measures the stability of a feature's classifier coefficient when a specific dataset is held out. Features with low retention ( $r_j \ll 1$ ) or sign flips are identified as "dataset shortcuts."
Baseline Comparisons: The study compares activation probes against production guardrails (PromptGuard 2, LlamaGuard), LLM-as-Judge (LJ) variants, and dedicated text-based prompt-injection detectors (ProtectAI, Deepset).

Key Results

1. Standard Evaluation Severely Overestimates Generalization

Standard 5-fold CV and held-out test set evaluations report AUC scores between 0.996 and 0.999. In contrast, LODO evaluation reveals a significant performance drop:

Aggregate Gap: The CV-LODO AUC gap ranges from 8.0 to 16.5 percentage points across the four models.
Per-Dataset Variance: The accuracy gap for specific held-out datasets ranges from 1.2% to 25.4%. For example, on the Jayavibhav dataset, accuracy drops from 94.5% (held-out test) to 69.1% (LODO).
Universality: This gap is consistent across all model families and sizes tested (2B to 70B parameters).

2. Dataset Identity and Safety Signals Are Entangled

Distinguishability: A logistic regression classifier can predict the dataset identity of a sample with 96.6% accuracy using SAE features and 89.1% using raw activations, confirming that datasets are trivially distinguishable in feature space.
Subspace Overlap: Projecting activations onto the orthogonal complement of the top dataset-predictive directions degrades LODO performance by 1.75 percentage points. This suggests that the subspace identifying the dataset partially overlaps with the subspace relevant to safety classification.
Shortcuts: Analysis of the top 50 SAE features reveals that 28% (Llama) to 44% (Gemma) are "shortcuts." These include "pure dataset shortcuts" (directly predicting dataset identity) and "context-dependent shortcuts" (features that fire across datasets but derive class signal only from specific dataset compositions).

3. Ineffectiveness of Standard Mitigations

The authors tested seven standard domain-generalization (DG) techniques, including adversarial training (DANN), subspace projection, sample reweighting, and class balancing. None of these methods successfully closed the CV-LODO gap; in some cases (e.g., subspace projection), performance degraded further. This implies that standard DG assumptions—that domain-identifying and task-relevant signals can be cleanly separated—do not hold for prompt-attack detection.

4. Comparison with Production Baselines

Activation Probes vs. Guardrails: At matched False Positive Rates (FPR), activation-based probes outperform PromptGuard 2 and LlamaGuard on indirect and agentic attacks, which text-only guardrails often fail to detect due to a lack of tool-schema support.
LLM-as-Judge Sensitivity: LLM-as-Judge performance is highly sensitive to prompt templates, with True Positive Rates (TPR) varying from 11% to 93% depending on the prompt framing, often at the cost of unacceptably high FPRs.
Text Detectors: Dedicated text classifiers (e.g., ProtectAI, Deepset) exhibit extreme trade-offs under LODO-style evaluation, either missing most attacks or generating unusable false positive rates.

5. LODO-Weighted Explanations

By weighting SAE attributions with the retention metric ( $z_j \cdot w_j \cdot r_j$ ), the authors filter out dataset-artifact features. This re-ranks the top-20 features for approximately 98% of samples, providing more reliable per-prompt explanations that focus on genuine attack patterns rather than dataset-specific formatting.

Significance and Claims

The paper claims that the current evaluation paradigm for prompt-attack classifiers is fundamentally flawed, leading to a false sense of security. The primary contributions are:

LODO Necessity: Establishing that Leave-One-Dataset-Out evaluation is required to surface the true generalization gap (8–16pp) that standard CV masks.
Shortcut Taxonomy: Defining a retention metric and taxonomy that identifies dataset-dependent features which firing-ratio heuristics miss.
Entanglement Evidence: Demonstrating that dataset-identity and safety-relevant subspaces partially overlap, explaining why standard domain-generalization fixes fail.
Diagnostic Utility: Providing a framework to filter dataset artifacts from explanations and benchmark classifiers under conditions that better reflect real-world deployment shifts.

The authors explicitly state that they do not provide a definitive solution to close the gap, noting that disentangling dataset identity from safety signals remains an open challenge. Instead, they offer a diagnostic tool to prevent overpromising on production performance based on inflated CV metrics.