🤖 machine learning

Before the Last Token: Diagnosing Final-Token Safety Probe Failures

This paper identifies that final-token safety probes fail to detect jailbreaks because unsafe evidence is often distributed across earlier prefill tokens rather than concentrated at the final state, and proposes that trajectory-aware analyses using PCA-HMM models can effectively recover these missed detections without the high false-positive rates of naive token pooling.

Original authors: Shravan Doda

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shravan Doda

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 have a very smart security guard (the AI model) who is supposed to stop dangerous requests before they happen. To help this guard, you hired a specialized inspector (the "probe") whose only job is to look at the very last word of a user's message and decide: "Is this dangerous?"

This paper investigates why this inspector sometimes fails to catch bad guys who are trying to sneak past the system using "jailbreaks" (tricks to bypass safety rules).

Here is the story of what the researchers found, explained simply:

1. The "Last Word" Problem

The inspector is trained to look at the final token (the last word) of a sentence to make a safety decision.

How it works: If the last word looks suspicious, the guard stops the conversation.
The flaw: Bad actors (jailbreakers) wrap their dangerous requests in fancy costumes, role-plays, or confusing instructions. By the time the sentence ends, the "last word" looks innocent, even though the middle of the sentence contained a clear danger signal.
The Analogy: Imagine a thief hiding a stolen diamond inside a box of teddy bears. The inspector only looks at the very top of the box (the last word). If the top looks like a cute bear, the inspector says, "All clear!" and lets the box pass, missing the diamond hidden deep inside.

2. Why the Inspector Misses the Clues

The researchers tested this on three different AI models and found that the inspector is very good at spotting obvious, plain-language bad requests. But when the bad request is wrapped in a "jailbreak" costume, the inspector gets confused.

It's not about "muscle": The researchers wondered if the inspector just needed to be "smarter" or have a bigger brain (more capacity). They made the inspector's brain much bigger, but it didn't help much.
The Real Issue: The problem is that the "danger signal" gets lost in the shuffle. When a bad request is wrapped in a trick, the final word shifts away from the "danger zone" in the AI's mind. The inspector is looking for a specific type of danger, but the trick moves the signal to a place the inspector can't see.

3. The "Hidden Clues" Earlier in the Sentence

The researchers then looked at the entire sentence, word by word, not just the last one.

The Discovery: In the cases where the inspector failed at the end, the "danger signal" was actually very loud and clear in the middle of the sentence (right where the bad request was hidden).
The Trap: However, simply looking at the loudest word anywhere in the sentence doesn't work either. Why? Because innocent, safe sentences can also have loud, "scary-sounding" words in the middle (like discussing a movie plot about a crime). If you flag every sentence with a loud word, you accidentally stop all the innocent people too.

4. The Solution: Watching the "Movie," Not Just the "Freeze Frame"

Instead of looking at a single freeze-frame (the last word) or just picking the loudest frame (max-pooling), the researchers tried watching the whole movie (the trajectory of the sentence).

The New Approach: They built a simple model that watches how the "danger score" changes from the first word to the last.
The Pattern: They found that jailbreaks have a specific pattern: The score goes high (danger!) when the bad request appears, and then drops (safe!) when the trick finishes. Innocent sentences don't follow this specific "spike and drop" pattern.
The Result: By watching this pattern, they could catch almost all the jailbreaks that the "last word" inspector missed, without accidentally stopping innocent conversations.

Summary

The paper concludes that relying on the last word to judge safety is like judging a book by its cover. Bad actors can make the cover look clean while hiding a dangerous story inside.

The researchers suggest that safety systems need to look at the whole story (the trajectory of the hidden states) to catch these tricks. They aren't saying this new method is a perfect, ready-to-use product yet, but it proves that the "danger clues" are there; we just need to look at the right place and in the right way to find them.

Technical Summary: Before the Last Token: Diagnosing Final-Token Safety Probe Failures

Problem Statement
Current safety monitoring systems for Large Language Models (LLMs), such as SafeSwitch, often rely on "final-token probes." These systems train classifiers on the hidden state of the last token in a prompt's prefill phase to detect unsafe content before generation begins. The paper identifies a critical failure mode in this architecture: while these probes achieve high recall on "clean" harmful prompts, they frequently fail to detect "jailbreak" prompts. Jailbreaks embed harmful requests within adversarial framing, role-play, or obfuscation. The authors argue that these wrappers shift the final-token representation away from the specific "harmful-vs-benign" contrast direction learned during training, causing the probe to miss the unsafe request even though the harmful content is present in the prompt. This failure is not merely a capacity issue, as increasing the probe's bottleneck width does not reliably resolve the detection gap.

