ML-Based Adversarial Prompt Detection, in depth: data splits by attack family, baselines, fine-tuned encoders, perplexity features and adversarial training

Define the task before the model

Decide what the classifier should flag before collecting data, because the label definition is the model. Three families are usually in scope. Direct injection is a user message that tries to override the system prompt. Indirect injection is the same intent hidden in a document, web page or tool result the model will read. Jailbreaks try to get policy-violating output, often through role-play, hypotheticals or encoding tricks. Optimised adversarial suffixes, such as those produced by the GCG attack described in universal adversarial suffixes, form a fourth family with a very different surface form.

Public models made different choices here, which is why scores do not transfer between them. ProtectAI's deberta-v3-base-prompt-injection-v2 is a binary classifier (benign or injection) fine-tuned from DeBERTa-v3-base. Meta's Llama Prompt Guard 2 comes in 86M and 22M parameter versions and labels text benign or malicious; the first Prompt Guard had separate injection and jailbreak classes. Read the model card of anything you adopt and test it on your own traffic.

Data: hard negatives and family splits

Data quality decides the outcome more than architecture. Build the training set from three sources: public injection and jailbreak collections, the output of your own red team and automated attack generators, and a large sample of real benign traffic from your product. The benign sample matters most, because production false positives come from legitimate text the model has never seen.

Add hard negatives on purpose: security articles that quote attacks, users asking how prompt injection works, support tickets that say "ignore my previous message", code that contains the word system. A detector trained without them learns that the phrase ignore previous instructions is the attack, and blocks every developer who pastes a red-team report.

Tag each attack example with its family, meaning the template, generator or campaign it came from. Then split by family, not by row. A random split puts near-identical variants of the same template in train and test, and you will report 99 percent accuracy on a model that fails on the first new template it meets. Deduplicate with normalised text and MinHash before splitting.

A ladder of models

Treat the table as relative, not as measurements; latency depends on hardware and text length. Most production systems combine the first three and reserve an LLM judge for escalation.

A baseline in twenty lines

Always start with a baseline. It takes minutes, sets the bar the expensive model must beat, and often stays in production as a cheap first stage. Character n-grams survive tricks such as inserted spaces better than word tokens do.

import unicodedata, re
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.metrics import average_precision_score

ZW = re.compile("[\u200b-\u200f\u2060\ufeff]")

def normalise(t):
    t = unicodedata.normalize("NFKC", t)
    return ZW.sub("", t).lower()

def train_baseline(train, test):
    # train/test: lists of (text, label, family); split by family upstream
    clf = make_pipeline(
        TfidfVectorizer(analyzer="char_wb", ngram_range=(3, 5),
                        min_df=2, sublinear_tf=True, preprocessor=normalise),
        LogisticRegression(C=4.0, class_weight="balanced", max_iter=2000))
    clf.fit([t for t, _, _ in train], [y for _, y, _ in train])
    scores = clf.predict_proba([t for t, _, _ in test])[:, 1]
    print("PR-AUC", average_precision_score([y for _, y, _ in test], scores))
    return clf

Report PR-AUC rather than accuracy, because attacks are rare in real traffic and accuracy rewards a model that says benign to everything. NFKC normalisation and zero-width removal belong in the serving path too, applied identically, or the model sees different text in production than in training.

Perplexity features for optimised suffixes

Optimised suffixes from attacks like GCG are usually high-perplexity token soup. Perplexity filtering was proposed as a baseline defence against them, and it works on unmodified suffixes. Compute it over sliding windows, because a short suffix appended to a long fluent prompt barely moves the whole-prompt average.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tok = AutoTokenizer.from_pretrained("gpt2")
lm = AutoModelForCausalLM.from_pretrained("gpt2").eval()

@torch.no_grad()
def window_perplexity(text, window=32, stride=16):
    ids = tok(text, return_tensors="pt").input_ids[0][:1024]
    logits = lm(ids.unsqueeze(0)).logits[0, :-1]
    nll = torch.nn.functional.cross_entropy(logits, ids[1:], reduction="none")
    worst = 0.0
    for start in range(0, max(1, len(nll) - window + 1), stride):
        worst = max(worst, nll[start:start + window].mean().exp().item())
    return worst   # use as a feature, not a verdict

Use the maximum window perplexity as one feature for the final classifier rather than a hard rule. On its own it flags base64, minified code, URLs and non-English text, and published work shows attackers can optimise suffixes for low perplexity once they know the filter exists.

Fine-tuning an encoder

The main model is usually a small encoder fine-tuned for binary classification. The loop is standard; what matters is the data, class weighting and long-input handling.

from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          TrainingArguments, Trainer)
from datasets import Dataset

name = "microsoft/deberta-v3-base"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name, num_labels=2)

def encode(batch):
    return tok([normalise(t) for t in batch["text"]], truncation=True, max_length=512)

