PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shravan Doda

arXiv:2605.12726v2cs.LG

TL;DR

Final-token safety probes read one hidden state, leaving open whether they detect harmful requests after jailbreak wrapping. This paper evaluates clean-trained probes and analyzes their representational geometry and prefill trajectories. The probes miss many wrapped requests, while trajectory diagnostics recover many final-token misses.

  • Problem

    It is unclear whether a final-prompt-token probe trained on clean harmfulness still detects the same harmful request after jailbreak wrapping.

    Source: §1 · p. 1

  • Method

    The study evaluates clean-trained SafeSwitch-style probes across three models and examines final-token geometry and user-content prefill trajectories, including a PCA-HMM diagnostic.

    Source: §1 · p. 1

  • Results

    Probes miss many wrapped harmful requests despite high recall on clean harmful prompts, while PCA-HMM trajectory models recover most final-token-probe misses across models.

    Source: §1 · p. 1

  • Takeaways & Limitations

    Trajectory-aware hidden-state analyses can serve as diagnostic complements to final-token probes for studying where safety-relevant information appears in prefill.

    Source: §1 · pp. 1–2

    The study is diagnostic rather than a deployed defense; false positives remain nonzero, and HMM states should not be interpreted as mechanistic safety states.

    Source: §8 · p. 5

Abstract

from arXiv · show

Final-token safety probes monitor a single hidden state after prompt prefill, but jailbreak prompts can contain probe-visible unsafe evidence distributed across earlier user-token representations that is missed by this readout. We study this prefill-time failure mode using SafeSwitch-style probes trained only on clean harmful and benign prompts across three instruction-tuned LLMs. The probes achieve high recall on clean harmful prompts, but miss many jailbreaks and can produce false positives on safety-adjacent benign prompts. Subspace analyses suggest that missed jailbreaks differ from clean benign prompts along directions that are poorly captured by the probe's representational subspace, and increasing probe bottleneck width does not reliably resolve this mismatch. Token-level prefill analyses reveal that probe-visible unsafe evidence often appears earlier in the sequence but is not exposed at the final-token readout, while naive max-pooling over token positions overfires on safe prompts. A simple PCA-HMM trajectory model, trained only on the same clean split, recovers many final-token misses from user-content prefill trajectories without the catastrophic false-positive behavior of naive token pooling, motivating trajectory-aware hidden-state analyses as diagnostic complements to final-token probes

1. Introduction

The paper examines why final-token safety probes miss jailbreaks despite detecting clean harmful prompts. It connects those misses to probe-visible directions and prefill trajectories, motivating trajectory-aware analysis as a diagnostic complement.

1. Introduction

  • Across three instruction-tuned models, probes trained on clean harmful and benign prompts achieve high recall on clean harmful prompts but miss many jailbreaks.The jailbreaks wrap held-out harmful requests in adversarial framing, role play, obfuscation, or instruction-following context.Source: §1 · p. 1
  • Missed jailbreaks differ from clean benign prompts along directions poorly captured by the probe-visible subspace, and wider bottlenecks do not reliably improve detection.The caught-versus-missed jailbreak contrast aligns with the clean harmful-versus-benign training direction, consistent with a shortcut-like readout.Source: §1 · p. 1
  • Missed jailbreaks often have earlier prefill positions with high probe scores, but naive token pooling also produces high scores on safe prompts.PCA-HMM trajectory models trained on the same clean split recover most final-token-probe jailbreak misses across models, as a diagnostic rather than a deployed detector.Source: §1 · pp. 1–2

2. Related Work

Prior work studies safety-related directions and hidden-state dynamics; this paper uses a narrower probe-geometry analysis and a single-prefill trajectory diagnostic.

Representation geometry of refusal and harmfulness

  • The geometry analysis uses difference-of-means contrasts and one trained binary final-token probe’s row-space, without treating its directions as universal harmfulness features or mechanistic control variables.This scope distinguishes the analysis from work on refusal directions, multiple safety-related directions, and harmfulness subcategory probes.Source: §2 · p. 2

Hidden-state trajectories for safety diagnostics

  • The trajectory analysis asks whether earlier probe-visible evidence appears within one prefill pass and treats PCA-HMM as a diagnostic complement, not a standalone safety monitor.Related trajectory work also examines layer-wise evolution, decoding-time trajectories, and latent trajectory classifiers.Source: §2 · p. 2

3. Experimental Setup

