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

Shravan Doda

REVIEW 4 major objections 3 minor 1 cited by

Final-token safety probes miss jailbreaks because unsafe evidence often appears earlier in the prompt and is invisible at the last-token readout.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · grok-4.5

2026-07-12 16:55 UTC pith:KIOBQZTC

load-bearing objection Abstract-only diagnosis of final-token probe misses on jailbreaks; coherent and field-relevant, but unverifiable without methods or numbers. the 4 major comments →

arxiv 2605.12726 v2 pith:KIOBQZTC submitted 2026-05-12 cs.LG

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

Shravan Doda This is my paper

classification cs.LG

keywords final-token probesjailbreak detectionprefill trajectoriessafety probeshidden-state monitoringPCA-HMMLLM safety

verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

This paper argues that final-token safety probes—monitors that inspect only the last hidden state after a prompt is filled—systematically miss many jailbreaks. The reason is that probe-visible signals of unsafety often sit in earlier user-token representations during prefill and never surface at the final readout. Using SafeSwitch-style linear probes trained solely on clean harmful versus benign prompts across three instruction-tuned language models, the authors show high recall on clean harmful prompts yet frequent misses on jailbreaks, plus false positives on safety-adjacent benign text. Subspace analysis indicates the missed cases lie outside the directions the probe has learned; widening the probe bottleneck does not fix the mismatch. Token-level prefill scans confirm the early evidence, while naive max-pooling over tokens overfires on safe prompts. A lightweight PCA-HMM trajectory model trained on the same clean data recovers many of those misses without the catastrophic false-positive rate of pooling, so the paper positions trajectory-aware analysis as a practical diagnostic complement to final-token probes.

Core claim

Final-token safety probes trained only on clean harmful and benign prompts achieve high recall on clean harmful inputs yet miss many jailbreaks, because probe-visible unsafe evidence is frequently present earlier in the user-token prefill trajectory and is not exposed at the final-token state; a PCA-HMM model of those trajectories, trained on the same clean split, recovers many of the misses without the false-positive collapse of naive token pooling.

What carries the argument

The PCA-HMM trajectory model: principal components of hidden states along the user-content prefill sequence are modeled as a simple hidden Markov chain trained only on the clean harmful/benign split; it diagnoses final-token probe misses by scoring entire prefill trajectories rather than a single last-token readout.

Load-bearing premise

That SafeSwitch-style linear probes trained only on clean harmful-versus-benign prompts, tested on three instruction-tuned models, fairly represent the final-token probe paradigm so the observed early-evidence miss pattern generalizes.

What would settle it

Train the same final-token probes and the PCA-HMM on the clean split of additional instruction-tuned models or a broader jailbreak suite; if final-token probes already recover most jailbreaks at the last token, or if the trajectory model fails to recover the residual misses without raising false positives, the claimed prefill-time failure mode is refuted.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

Final-token probes alone are insufficient for reliable jailbreak detection even when they look strong on clean harmful prompts.
Trajectory diagnostics can surface unsafe evidence that never reaches the last token without the over-firing of simple max-pooling.
Widening a linear probe bottleneck does not automatically capture the directions that separate jailbreaks from clean benign prompts.
Safety monitors that inspect only the final prefill state systematically under-estimate risk on distributed-evidence jailbreaks.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

Safety-critical deployments may need lightweight trajectory monitors as a standard second-stage check rather than relying solely on last-token classifiers.
Jailbreak designers could deliberately bury unsafe cues early and restore a benign final token to evade final-token probes.
The same early-evidence pattern may appear in other single-state monitors (toxicity, factuality) that read only the last hidden state.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Desk Editor's Note

Abstract-only diagnosis of final-token probe misses on jailbreaks; coherent and field-relevant, but unverifiable without methods or numbers.

read the letter

The one thing to know is that this is an abstract-only empirical diagnosis: final-token safety probes (SafeSwitch-style, clean-trained) hit high recall on clean harmful prompts but miss many jailbreaks because probe-visible unsafe signal often sits earlier in user-token prefill and never shows up at the final readout. A clean-trained PCA-HMM on prefill trajectories recovers many of those misses without the FP blowup of naive max-pooling. That is the whole load-bearing claim.