train_ds = Dataset.from_list(train_rows).map(encode, batched=True)
eval_ds = Dataset.from_list(heldout_family_rows).map(encode, batched=True)

args = TrainingArguments(output_dir="detector", learning_rate=2e-5,
                         per_device_train_batch_size=16, num_train_epochs=3,
                         weight_decay=0.01, warmup_ratio=0.1, bf16=True)
Trainer(model=model, args=args, train_dataset=train_ds,
        eval_dataset=eval_ds, processing_class=tok).train()

At serving time an encoder sees at most 512 tokens, while a retrieved document can be 20,000. Chunk the input into overlapping windows, score each, and take the maximum; an attack hidden on page nine must not be truncated away. Truncating to the first 512 tokens is the most common way a well-trained detector is defeated in production.

Adversarial training

A detector that has only seen published attacks has learned their surface. Make it robust by attacking it yourself and training on what gets through:

for round in range(R):
    model = train(train_set)
    misses = []
    for x in attack_seed_set:                 # known malicious prompts
        for x_adv in mutate(x):               # synonyms, paraphrase by an LLM,
                                              # homoglyphs, spacing, encodings,
                                              # benign padding, language switch
            if model.score(x_adv) < threshold and still_malicious(x_adv):
                misses.append((x_adv, 1, family_of(x)))
    train_set += dedupe(misses)
    report(recall_on(heldout_families), fpr_on(benign_holdout))

The check still_malicious is essential: a mutation that destroys the attack is a benign example, and labelling it malicious teaches the model to block noise. Run the target LLM on the mutated prompt or have a reviewer confirm. Watch the false positive rate on the benign holdout every round, because adversarial training tends to trade it away. Tools for generating the attack side are covered in automated red teaming.

Worked example and thresholds

A hypothetical walk-through, not a benchmark: suppose you fine-tune an encoder on a few tens of thousands of labelled prompts. A random row split will typically report recall close to 1.0. Re-split by family, holding out several attack templates entirely, and recall on the held-out families is usually much lower, with misses clustered in templates that disguise the attack, for example as a translation task. Adversarial rounds with translation and role-play mutations close part of that gap, and benign false positives on developer traffic may rise until you add hard negatives such as code review comments. Measure your own numbers with the method in prompt injection evaluation.

Thresholds deserve one rule here: choose them on the benign rate of real traffic, not on a balanced test set, because at a 0.1 percent attack base rate even a 1 percent false positive rate means most alerts are false. The scanners article works through that arithmetic.

Serving the detector

A detector that adds 300 ms to every request will be switched off by the first team that owns a latency target, so design the serving path as carefully as the model. Run the char n-gram stage inline on every message; it costs almost nothing and can block the crudest attacks outright. Send the encoder stage to a dedicated classification service that batches requests, and export the model to ONNX or a similar runtime so it can run on CPU for modest traffic. A small encoder on a shared GPU handles far more volume, but then the detector competes with generation for capacity.

Score retrieved content when it is ingested, not on every request that reads it. A document in a knowledge base is read thousands of times; scoring it once at indexing time, storing the result with the chunk, and rescoring only when the model version changes removes most of the cost of indirect injection detection. Cache scores for repeated user text by a hash of the normalised string.

Decide in advance what happens when the detector is down or times out. Fail open keeps the product working and accepts more risk; fail closed blocks traffic during an outage. Agents with powerful tools usually justify failing closed for tool-calling paths and open for plain chat. Finally, keep the score internal: return a generic refusal, log the score and model version, and never echo which phrase triggered the block, or you hand attackers a gradient to follow.

Drift and retraining

Attack styles move monthly and product traffic moves with every feature launch. Log scores for all traffic, sample flagged and unflagged items for human review each week, and track three numbers: the score distribution on benign traffic, recall on a frozen regression set of past attacks, and recall on the newest red-team batch. Retrain when the newest batch drops below target, and never ship a retrained model that regresses on the frozen set. Version the model, the normaliser and the threshold together; changing one without the others silently moves your operating point.

Failure modes

• Leaky split. Template variants in both train and test, giving inflated metrics.
• Keyword model. No hard negatives, so it blocks people discussing security.
• Truncation. Only the first 512 tokens scored; the attack sits later.
• Normalisation mismatch. Training normalised text, serving raw text, or the reverse.
• Oracle exposure. Returning the score or a detailed block reason lets attackers optimise against it.
• Detector as the only control. A miss then means full compromise; pair it with layered jailbreak defences.

What to do next

1. Write down the label definition and the attack families in scope.
2. Collect benign production samples and hard negatives before collecting more attacks.
3. Tag every attack with a family and split by family; deduplicate first.
4. Train the char n-gram baseline and record PR-AUC on held-out families.
5. Fine-tune an encoder, score long inputs in overlapping windows, and add window perplexity as a feature.
6. Run two rounds of adversarial training with a still-malicious check, watching benign false positives.
7. Set the threshold on real traffic base rates and schedule weekly review and drift tracking.