Building a System One Model for Prompt Injection

Defender, our prompt injection guard for tool-calling agents, checks every tool result in tiers. Tier 1 is pattern matching and finishes in about a millisecond. Tier 2 is the 22MB classifier from the first post in this series, which scores each string on CPU in tens of milliseconds. Tier 3 is the careful reviewer for the cases Tier 2 isn’t sure about, and until recently it was a small generative LLM that read the tool result and wrote back a verdict.

Tier 3 caught attacks well enough, but it was expensive to run. A generative judge needs a GPU, a few hundred milliseconds and a parsed block of text for every call, so it only ever saw what Tier 2 escalated. If a judge were cheap enough, it could look at every tool result next to Tier 2 instead of after it.

This post covers how we replaced that judge with a decision model, the kind TypeSafe calls a System One model and that Jev made popular last month.

System One models

The term comes from Daniel Kahneman, who split thinking into a fast, intuitive System 1 and a slow, deliberate System 2. TypeSafe uses the label for a model that makes decisions and does nothing else. You send it a state (text or JSON) and a set of typed questions. Each question is a choice between named options, a score on a scale, or a “noul”, which is a yes/no probability. The model returns a probability for every option and never writes text.

The best-known System One model is Jev, which TypeSafe released on 15 September. Cloudflare followed on 1 October with Clef and Clef-flash, open weights under Apache 2.0, which add a scoring head to a Qwen backbone and answer every question in one forward pass. On the research side, Laya showed that a small encoder can answer the same kind of typed questions. None of these was built for prompt injection, though all of them fit the kind of decision a guard makes.

For Tier 3, two properties of this design stood out.

The output space is closed. A generative judge reads attacker-controlled text and then writes freely, so a good enough injection can push it into refusing, rambling, or producing something the parser chokes on. A decision model can only put probability on options we declared. Injected text can still push those probabilities toward the wrong answer, so the model can be fooled. What it can’t do is respond in any other way.

The output is already a probability. A generative judge that writes “block” or “allow” is making a decision nobody configured: whichever word is more likely wins, which works out to a fixed threshold of 0.5. Every retrain moves the scores underneath that cut. When we fine-tuned a generative LLM version of the judge, small changes to the training data moved its recall a long way while its ranking quality barely changed. We’d rather choose the threshold ourselves, per deployment, from a target false positive rate, and a decision model hands us the number to threshold.

Why not just call Jev

Jev was the obvious first thing to try, but it only runs as a hosted API. Tool results carry employee records, emails and CRM notes, and Defender is designed so that data stays inside the infrastructure that’s already handling it. Sending every tool result to a third party for classification is the same objection that ruled out calling a frontier LLM in the first post.

We still tried it through OpenRouter, to see whether the format fit our problem, and it did: all of our questions fit into one call of about 600 input tokens, which at Jev’s listed price is $0.000025. On three of our rows, written with placeholder criteria, it flagged benign tool output as injection and was most confident on exactly the rows it got wrong. Three rows can’t judge Jev, but they were enough to tell us not to count on zero-shot transfer to tool payloads.

That pushed us toward open weights we could fine-tune and host: a decision model on a ModernBERT backbone, which speaks the same choice, score and noul interface as Jev and fits on a single small GPU. Zero-shot on our held-out test split it was right 54.5% of the time, so it needed training before it could sit in a data path.

Choosing the questions

Asking a model “is this a prompt injection?” gets a single, fuzzy answer. We split it into the things an injection has to do to be dangerous inside a tool result, and asked each one as a three-way choice:

• Does the text address the agent with instructions, implicitly or explicitly?
• Does it ask for data to be sent somewhere, and is that an actionable instruction or only a mention?
• Does it claim identity or authority, and is it an outright impersonation?

To turn three answers into one decision, we take the highest probability any question puts on its most dangerous option. An email asking a colleague to send the Q3 numbers to finance mentions moving data but addresses a person. An email body telling the assistant to forward the same numbers to an outside address scores high on two questions at once.

Splitting the verdict also makes blocks easier to triage. In our tests we can see which question drove each block and how confident the model was, which says more than a single score.

