Safety heads are lightweight classifiers attached to pretrained language models for flagging harmful inputs before generation. Their empirical detection performance has been studied, but their formal robustness properties remain largely unexplored. We ask when a State Space Model (SSM)-based safety head can be certified to produce the same prediction for all inputs within a bounded embedding-space perturbation. We prove that the answer turns on a single condition: the l∞ norm of the state transition matrix must satisfy \normA∞<1 (the contraction condition), which enables exact interval bound propagation (IBP) certification for linear time-invariant classifiers. When the contraction condition holds, the reachable output interval has bounded steady-state width and examples can be certified as robustly classified. When it fails, the interval grows exponentially with sequence length and certification is impossible at any practical perturbation radius. We enforce contraction with a hinge penalty and show on toxic comment data that certified fraction improves from 41% to 59%, with a sharp empirical phase transition at \normA∞=1 matching the theory. Applying a contraction-regularized S4 head to jailbreak detection on JailbreakBench, we achieve a zero-shot transfer to AdvBench (DR=0.994) and HarmBench (DR=0.988). A logistic regression on mean-pooled Mamba-130M embeddings matches or exceeds the S4 head on every detection metric, confirming that harmful intent is already linearly separable in the embedding space. The S4 safety head's contribution is not superior discrimination but the formal certification that no probe-based approach provides.
Figures & tables
Figure 1: SSM-head detection walkthroughs for six representative prompts. Each panel shows the running logit score as tokens are added (top) and the full prompt text (bottom, colored border). Red border = harmful ground truth; green border = benign.
Figure 2: Proxy-gradient attack: mean score trajectory (left, 3.681 → 3.675) and score distribution before vs. after (right). JDR=1.0 before and after; this is a continuous embedding-space relaxation, not a discrete-token attack.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Output reach tube [stl,sth] (shaded) and nominal trajectory stnom (solid) over time steps t=1,⋯,64 for three representative examples under the constrained ( λ=0.10 , left) and unconstrained ( λ=0 , right) models.
Figure 4: Certified fraction vs. perturbation radius δ (log scale, 50 log -spaced points in [10−5,3×10−2] ) for constrained (blue) and unconstrained (red) models. Left: toxic examples only. Center: non-toxic examples only. Right: overall (all correctly classified).
δ
Con. Toxic
Con. Non-tox
Con. Overall
Unc. Toxic
Unc. Non-tox
Unc. Overall
1e−4
58.8%
59.3%
59.2% ( ± 0.24)
24.2%
43.1%
41.0%
2e−4
44.1%
49.2%
48.6% ( ± 0.26)
9.1%
33.7%
31.0%
5e−4
23.5%
38.4%
36.9% ( ± 0.24)
6.1%
10.1%
9.7%
1e−3
11.8%
27.3%
25.7% ( ± 0.16)
0%
0.4%
0.3%
2e−3
5.9%
13.5%
12.7% ( ± 0.07)
0%
0%
0%
Appendix
Table 1: Certified fractions at representative δ values. Constrained ( ∥A∥∞=0.999 ) vs. unconstrained ( ∥A∥∞=1.091 ). Overall column: mean ± std across 3 seeds.
Figure 5: δ calibration to text-edit distances. Left: Violin plot of l∞ embedding distances by pair type. Gray dotted lines mark the reference δ values used in the experiments. Right: Distribution histograms of pair distances with the same δ reference lines.
Figure 6: Results Gallery for ToySSMClassifier : each cell shows the text, clean logit, and output interval [sl,sh] . Border color encodes outcome: red = certified toxic, blue = certified non-toxic, orange = uncertain correct, gray = misclassified.
Figure 7: Growth-rate sweep across 10 λ values: ∥A∥∞ vs. λ (left), certified fraction at two δ values (center), and mean output interval width vs. ∥A∥∞ with the analytic steady-state bound overlaid (right).
Figure 8: Effect of contraction penalty λ on ∥A∥∞ (left), accuracy and recall (center), and certified fraction at δ=10−4 and δ=10−3 (right) across 12 λ values.
Figure 9: Certified fraction at δ=10−4 for various T values (primary y-axis) and analytic bound ( log scale, right axis). The constrained model’s cert rate is nearly flat; the unconstrained model reaches zero by T=32 .
Figure 10: Clean logit sTnom vs. output interval half-width scatter at delta=5×10−5 (left) and δ=2×10−3 (right); colors encode certification status (blue = certified toxic, light blue = certified non-toxic, orange = uncertain, gray = misclassified).
Figure 11: Horizon sensitivity sweep (16 points, T∈[4,128] ): empirical certified fraction at δ=10−4 for constrained vs. unconstrained models (left) and analytic bound overlay (right).
Figure 12: SSM-head logit score KDEs by attack family and benign prompts; dashed line = τ=0 . All attack families concentrate above zero; benign prompts cluster below zero, showing consistent score separation regardless of attack template.
Figure 13: Marginal risk contribution per token: each cell is one sub-word token; color intensity = ∣Δscore∣ when that token is added (red = positive contribution = raises harm score; blue = negative contribution = lowers harm score). Rows: three direct attack variants, two DAN variants, and four benign prompts. All cells are normalized to the same colormap range per figure. Harmful tokens (verbs like ”malware”, ”weapon”) consistently receive high positive attribution; DAN preamble tokens (”from now on”, ”act as”) receive lower or negative attribution
Figure 14: 2D PCA projection of mean-pooled Mamba-130M embeddings (768 → 2 dims; PC1=9.3%, PC2=5.0%) with harmful JBB (red, n=100), easy-benign (green, n=30), and hard-negative (orange, n=70) prompts; two starred points show l∞ (0.8) perturbation halos (right: 1D projection onto probe normal, AUROC=1.0). Hard-negative prompts straddle the dashed probe boundary, consistent with FPR=0.557 at τ=0 .
Figure 15: Projected l∞ (0.8) perturbation clouds (200 samples each) for four prompts in the same 2D PCA subspace as Fig. 14 : harmful high-confidence (cloud fully in detected region), easy-benign (cloud fully in safe region), hard-negative near-boundary (cloud straddles the probe boundary), and harmful near-boundary (cloud partially crosses safe region).
Pooling strategy
AUROC (5-fold CV)
σ
Remark
Mean pool (all tokens)
0.998
0.002
Aggregates full prompt context
Final non-pad token
0.718
0.029
Much weaker: last token carries little signal alone
First token (position 0)
0.515
0.025
Near chance: first token is almost always punctuation/BOS
Appendix
Table 2: Representation ablation: LR probe AUROC (5-fold CV ±σ ) across three token-embedding pooling strategies, on 100 harmful JBB + 100 benign prompts.
Figure 16: Mean ± 10–90th percentile (left) and within-prompt score std across lengths (right). Harmful prompts remain above τ=0 at all lengths tested; score instability is highest at 5-15 tokens.
Figure 17: Black-box greedy token substitution attack (10 steps, 30 direct-attack prompts): score trajectories (blue = not evaded, orange = evaded, ASR=0.933) at left, score before vs. after at center, and per-family JDR before vs. after at right. Template attacks maintain JDR=1.0; direct attacks drop to JDR=0.067 after 10 steps.
Figure 18: White-box HotFlip-style attack (10 steps, top-5 candidate positions): ASR comparison between black-box (ASR=0.933) and white-box (ASR=1.0) on direct attacks, both ASR=0.0 on DAN/roleplay (left); score trajectories for 30 direct-attack prompts (center); semantic similarity vs. token change rate showing evasion does not require large surface changes (right).
Guardrail Classifiers defend production language models against harmful behavior, but although results seem promising in testing, they provide no formal guarantees. Providing formal guarantees for such models is hard because "harmful behavior" has no natural specification in a discrete input space: and the standard epsilon-ball properties used in other domains do not carry semantic meaning. We close this gap by shifting verification from the discrete input space to the classifier's pre-activation space, where we define a harmful region as a convex shape enclosing the representations of known harmful prompts. Because the sigmoid classification head is monotonic, certifying the worst-case point is sufficient to certify the entire region, yielding a closed-form soundness proof without approximation in O(d) time. To formally evaluate these classifiers, we propose two constructions of such regions: SVD-aligned hyper-rectangles, which yield exact SAT/UNSAT certificates, and Gaussian Mixture Models, which yield probabilistic certificates over semantically coherent clusters. Applying this framework to three author-trained Guardrail Classifiers on the toxicity domain, every hyper-rectangle configuration returns SAT, exposing verifiable safety holes across all classifiers, despite seemingly high empirical metrics. Probabilistic GMM certificates also expose a divergent structural stability in how these models represent harm. While GPT-2 and Llama-3.1-8B maintain robust coverage of 90% and 80% across varying boundaries, BERT's safety guarantees prove uniquely volatile. This 'coverage collapse' to 55% at the optimal threshold reveals a sparsely populated safety margin in BERT, which only achieves full coverage by adopting an extremely conservative pessimistic threshold. These approaches combined, provide new insights on how effective Guardrail Classifiers really are, beyond traditional red-teaming.
Nikita Kezins, Urbas Ekka, Pascal Berrang +1
Delft University of Technology · University of Birmingham & Zeroth Research
Multi-turn jailbreak attacks progressively erode LLM safety alignment across seemingly innocuous conversation turns, achieving success rates exceeding 90% against state-of-the-art models. Existing alignment-based and guardrail methods suffer from three key limitations: they require costly weight modification, evaluate each turn independently without modeling cumulative safety erosion, and detect attacks only after harmful content has been generated. To address these limitations, we first formulate the proactive early jailbreak detection problem with a new metric, detection lead, that measures how early an attack can be detected before the LLM complies. We then propose SAFEDREAM, a lightweight world-model-based framework that operates as an external module without modifying the LLM's weights. SAFEDREAM introduces three components: (1) a safety state world model that encodes LLM hidden states into a compact safety representation and predicts how it evolves across turns, (2) CUSUM detection that accumulates weak per-turn risk signals into reliable evidence, and (3) contrastive imagination that simultaneously rolls out attack and benign futures in latent space to issue early alarms before jailbreaks occur. On three multi-turn jailbreak benchmarks (XGuard-Train, SafeDialBench, SafeMTData) against 8 baselines, SAFEDREAM achieves the best detection timeliness across all benchmarks (1.06-1.20 turns before compliance) while maintaining competitive false positive rates and outperforming baselines in detection quality.
Bo Yan, Weikai Lin, Yada Zhu +1
University of Central Florida · University of Rochester · IBM Research
Autonomous systems increasingly rely on Large Language Models (LLMs) yet the safety infrastructure surrounding these models introduces latency and compute overhead. This limits utility in resource-constrained, time-critical deployments. Existing external guardrail models remain blind to the model's internal workings, creating a fundamental assurance gap. We ask: does the model already know when the content is harmful? We extract activations from LLaMA-3.1-8B and train lightweight MLP classifier probes (12.6M parameters) to detect harmful prompts. Evaluated on WildJailbreak, Beavertails, and AEGIS 2.0, our probes achieve F1 scores of 99%, 83%, and 84%, respectively competitive with 1000x larger guard models while cutting latency and compute costs.