State space models improve safe input detection with formal guarantees
Bounded Reachability & Jailbreak Detection via Contraction-Constrained State Space Models
Artificial Intelligence
Summary
Detecting harmful inputs before a language model responds is important for safety. The paper studies a type of model called State Space Models (SSMs) used for this detection and finds a key mathematical condition that ensures the detection is reliable even with small changes in input. The authors show that when this condition holds, the model’s safety predictions can be proven robust. They test this on toxic comments and jailbreak detection data, improving the fraction of reliably detected harmful inputs and confirming the model’s predictions can be formally certified.
What this means in practice
- •For ai safety engineers: Improve harmful input detection systems by integrating safety heads with formal robustness guarantees against small input changes.
- •For security teams in ai services: Use contraction-regularized state space models to certify and improve detection of jailbreak attempts in language models deployed in security-critical settings.
Authors
Omanshu Thapliyal
Abstract
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.