What is new, if the numbers hold, is the clean split of the failure mode: not just “probes can fail,” but that the miss is often prefill-distributed evidence outside the probe subspace, that wider bottlenecks do not fix it, and that a simple trajectory model trained only on the same clean data recovers signal without collapsing specificity. The experimental structure described is the right non-circular one—train on clean harmful/benign, evaluate on jailbreaks and safety-adjacent benign—and the abstract is careful not to overclaim a new framework. Credit for naming a concrete, monitor-design-relevant failure and for proposing a lightweight diagnostic complement rather than another end-to-end detector.

Soft spots are almost entirely availability. We have no models, no jailbreak suite, no miss rates, no subspace plots, no HMM details, no thresholds, no error bars. The representativeness of SafeSwitch-style linear/bottleneck probes on three instruction-tuned LLMs is a real but secondary concern; the primary one is that the causal story and the recovery claim cannot be checked. Free parameters (bottleneck width, PCA dim, HMM states, thresholds) are free until the paper shows they are not driving the result. Nothing in the abstract looks circular or invented, but nothing is demonstrated either.

This is for people who build or evaluate runtime safety monitors and representation probes. If the full paper ships the numbers and ablations cleanly, it is worth a serious referee and a reading-group slot. Right now I would not cite it and would not bring the abstract alone. Send it to peer review if the full text appears with reproducible methods; desk-reject only if the full text never materializes or the numbers evaporate. My own posture: interesting diagnosis, currently unverified.

Referee Report

4 major / 3 minor

Summary. The manuscript diagnoses a prefill-time failure mode of final-token safety probes (SafeSwitch-style linear/bottleneck probes trained only on clean harmful vs. benign prompts) across three instruction-tuned LLMs. It claims these probes achieve high recall on clean harmful prompts but miss many jailbreaks because probe-visible unsafe evidence often appears earlier in user-token prefill representations and is not exposed at the final-token readout; subspace analyses are said to show missed jailbreaks lie outside the probe subspace, and bottleneck-width increases do not reliably fix the mismatch. Naive token max-pooling is reported to overfire on safe prompts, while a PCA-HMM trajectory model trained on the same clean split is claimed to recover many final-token misses from user-content prefill trajectories without catastrophic false positives, motivating trajectory-aware diagnostics as complements to final-token probes.

Significance. If the empirical diagnosis holds under full methods and results, the work would be a useful contribution to LLM safety monitoring: it would identify a concrete, non-circular limitation of a widely used final-token probe paradigm, provide token-level timing evidence for when unsafe signal is probe-visible, and offer a simple clean-trained trajectory baseline (PCA-HMM) that recovers misses without the false-positive pathology of naive pooling. That combination of failure diagnosis plus a constructive, same-split diagnostic complement would matter for both probe design and evaluation practice. With only the abstract available, however, none of these load-bearing claims can be confirmed, so significance remains conditional on unreported evidence.

major comments (4)

Abstract-only review: the central empirical claim (clean-trained final-token probes miss many jailbreaks with earlier prefill evidence; clean-trained PCA-HMM recovers them without pooling-level FPs) cannot be verified. No methods, model identities, jailbreak suite, miss/FP definitions, metrics, error bars, baselines, or artifacts are available. Until the full experimental section is inspectable, the diagnosis is an unverified assertion rather than a demonstrated result.
Abstract claim that 'missed jailbreaks differ from clean benign prompts along directions that are poorly captured by the probe's representational subspace': without the subspace analysis method (e.g., how the probe subspace is defined, projection residual statistics, controls for prompt length/style), it is impossible to judge whether this is a genuine geometric mismatch or a confound. This is load-bearing for the causal story that final-token readout fails because of subspace mismatch rather than label noise or distribution shift alone.
Abstract claim that 'increasing probe bottleneck width does not reliably resolve this mismatch': bottleneck width is a free parameter of the probe family. Without the ablation design (widths tried, layers, training protocol, and quantitative recall/FP tables), the claim that capacity is not the fix cannot be assessed and undercuts the argument that the failure is structural rather than underparameterization.
Abstract claim that a 'simple PCA-HMM trajectory model, trained only on the same clean split, recovers many final-token misses ... without the catastrophic false-positive behavior of naive token pooling': PCA dimensionality, HMM state count/transition structure, and decision thresholds are free parameters. Without reported recovery rates, FP rates on safety-adjacent benign prompts, and a clear comparison protocol to max-pooling, the constructive half of the paper cannot be evaluated for soundness or overfit to the clean split.

minor comments (3)