The experiments train final-token probes on clean harmful and benign prompts, then evaluate harmful-request detection and benign-prompt false positives across several test sources.

Models and probes

  • Three instruction-tuned models use SafeSwitch-style probes trained on final prompt-token hidden states from the original SafeSwitch train split, with a default threshold of 0.5.The models are Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.1, and OLMo3-7B-Instruct.Source: §3 · p. 2

Evaluation sets

  • Evaluation uses held-out SorryBench harmful prompts, the same requests wrapped with HarmBench jailbreak templates, AdvBench harmful requests, and benign safety-adjacent XSTest prompts.The reported measures are harmful-source and jailbreak detection rates and false-positive rate on XSTest.Source: §3 · p. 2
  • Table 1 covers harmful-source detection and benign XSTest false positives, comparing jailbreak prompts made by wrapping the same held-out SorryBench requests used in the clean condition.Source: Table 1 · p. 2

4. Final-Token Probes

Final-token probes detect clean harmful prompts but miss many wrapped versions and also flag some benign safety-adjacent prompts. The section tests threshold trade-offs, downstream implications, and whether wider bottlenecks resolve the misses.

4.2. Probes miss wrapped requests

  • Probes retain high recall on clean SorryBench prompts but miss a substantial fraction when the same harmful requests are wrapped, while also producing non-trivial false positives on XSTest.Lowering the threshold to recover wrapped requests would further increase false positives on benign safety-adjacent instructions.Source: §4.2 · p. 2
  • A final-token miss prevents the stage-2 refusal head from being invoked, allowing the wrapped harmful request to flow through to generation.Source: §4.2 · p. 2
  • The same probes detect AdvBench at 99.2–99.8% and HarmBench at 96.2–100.0% across three models, contrasting with misses caused by wrapping the same SorryBench prompts.Source: §4.2 · pp. 2–3

4.3. Width is not the bottleneck

  • 72.2% to 78.0% for Llama, 86.9% to 87.6% for Mistral, and 63.6% to 55.7% for OLMo3 as width rises from 64 to 1024, so wider readouts do not reliably improve jailbreak detection.The sweep uses the same training procedure and split.Source: §4.3 · p. 3

5. Geometry of Missed Jailbreaks

Missed jailbreaks differ from clean benign prompts along a direction poorly captured by the final-token probe’s row-space. The geometry is consistent with a shortcut-like readout that relies chiefly on the clean harmfulness contrast.

5. Geometry of Missed Jailbreaks

  • The authors project contrast directions into the probe row-space and compare their probe-visible energy and projected cosine alignments.The final-token hidden state motivates the geometry analysis; the probe’s row-space is constructed from the first-layer weight matrix’s right singular vectors.Source: §5 · p. 3
  • The analysis contrasts clean training harm, safety-adjacent benign prompts, missed jailbreaks versus clean benign prompts, and caught versus missed jailbreaks.dharm is the training-set harmful-versus-benign contrast; dsafe contrasts harmful SorryBench prompts with safety-adjacent benign XSTest prompts.Source: §5 · p. 3
  • 3.5%, 2.4%, and 3.2% of the missed-jailbreak direction’s squared norm is captured by the probe row-space, the lowest energy among four directions across all three models.The remaining 96–98% lies outside the probe row-space; no single SVD basis vector aligns strongly with any tested direction.Source: §5 · p. 3, Table 2 · p. 3
  • Probe projection sharpens alignment with the harmful proxy from 0.32–0.39 to 0.68–0.78 and with the safety proxy from 0.12–0.19 to 0.41–0.51.After projection, the caught-versus-missed contrast has only weak alignment with dmiss (0.21, 0.13, 0.06).Source: §5 · p. 3
  • Caught jailbreaks retain stronger projection onto the clean harmful-versus-benign direction, while the wrapper appears to reduce exposure of the relevant feature at the final token.Safety-related cues contribute partially, but the strongest projected alignment is with dharm.Source: §5 · pp. 3–4

6. Unsafe Evidence Appears Before the Last Token

Final-token readouts can miss unsafe evidence visible earlier in jailbreak prompts. Token-position results point to a positional failure, while trajectory diagnostics recover many misses and reduce XSTest false positives at the reported operating point.