Methodology
The study evaluates three instruction-tuned models: Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.1, and OLMo3-7B-Instruct.

Probe Training: SafeSwitch-style probes were trained exclusively on "clean" harmful and benign prompts from the original SafeSwitch training split. These probes read the final prompt-token hidden state, project it through a bottleneck (initially width $w=64$ ), and predict safety.
Evaluation: The probes were tested on:
- Clean held-out harmful prompts (SorryBench).
- The same harmful prompts wrapped with jailbreak templates from HarmBench.
- Benign, safety-adjacent prompts (XSTest) to measure false positives.
- Direct harmful benchmarks (AdvBench) to rule out generic distribution shifts.
Geometric Analysis: The authors analyzed the representational subspace of the probes. They computed difference-of-means vectors for various sets (e.g., missed jailbreaks vs. benign prompts) and measured their "probe-visible energy"—the fraction of a direction's norm captured by the probe's row-space (via SVD).
Trajectory Analysis: Instead of relying solely on the final token, the authors analyzed token-by-token hidden-state trajectories during the prefill phase. They investigated whether unsafe evidence appeared earlier in the sequence but collapsed by the final token.
Diagnostic Modeling: To test if trajectory information could recover missed detections, the authors fitted a simple PCA-HMM (Principal Component Analysis - Hidden Markov Model) on the clean training data. This model learned the trajectory dynamics of clean harmful vs. benign prompts and was applied to jailbreak trajectories without re-tuning on jailbreak data.

Key Results

Probe Failure on Jailbreaks: The final-token probes maintained high recall (~95%) on clean harmful prompts but dropped significantly on jailbreaks (ranging from ~63% to ~87% depending on the model). Conversely, they produced non-trivial false positives (23%–36%) on benign safety-adjacent prompts.
Capacity is Not the Bottleneck: Increasing the probe bottleneck width from 64 to 1024 did not reliably improve jailbreak detection. For OLMo3, performance actually degraded. This suggests the failure is geometric (where the signal lies) rather than a lack of representational capacity.
Geometric Misalignment: The direction separating missed jailbreaks from benign prompts ( $d_{miss}$ ) had the lowest probe-visible energy (approx. 2.4%–3.5%) compared to the clean training contrast ( $d_{harm}$ ). The probe row-space is effectively orthogonal to the specific features that distinguish missed jailbreaks from benign inputs.
Positional Failure: Unsafe evidence in missed jailbreaks is often visible at intermediate token positions (specifically near the embedded harmful request) but disappears or "collapses" by the final token.
Naive Pooling Failure: Simply taking the maximum score across all token positions is not a viable solution, as it triggers on 100% of benign XSTest prompts (false positives), indicating that high intermediate scores are not semantically reliable on their own.
Trajectory Recovery: The PCA-HMM diagnostic, trained only on clean data, successfully recovered a large majority of the final-token misses (97%–98% recovery rate) while maintaining a significantly lower false-positive rate on XSTest (6%–15%) compared to naive pooling. The missed jailbreak direction showed significantly higher energy in the trajectory PCA subspace than in the final-token probe row-space.

Significance and Claims
The paper positions itself as a diagnostic study rather than a deployed defense mechanism. Its primary contribution is identifying why final-token probes fail: they rely on a "shortcut-like" readout that is sensitive to the clean training distribution but blind to the specific geometric shifts induced by jailbreak wrappers.

The authors claim that:

Final-token probes are insufficient for detecting wrapped harmful requests because the wrapper-induced shift moves the representation out of the probe's sensitive subspace.
Safety-relevant information often exists in the hidden-state trajectory before the final token but is lost by the time the final readout occurs.
Simple trajectory-aware models (like the PCA-HMM), trained only on clean data, can serve as effective diagnostic complements to final-token probes. They can recover missed detections by analyzing the evolution of safety signals across the prompt, offering a more robust view of the prompt's internal state without the catastrophic false-positive rates of naive token pooling.

The work motivates a shift in safety diagnostics from single-point readouts to trajectory-aware analyses to better understand and mitigate the limitations of current probe-based safety systems.