We started with a fourth question: does this payload look structurally unusual for its tool? The three frontier models we used as labelling teachers disagreed on it more than on any other question. It also transferred worst zero-shot, and after fine-tuning it was the worst calibrated of the four (expected calibration error 0.153, against 0.04 to 0.05 for the rest). We dropped it from the gate, because structure is cheaper to measure from the payload directly.

Fine-tuning

We fine-tuned in three rounds. The first used a few thousand tool results, each labelled per question by three frontier teacher models. It moved accuracy from 54.5% to 88.1%, with per-question temperature scaling bringing calibration error down to 0.04 to 0.05.

The second round added rows shaped like BIPIA, half of them benign. Those benign rows made the model less trigger-happy everywhere, including on payloads that look nothing like BIPIA.

In the third round we added our own data, including connector hard negatives of the kind described in Detection Isn’t the Hard Part, which grew the corpus 2.5x. We retrained from scratch, and BIPIA AUC fell from 0.86 to 0.65. Recalibrating didn’t help, because the model had stopped ranking attacks above benign text, so there was no threshold left to fix. The new data was real, clean and deduplicated. It was also shaped differently, and it cut the share of injection signal the model saw per epoch to about 40% of what it had been.

What fixed it was continuing training instead of starting over. We continued from the previous checkpoint and replayed its training split alongside the new rows, so the old signal couldn’t be outnumbered. BIPIA AUC came back to 0.86, and the model flagged none of the connector payloads in our false positive benchmark.

The continued model also needed its calibration revisited. The temperatures fitted to the continued model came out at 2.2 to 4.8, roughly double the earlier ones, and applying them in full cut recall on our Internal Benchmark from 72% to 33%. Applying a quarter of the correction, T' = 1 + 0.25 × (T − 1), kept most of the false positive gains without giving the recall away. If you fine-tune continually, refit the temperatures after every round and check them against recall as well as calibration error.

Serving

Cheaper answers only help if the model can serve them quickly, and our first deployment managed 8 requests per second on an L4.

The first doubling came from fixing a bug: our batch path skipped autocast, so the model had been running in fp32. Fixing that took us to 18.5. Moving to vllm-jev, a vLLM engine for Jev-style decision models, took us to about 33.

Next we tried TensorRT, expecting it to beat vLLM, but it reached only 28.5 on the same GPU. The profile explained why: we were losing to padding, not to vLLM’s scheduler. Batching tool results of mixed length means padding every row to the longest one, and that costs more per request than running them one at a time. vLLM packs tokens and never pays for padding.

Profiling the requests themselves turned up a bigger waste. Each question was a separate row, so the tool result got encoded once per question, and 70% of every request’s 1,004 tokens was the same state text repeated three times. We packed all three question heads and one copy of the state into a single sequence, which cut a request to 536 tokens.

Changing the input format meant retraining the model. Retraining on the original teacher labels kept AUC within a point but lost 11 to 14 points of recall at 1 to 2% false positive rate. Distilling the packed model from the existing model’s own output distributions worked, and it matched or beat that model on every AUC we track. At 1% false positive rate, its recall held on BIPIA and went up on AgentDojo.

With torch.compile, bf16 weights and 64-token length buckets to keep padding down to about 6%, the packed model serves 55.9 requests per second on one L4, at 16ms of GPU time per request. That’s 1.7x vllm-jev on the same card, and an older T4 still manages 36.

Distillation and replay worked for the same reason: both preserve the function the previous model had already learned. Fresh labels only describe the target, and leave the model free to reach it a different way.

Current state

We’re running our initial tests of the decision model behind Defender’s Tier 3 interface. Before it replaces anything, we want more training runs to show the results hold, and a block threshold chosen on data we don’t report on.

The main problem we haven’t solved yet is long tool results, which need either a longer context window or a way for Tier 3 to read in chunks.

In Defender, Tier 3 is available as Deep Scan. If you’re interested in the model, request access through the form on defender.stackone.com and tell us where you’d want it to run.

---

Defender is open source: github.com/StackOneHQ/defender. Tier 3 plugs in through the Tier3Provider interface, so you can bring your own judge, decision model or otherwise.