PaperlayerView exact revision on arXiv

Source-linked AI summary

Bounded Reachability & Jailbreak Detection via Contraction-Constrained State Space Models

Omanshu Thapliyal

arXiv:2610.02853v1cs.AI

TL;DR

Formal robustness guarantees for prompt-level safety classifiers on sequence models remain comparatively unexplored. The paper proves and uses a contraction condition for exact certification in linear SSMs, then evaluates a contraction-regularized S4 head for jailbreak detection. Jailbreak detection results show strong benchmark performance, while a linear probe matches or exceeds the head’s discrimination metrics; the S4 head’s distinct contribution is certification.

  • Problem

    The paper asks when an SSM safety classifier can be certified to keep the same prediction for all embedding sequences within a bounded perturbation, addressing a gap in formal robustness analysis for prompt-level sequence classifiers.

    Source: §1 · p. 1, §2 · p. 2

  • Method

    The paper derives exact interval propagation for a linear SSM classifier and uses a contraction penalty to encourage ||A||∞ < 1, then applies a contraction-regularized S4 head to jailbreak detection.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    JDR=1.0 across six JailbreakBench attack families, while logistic regression on mean-pooled Mamba-130M embeddings achieves AUROC=1.0 on all three benchmarks and exceeds the SSM head on every discrimination metric.

    Source: §5 · p. 8

  • Takeaways & Limitations

    The S4 head’s distinctive contribution is formal certification of a classification decision at a specified perturbation radius, rather than superior discrimination.

    Source: §5 · pp. 8–9

    The JailbreakBench results use a static threat model and do not guarantee robustness against adaptive adversaries.

    Source: §4 · p. 6

Abstract