6. Unsafe Evidence Appears Before the Last Token

  • 87 Llama, 65 Mistral, and 127 OLMo3 final-token-missed jailbreaks had recoverable embedded harmful-request spans, showing probe-visible evidence can occur earlier in the prompt.The failure is positional: by the final token, the wrapped prompt no longer exposes the same probe-visible feature.Source: §6 · p. 4
  • SafeSwitch uses only the final-token Stage-1 decision to gate Stage 2, so earlier probe-visible evidence cannot be recovered by the downstream refusal head.The head can only refine prompts that pass the final-token gate.Source: §6 · p. 4
  • Across all three models, the user-window PCA-HMM trajectory diagnostic recovers many final-token misses while reducing XSTest false positives at this operating point.Missed jailbreaks can show a high harmful-request score followed by collapse at the final token, motivating trajectory-level diagnostics rather than max-pooling.Source: Figure 1 · p. 4, §6 · p. 4

7. Token Trajectories Recover Probe Misses

A clean-trained PCA-HMM trajectory diagnostic recovers many jailbreaks missed by final-token probes, without inheriting max-pooling’s XSTest false-positive rate. Its sequence scores compare user-content trajectories against clean harmful and benign trajectories.

7. Token Trajectories Recover Probe Misses

  • The PCA-HMM diagnostic tests whether a user-content trajectory is more likely under clean harmful trajectories or clean benign trajectories.Each model uses a 64-dimensional PCA and two 2-state diagonal-covariance Gaussian HMMs fit on balanced clean training prompts.Source: §7 · p. 4
  • Trajectory scores are length-normalized log-likelihood ratios comparing harmful and benign HMM likelihoods.The threshold is selected to maximize harmful-versus-benign accuracy on the clean training split, without tuning on jailbreaks or XSTest.Source: §7 · p. 4
  • 236 of 250 Llama misses, 106 of 118 Mistral misses, and 310 of 328 OLMo3 misses are recovered without max-pooling’s 100% XSTest false-positive rate.Table 3 reports recovery only among jailbreaks missed by the final-token probe.Source: §7 · pp. 4–5, Table 3 · p. 5
  • Missed-jailbreak direction energy is higher in the trajectory PCA subspace than in the final-token probe row-space: 3.5% → 11.5% for Llama, 2.4% → 29.2% for Mistral, and 3.2% → 33.7% for OLMo3.The authors interpret the PCA-HMM as a diagnostic bridge, not a deployed detector or a set of mechanistic features.Source: §7 · p. 5

8. Limitations and Future Work

The study is diagnostic rather than a deployed defense, and the authors identify false positives, interpretation boundaries, and future evaluation needs. Wider final-token probe bottlenecks do not reliably improve jailbreak detection.

8. Limitations and Future Work

  • The study is diagnostic rather than a deployed defense; false positives remain nonzero, and HMM states should not be interpreted as mechanistic safety states.Future work should examine calibrated fusion, adaptive and length-controlled jailbreaks, and causal interventions on trajectory features.Source: §8 · p. 5

A. Bottleneck Width Sweep

  • Wider final-token probe bottlenecks do not reliably improve jailbreak detection: Llama gains modestly, Mistral is roughly flat, and OLMo3 degrades.Table 4 reports detection rate as a function of bottleneck width.Source: A. Bottleneck Width Sweep · p. 7, Table 4 · p. 7

B. Geometry Alignment Results

The geometry analysis compares cosine alignments in full and probe-projected hidden-state spaces; trajectory complementarity and within-source length correlations provide additional views of the diagnostic.

B. Geometry Alignment Results

  • Table 5 reports cosine alignment of ∆cm with the harmful, safety-proxy, and missed-jailbreak directions in the full hidden-state space; the note also identifies projection onto the probe row-space.Source: Table 5 · p. 7, B. Geometry Alignment Results · p. 7

C. Trajectory Complementarity

  • PCA-HMM catches 236 Llama, 106 Mistral, and 310 OLMo3 jailbreak prompts missed by the final-token probe, among a set of 900 jailbreak prompts.Figure 2 partitions prompts into those caught by the final-token probe, the PCA-HMM, both, or neither.Source: Figure 2 · p. 7

D. PCA-HMM Length Correlations

  • Within each evaluation source, window-length and PCA-HMM score correlations are weak to moderate, so the score is not purely a length proxy within a source.Table 6 reports within-source Spearman rank correlations.Source: D. PCA-HMM Length Correlations · p. 7, Table 6 · p. 8

E. Token-Position Results

Token-position comparisons show that located harmful-request spans in jailbreak prompts are detected even when the final-token readout misses them. Clean harmful prompts remain high at the final token.

E. Token-Position Results

Related papers

This paper cites

Loading 2605.12726v2…