Abstract: 'SafeSwitch-style probes' and 'three instruction-tuned LLMs' should be named explicitly so readers can judge representativeness of the probe family and model set without the full text.
Abstract: 'high recall', 'miss many jailbreaks', and 'recovers many' are qualitative; even a one-line quantitative teaser (e.g., clean recall, jailbreak miss rate, PCA-HMM recovery and FP rates) would make the abstract falsifiable.
Abstract: clarify whether 'user-content prefill trajectories' exclude system/template tokens and how multi-turn or tool-use prompts are handled, if at all, in the claimed setting.

Circularity Check

0 steps flagged

No circularity detectable from abstract; clean-train / jailbreak-eval structure is non-circular by design.

full rationale

Only the abstract is available. It states that SafeSwitch-style final-token probes and a PCA-HMM trajectory model are trained solely on clean harmful/benign prompts, then evaluated on jailbreaks and safety-adjacent benign prompts. That train/eval split is the standard non-circular empirical structure: the models are not fitted to the test labels, and no equation or definition forces the reported miss/recovery pattern by construction. There are no self-citations, uniqueness theorems, fitted parameters renamed as predictions, or ansatzes smuggled via prior author work visible in the abstract. Residual risks (template leakage between clean and jailbreak sets, representativeness of the probe family) are ordinary empirical validity concerns, not circularity. With no full text, methods, or equations to inspect, no circular step can be quoted or exhibited; the honest finding is score 0 and empty steps.

Axiom & Free-Parameter Ledger

3 free parameters · 3 axioms · 0 invented entities

Abstract-only; free parameters and modeling choices for probes and PCA-HMM are implied but not quantified. No new physical entities. Domain assumptions are standard for LLM probing work. Invented entities are none beyond the diagnostic use of a PCA-HMM trajectory model as a complement.

free parameters (3)

probe bottleneck width
Abstract states increasing bottleneck width does not reliably fix the mismatch; width is a design/hyperparameter of the probe that affects capacity and reported failure.
PCA dimensionality / HMM state count and transition structure
A PCA-HMM trajectory model necessarily chooses latent dimension and number of hidden states; abstract does not report values but the recovery claim depends on those choices.
decision thresholds for probe and trajectory scores
Recall, miss, and false-positive claims require operating thresholds on probe logits or HMM likelihoods; not specified in the abstract.

axioms (3)

domain assumption Final-token hidden states after prompt prefill are a valid and standard readout for safety probes (SafeSwitch-style).
The entire failure-mode diagnosis is relative to this common design choice stated in the opening sentence.
domain assumption Clean harmful vs. benign labeled prompts are a sufficient training distribution for learning a safety-relevant linear/bottleneck probe.
Probes and the PCA-HMM are trained only on the clean split; generalization claims rest on this.
ad hoc to paper User-content prefill token hidden states form trajectories that can be summarized by PCA then modeled as an HMM for safety diagnostics.
The positive recovery result depends on this modeling choice introduced as the diagnostic complement.

pith-pipeline@v1.1.0-grok45 · 6101 in / 2714 out tokens · 29551 ms · 2026-07-12T16:55:49.494080+00:00 · methodology

0 comments

read the original abstract

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

Figures

Figures reproduced from arXiv: 2605.12726 by Shravan Doda.

**Figure 1.** Figure 1: Jailbreak–XSTest operating-point shift from the finaltoken probe to the user-window PCA-HMM trajectory diagnostic. Across all three models, the trajectory diagnostic recovers many final-token misses while reducing XSTest false positives at this operating point. is not whether any token score is high, but how the score evolves across the prompt: missed jailbreaks exhibit a high harmful-request score follow… view at source ↗

**Figure 2.** Figure 2: Complementarity on jailbreak prompts. Stacked bars partition 900 jailbreak prompts by whether they are caught by the final-token probe, the PCA-HMM trajectory diagnostic, both, or neither. Percentages inside bars are normalized by the 900-prompt jailbreak set. PCA-HMM catches many prompts missed by the final-token probe: 236 for Llama, 106 for Mistral, and 310 for OLMo3. D. PCA-HMM Length Correlations [PI… view at source ↗

Review history (2 revisions) →

discussion (0)

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

Breaking Refusal in the First Half: A Mechanistic Study of the Prefill Jailbreak
cs.CL 2026-07 conditional novelty 6.0

Refusal under a prefill jailbreak is a shallow response-site computation: the harm representation stays intact, the failure lives in an early response window, and the dominant mechanism is passive autoregressive conditioning.