from arXiv · show

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_\infty$ norm of the state transition matrix must satisfy $\norm{A}_\infty<1$ (the \emph{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 $\norm{A}_\infty=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.

1 Introduction

The paper asks how prompt-level safety classifiers for sequence models can receive formal robustness guarantees. It analyzes exact certification for linear SSMs and applies a contraction-regularized S4 head to jailbreak detection.

Motivation

  • Prompt-level classifiers operate on inputs alone but have rarely received formal robustness analysis for sequence models, while response-level judges can fail on base models without structured refusal phrases.The paper frames formal certification as a safety mechanism beyond empirical evaluation.Source: §1 · p. 1

Contributions

  • For a toy LTI SSM classifier, ||A||∞ < 1 is sufficient for bounded reachability and necessary for a horizon-independent bound; a hinge penalty can enforce contraction during training.The bound is δ ||B||∞/(1 − ||A||∞).Source: §1 · p. 2
  • The S4 head achieves JDR=1.0 across six JailbreakBench attack families, while logistic regression on the same mean-pooled embeddings achieves AUROC=1.0 on all three benchmarks and equals or exceeds the head on every detection metric.The paper identifies formal certifiability, rather than superior discrimination, as the S4 head’s distinctive contribution.Source: §1 · p. 2

Scope

  • The formal results apply only to the toy LTI setting, whereas the JailbreakBench results are empirical claims tied to specific benchmarks and threat models.The paper points to appendices for detailed methodology and results.Source: §1 · p. 2

2 Problem Setup

This section defines certification as invariance of an LTI SSM classifier’s prediction over a bounded embedding perturbation set. It sets out the classifier, adversary, and interval-based criterion used to pose the problem.

2 Problem Setup

  • The central question is whether a trained LTI SSM classifier’s predicted label stays identical for every embedding sequence within an l∞ ball of radius δ around the input.This is a worst-case guarantee over the threat set, rather than an evaluation against a particular attack.Source: §2 · p. 2

2.1 Linear SSM Classifier

  • The classifier processes token embeddings sequentially, updates a hidden state, and maps the state and current input to a scalar logit and binary prediction.Its state is a weighted sum of past inputs, with weights determined by A.Source: §2.1 · p. 3

2.2 Perturbation Model

  • The adversary can independently shift each coordinate of every token embedding by at most δ in l∞ distance, with perturbations applied directly in embedding space.Relating δ to particular text edits requires a calibration procedure.Source: §2.2 · p. 3

2.3 Certification Criterion

  • Because the classifier is linear in each embedding, its reachable logits form a closed bounded interval; certification requires that the entire interval lie on the correct side of zero.An interval straddling zero means a perturbation in the set could flip the label, though uncertified examples may still be correct under most perturbations.Source: §2.3 · pp. 3–4, §2.2 · p. 3

3 Methodology

The methodology propagates perturbation intervals exactly through linear SSM dynamics, derives the contraction condition for bounded reachability, and uses a penalty to encourage it. It also describes the S4 safety head used for jailbreak detection.

3.1 Interval Bound Propagation for LTI Systems

  • Exact interval bound propagation splits A into positive and negative parts and routes interval endpoints according to sign, avoiding relaxation or enclosure error.The resulting output reach tube tracks how possible outputs evolve as tokens are processed.Source: §3.1 · p. 4

3.1.1 Contraction and Bounded Reachability

  • When ||A||∞ < 1, interval width converges to a finite steady state; when ||A||∞ ≥ 1, it grows unbounded as sequence length increases.The state half-width is bounded by δ ||B||∞/(1 − ||A||∞) in the contracting case.Source: §3.1.1 · pp. 4–5

3.1.2 Contraction Regularization

  • A hinge penalty pushes ||A||∞ below 1, with gradient on the maximum-row-sum row when the norm exceeds 1 and zero gradient when it is below 1.A sweep over 12 penalty values shows monotone improvement in certified fraction.Source: §3.1.2 · p. 5

3.2 S4 Safety Head for Jailbreak Detection

  • The S4 head processes token embeddings sequentially, averages its per-token hidden states over real-token positions, then maps the average through a linear layer to a scalar score.Its two-layer diagonal S4 uses zero-order hold discretization; the recurrent representations are order-dependent, while identical pooled states receive identical scores.Source: §3.2 · pp. 5–6

4 Certifiable Jailbreak Detection

The section evaluates a prompt-only SSM safety head on jailbreak benchmarks, finding strong transfer alongside a linear-probe advantage and a distinct certification contribution.

4 Certifiable Jailbreak Detection

  • Evaluation uses a prompt-only head trained on 770 examples, with a static primary threat model in which attack templates do not adapt to detector scores.The response-based reference judges return JDR=0.0 on the base Mamba model because it produces no structured refusals, a protocol mismatch rather than a meaningful comparison.Source: §4 · p. 6

4.1.1 Detection Walkthrough

  • Across attack formulations, scores are similar, while the benign distribution has mean µ ≈−1.5 and is separated from harmful families.The paper interprets this pattern as consistent with semantic rather than syntactic operation.Source: §4.1.1 · p. 6
  • LR-mean-pool reaches AUROC=1.0 on all three benchmarks, while SSM-head records 0.993, 0.981, and 0.947 on JBB, AdvBench, and HarmBench.The HarmBench AUROC gap is concentrated in copyright and harassment/bullying categories; adding MLP capacity does not close it.Source: §4.1.1 · pp. 6–7
  • Figure 1 walkthroughs show direct attacks rising within 10–20 tokens, wrapper attacks rising later, and all four harmful and both benign examples classified correctly.Benign prompts stay near or below τ=0 as tokens are processed.Source: §4.1.1 · p. 6, Figure 1 · p. 7
  • Contraction regularization raises certified fraction from 41% to 68% while accuracy remains essentially flat; neither the probe nor MLP provides this guarantee.The paper identifies certification, rather than superior discrimination, as the SSM-head's contribution.Source: §4.1.1 · p. 7

4.2 Adaptive Threat Model

  • A continuous proxy-gradient attack left JDR=1.0 and lowered mean score only from 3.681 to 3.675; separate black- and white-box attacks operate on discrete tokens.The proxy attack used an l2-ball ε =0.5; no example crossed the τ =0 boundary.Source: §4.2 · pp. 7–8

4.3 Cross-Benchmark Transfer

  • On AdvBench, SSM-head reaches DR=0.994 versus the probe's DR=1.000; on HarmBench, it reaches DR=0.988 versus 0.995.SSM-head misses 3/520 AdvBench instructions and 5/400 HarmBench behaviors; HarmBench misses are concentrated in copyright and harassment/bullying.Source: §4.3 · p. 8

5 Conclusion

The paper establishes contraction as the condition for bounded reachability in a toy LTI SSM classifier and applies a contraction-regularized S4 head to jailbreak detection. Its contribution is formal certification, while a linear probe performs as well or better at discrimination.

5 Conclusion

  • Contraction of the state transition matrix is sufficient and necessary for bounded reachability in the toy LTI SSM classifier.The conclusion also reports applying a contraction-regularized S4 safety head to jailbreak detection on JailbreakBench.Source: §5 · p. 8
  • Under a static threat model, SSM-head reaches JDR=1.0 across six JBB attack families and AUROC=0.994, but adaptive attacks achieve ASR=0.933 and ASR=1.0.The conclusion also reports zero-shot transfer to AdvBench and HarmBench in the preceding evaluation.Source: §5 · p. 8
  • Logistic regression on mean-pooled Mamba-130M embeddings achieves AUROC=1.0 on all three benchmarks and exceeds SSM-head on every discrimination metric.The paper attributes transfer to the embedding space and identifies formal certifiability, rather than superior discrimination, as the SSM-head’s contribution.Source: §5 · pp. 8–9

A.1 Related Work

  • The paper situates its analysis in LTI stability theory and contrasts classifier certification with representation-based safety work that lacks formal robustness analysis.It frames its contribution as certification of the classifier, not a claim that simpler baselines lack safety signals.Source: §A.1 · p. 10

A.2 Contributions

  • The box-propagation framework computes the tightest axis-aligned reachable interval for LTI systems, while the hinge penalty enforces contraction without constrained optimization or projection steps.The framework supports both final-step certification and trajectory visualization.Source: §A.2 · p. 10
  • The prompt-only SSM-head uses token embeddings without model responses, and reach tubes visualize the temporal geometry behind scalar certification results.These tubes are presented as diagnostics for LTI SSM safety heads.Source: §A.2 · p. 11

A.3 Limitations

The stated limitations concern the model class, certification margins, embedding-space threat model, and conservatism of the certified set.

A.3 Limitations

  • Exact box propagation applies to the toy LTI model, not SSMs with nonlinear gating, selection mechanisms, or attention.Extending it to nonlinear settings requires bounding activations and introduces relaxation error.Source: §A.3 · p. 11
  • Certification rates depend on logit margins: at δ = 10^-3, constrained-model half-widths of 0.11-0.18 are comparable to the mean clean logit of 0.14.Larger margins from more data, a larger model, or a margin-maximizing loss could improve rates at the same δ.Source: §A.3 · p. 11
  • The δ-ball is defined in token embedding space, and bridging it to character or token space is outside the work’s scope.That bridge would require composing interval propagation with the tokenizer’s discrete-to-continuous map.Source: §A.3 · p. 11
  • The certified set is a conservative subset of the true robust set, so non-certification does not establish that an example is non-robust.Tighter LP relaxations could certify more examples at the same δ but would cost more to compute.Source: §A.3 · p. 11

A.4 Sketch of proof for Proposition 1

The proof tracks exact interval widths through the LTI recurrence and uses geometric-series convergence to establish bounded reachability under contraction. Certification then follows when the output interval lies entirely on the correct side of the decision boundary.

A.4 Sketch of proof for Proposition 1

  • The state half-width follows w_t = |A|w_(t-1) + δ|B|1, yielding a geometric-series bound that converges when ∥A∥∞ < 1.The proof passage states that the series diverges when ∥A∥∞ ≥ 1 and r_0 > 0.Source: §A.4 · p. 11
  • A label is certified exactly when the reachable output interval stays on the correct side of zero: lower bound positive for label 1 or upper bound negative for label 0.The interval is exact for the LTI system, so this criterion introduces no relaxation.Source: §A.4 · p. 11
  • The output interval half-width adds direct input uncertainty δ∥D∥∞ to the propagated state uncertainty scaled by ∥C∥∞.The output equation is s_T = Cx_T + Du_T.Source: §A.4 · pp. 11–12
  • Figure 3 depicts shaded output reach tubes and solid nominal trajectories across 64 steps for constrained and unconstrained models.It shows three representative examples for each model condition.Source: Figure 3 · p. 12

B A Toy Example: Empirical Results

The toy LTI classifier illustrates how contraction bounds reachable intervals and supports certification. Certification is stronger for the constrained model, especially as the perturbation radius grows.

B A Toy Example: Empirical Results

  • The ToySSMClassifier uses a word-level tokenizer, 32-dimensional embeddings, an 8-state LTI SSM, and sequences up to T = 64, with 10,000 training and 600 held-out examples.The constrained model uses λ = 0.10 and has ∥A∥∞ below 1; the unconstrained model uses λ = 0 and has ∥A∥∞ = 1.091.Source: B A Toy Example: Empirical Results · p. 12

B.1 Geometric Structure of the Reach Tube

  • With ∥A∥∞ = 0.998, the constrained model’s reach tube stabilizes before T = 64, while the unconstrained model’s tube widens and can cross the decision boundary.The comparison illustrates the bounded-versus-growing reach-tube behavior for toxic and non-toxic examples.Source: §B.1 · pp. 12–13

B.2 Certification vs. Perturbation Radius

  • At δ = 10−4, the constrained model certifies 59% of examples versus 41% for the unconstrained model; at δ = 2 × 10−3, the constrained model retains 12% certification.The unconstrained model certifies zero examples beyond δ = 5 × 10−4.Source: §B.2 · p. 13, Table 1 · p. 13, Figure 4 · p. 13

B.2.1 Interpreting the Perturbation Radius

The experiments relate embedding-space perturbation radii to text-edit distances, then examine how contraction and sequence length affect certification. The radius-to-edit mapping is a heuristic, not a formal guarantee against token-level attacks.

B.2.1 Interpreting the Perturbation Radius

  • For the toy classifier, δ = 10−4 is below the mean typo-pair distance, δ = 10−3 reaches the lower end of the synonym-pair distribution, and δ = 5 × 10−3 overlaps it.The calibration measures l∞ distances for typo, synonym, and antonym word pairs; the figure shows pair-type distributions and reference δ values.Source: §B.2.1 · pp. 13–14, Figure 5 · p. 14
  • The δ calibration is a heuristic proxy: isolated-word embedding distances and token-substitution difficulty are correlated but not equivalent.The δ-ball perturbs all positions simultaneously, and the calibration uses a toy word-level vocabulary rather than the Mamba-130M BPE tokenizer.Source: §B.2.1 · p. 14
  • Across 10 trained models, interval width rises sharply near ∥A∥∞ = 1 and certification collapses above 1 at δ = 10−3; the measured width matches the analytic bound near the threshold.The hinge penalty produces a sharp certification step at ∥A∥∞ = 1, while accuracy remains flat or slightly increases.Source: §B.2.1 · p. 14, §B.3 · p. 14, Figure 7 · p. 15
  • At δ = 10−4 across T ∈ {8, 16, 24, 32, 40, 48, 56, 64}, constrained-model certification remains nearly flat, while unconstrained-model certified fraction falls to zero by T = 32 as its tube widens beyond the logit margin.Source: §B.3 · p. 14

B.4 Certification Boundary Analysis

Certification depends on both the reachable interval and the example’s clean-logit margin. In the horizon sweep, constrained certification plateaus after the reachable set reaches steady state, while unconstrained certification declines more steeply.

B.4 Certification Boundary Analysis

  • Examples with large positive or negative clean logits are certified, while examples near the decision boundary remain uncertain regardless of contraction status.The mean clean logit for toxic examples is approximately 0.14, leaving many examples in the margin-sensitive zone.Source: §B.4 · p. 16

B.5 Horizon Sweep T ∈[4, 128]

  • The constrained model’s certified fraction falls from 0.996 at T = 4 to 0.629 at T = 64, then plateaus; the unconstrained model falls from 0.992 to 0.421 by T = 64.The analytic bound is constant across T for the constrained model and grows exponentially for the unconstrained model.Source: §B.5 · pp. 16–17, Figure 11 · p. 17

C Jailbreaking Example: Supplementary Results

Across six jailbreak attack families, score distributions cluster at similar positive logits, while benign prompts cluster near −1.5.

C Jailbreaking Example: Supplementary Results

  • All six attack-family score distributions concentrate around logits +2 to +5 with nearly identical shapes and means, while the benign distribution concentrates near −1.5.The families include direct, prefix injection, DAN, developer mode, roleplay, and completion attacks.Source: C Jailbreaking Example: Supplementary… · p. 17

C.1 Static Threat Model

The static-threat analysis examines how the SSM-head separates harmful from benign prompts and how harmful intent is represented across token sequences. Under adaptive attacks, direct-attack prompts can be evaded, while template-wrapped prompts remain unaffected in the reported attacks.

C.1 Static Threat Model

  • The token heatmap attributes higher positive contributions to harmful request tokens than to DAN and roleplay wrapper text, indicating classification driven by harmful content rather than surface syntax.The figure describes token-level |∆score|, with positive contributions raising the harm score and negative contributions lowering it.Source: §C.1 · p. 17, Figure 13 · p. 19
  • The formal reachability analysis applies to a toy LTI classifier, not the nonlinear Mamba backbone; the plotted high-dimensional reachable regions are illustrative projections, not certified bounds.The analysis describes mean-pooled Mamba-130M embeddings as 768-dimensional and visualizes projections into two dimensions.Source: §C.1 · p. 17
  • Across attack families, SSM-head logits concentrate above zero while benign-prompt scores cluster below zero, showing separation regardless of attack template.The figure uses τ=0 as its dashed reference line.Source: Figure 12 · p. 18
  • The harmful/benign signal is distributed across the token sequence: mean pooling aggregates it effectively, whereas single-token representations lose almost all of it.The reported comparison uses logistic regression probes on mean, final-token, and first-token summaries, evaluated with 5-fold cross-validation on 100 harmful JBB and 100 benign prompts.Source: §C.1 · pp. 17–18, Table 2 · p. 18

C.2 Adaptive Threat Model

  • ASR=0.933 for black-box and ASR=1.0 for white-box attacks on direct prompts, while both attacks have ASR=0.000 on template-wrapped prompts; τ=2.46 yields hard-negative FPR=0.010 at JDR=0.990.The black-box attack reduces direct-attack JDR to 0.067 after 10 steps, while template attacks maintain JDR=1.0.Source: §C.2 · p. 18, Figure 17 · p. 20, Figure 18 · p. 21
  • At τ=0, hard-negative prompts straddle the probe boundary with FPR=0.557; projected perturbation clouds show high-confidence and easy-benign prompts staying on one side while near-boundary clouds cross it.The clouds are l∞(0.8) projections; Figure 14 reports AUROC=1.0 for a one-dimensional projection onto the probe normal.Source: Figure 14 · p. 19, Figure 15 · p. 20
  • Harmful prompts remain above τ=0 across the tested 5–120-token lengths, while score instability is highest at 5–15 tokens.The length test covers 11 lengths and compares direct attacks, DAN attacks, and benign prompts.Source: §C.1 · p. 18, Figure 16 · p. 20

Related papers

Loading 2610.02853v1…