01 · Results

Jailbreak attacks succeed on 9 to 21% of prompts with diffusion models and on 48 to 62% with comparable autoregressive models.

Sections 2.1, 2.3, 2.4 and 4.3 of the paper, with Appendices B and C.3.

The paper compares six models in the 7 to 8B range: three autoregressive (LLaMA-3, Qwen-2.5, Gemma) and three diffusion (LLaDA, LLaDA-1.5, Dream). Each answers 100 plain harmful prompts and 100 prompts for each of five jailbreak attacks. An LLM judge scores attack success, and a refusal dictionary scores refusal. Pick an attack and a metric to compare the models.

Jailbreak robustness, autoregressive vs. diffusionTable 3 · six models, five attacks

What you learn

Across diverse jailbreak attacks, metrics, and model families, DLMs consistently exhibit better robustness than the evaluated AR models of comparable scale.

The same weights become safer when only the sampler changes.

The gap could still come from the weights. LLaDA and LLaDA-1.5 run under either sampler, so the paper switches the sampler and keeps the weights fixed. Safety improves across models, attacks and metrics, and no value in Table 2 is negative.

Improvement from switching AR sampling → diffusion remaskingTable 2 · percentage points, positive is safer

Every value is non-negative; PAIR moves least and Refusal Suppression most in refusal rate. Greedy, random, static and dynamic remasking variants all improve on AR sampling (Appendix A), and an LLM judge also finds higher refusal rates under diffusion sampling for every attack (Table 15).

What you learn

Under fixed model weights, enabling recovery by switching from AR sampling to diffusion remasking consistently improves jailbreak robustness across all the evaluated attacks and metrics.

SRI Guard matches or outperforms existing defenses at a fraction of their cost.

The compared defenses (LlamaGuard 3, Self-Examine, perplexity filtering) run an extra model or an extra pass over the text. SRI Guard reads the last-layer hidden states already produced during generation, so it adds only 0.01 to 0.04% to generation time. Averaged over the six models it has the highest refusal rate and the lowest attack success of the compared defenses, with a 5% false-positive rate on harmless prompts.

Defenses at a glanceTable 13 · mean over the six models · overhead on a log scale

02 · Use it

Run SRI on your own model with the released code.

Code by the authors; the Colab notebook accompanies this page.

You need a model whose last-layer hidden states you can read at every step, plus a harmless and a harmful prompt set for the anchors. The repository wraps LLaDA and LLaMA-3, each with a runnable notebook. A new model means subclassing BaseWrapper and implementing two methods. The Colab notebook runs the pipeline on a small open model and animates the signal like the cover.

  1. Anchors. Generate from a harmless and a harmful prompt set and store the mean last-layer activation of each set at every step. The paper uses 400 prompts per set.
  2. Signal. For a new prompt, mean-pool the last-layer activations of the generated tokens at each step, take the cosine distance to both anchors and pass the temperature-scaled log-ratio (τ = 0.1) through a sigmoid. The T = 32 values form the SRI signal.
  3. Guard (optional). Fit a small autoencoder on harmless signals and flag any generation whose reconstruction error exceeds the 99% quantile of a held-out harmless set.
Open source

Generate SRI signals for your own model.

Clone the repository, build the two anchor centers from your prompt sets and generate a step-wise SRI signal for any prompt.

# 1. clone and install
git clone --recurse-submodules https://github.com/ElironRahimi/sri-signal
cd sri-signal && pip install -r requirements.txt

# 2. wrap a model, build the anchor centers, generate signals
from utils import read_alpaca, read_advbench, GenerationProfile
from Models.LLaDAWrapper import LLADAModelWrapper
modelWrapper = LLADAModelWrapper(
    generation_profile=GenerationProfile(), device="cuda"
)
modelWrapper._load_model()
modelWrapper.generateCenter(
    harmless_prompts=read_alpaca(400), harmful_prompts=read_advbench(400)
)
signals = modelWrapper.generateSignal(prompts=test_prompts)
# one SRI signal per prompt: T values in (0, 1), 1 compliant, 0 refusal-aligned
LLaMA-3Qwen-2.5Gemma LLaDALLaDA-1.5Dream
03 · How it works

Remasking can revise a harmful token before it is committed, and SRI reads this process from the hidden states.

Sections 2.2 and 3 of the paper.

A model that has started a harmful answer should still be able to stop. Autoregressive decoding cannot revise: each token is committed as soon as it is generated, so an early harmful token shapes the rest of the response. Remasking diffusion predicts every uncommitted position at each step and commits only some of them. A harmful prediction is therefore often replaced before it is committed, and each committed harmless token makes harmful continuations less likely.

Autoregressive samplingSure,hereis▇▇▇▇→▇▇▇░░░
Diffusion remaskingSure,▇▇▇▇sorry,→Icannot

Solid pills are committed tokens and dashed pills are predictions that can still change. The diagram is a schematic; the cover animates the same mechanism step by step.

To study recovery, the paper needs to know at every step whether the generation leans toward compliance or refusal. The text cannot tell: two generations can look identical while evolving very differently inside the model. SRI compares the last-layer activations at each step with those of typical harmless and harmful generations. The result is a score from 1, compliant, to 0, refusal-aligned, and the T = 32 steps form the signal.

1

Anchor the space

Mean last-layer activation of harmless and of harmful generations, stored for every step.

2

Read the step

Mean-pool the last-layer activations of the tokens generated so far into φt.

3

Score it

Cosine distance to both anchors, log-ratio, sigmoid: σt runs from 1 (compliant) to 0 (refusal-aligned).

4

Guard

An autoencoder trained on harmless SRI signals flags anomalous trajectories. It uses no jailbreak data and leaves sampling unchanged.

What you learn

The remasking diffusion sampling mechanism naturally promotes recovery from harmful intermediate generations.

04 · Details

What makes the signal work, and where its limits are.

Sections 2.2, 3.1, 3.3 and 4.2 to 4.4 of the paper, with Appendices A, C and E and the Limitations section.

Three design choices carry the signal.

SRI uses internal activations, reads them from the last layer and keeps the whole trajectory over steps. Giving up any one of these lowers the mean AUROC, and a text-based signal reaches only 0.57 on average, with large variance across models.

Ablation of the SRI design choicesTable 4 · mean AUROC across the six models

Recovery by revision is frequent and persistent.

The Harmful Remasking Rate (HRR) is how often a harmful intermediate generation is later revised, and the Full Recovery Rate (FRR) is how often the final output is then harmless. Across LLaDA, LLaDA-1.5 and Dream, HRR is 0.81 to 0.96 and FRR is 0.63 to 0.73.

Recovery-by-revision statisticsTable 1 · jailbreak test set, LLM-judged harmfulness

HRR: how often an intermediate harmful generation is later revised. FRR: how often the final output is then non-harmful. Recovery persists in the hybrid model SDAR and in LLaDA-2 at 16B (Appendix A).

Four stacked panels from the paper. A jailbreak prompt with an attack prefix asks a harmful chemistry question; the intermediate LLaDA output at step 16 is labelled harmful, the output at step 25 is labelled transition as refusal language begins to appear, and the final output at step 32 is labelled safe and recovered, an explicit refusal.

Figure 1. Recovery in LLaDA: harmful intermediate tokens are revised into a safe refusal, the mechanism the cover animates.

Internal recovery tracks the robustness gap.

The Internal Recovery Rate (IRR) counts generations that are internally compliant at some step and refusal-aligned at the end. Diffusion models recover more often at every refusal threshold, and running LLaDA or LLaDA-1.5 under autoregressive sampling lowers IRR by 0.29 to 0.48.

Bar chart of Internal Recovery Rate per model at three refusal thresholds. The autoregressive models Qwen, LLaMA-3 and Gemma stay below about 0.15 at every threshold, while the diffusion models LLaDA, Dream and LLaDA-1.5 reach roughly 0.4 to 0.6 at the loosest threshold and stay above the autoregressive models as the threshold tightens.

Figure 5. Per-model IRR at three refusal thresholds, AR in blue, diffusion in red. Diffusion models recover more often at every threshold.

Grouped bars for LLaDA and LLaDA-1.5 at three refusal thresholds. Under AR sampling of the same weights the Internal Recovery Rate stays below about 0.17, while under diffusion sampling it reaches 0.39 to 0.63; the labelled differences range from 0.29 to 0.48.

Figure 6. The same weights under both samplers. Switching LLaDA or LLaDA-1.5 to AR sampling lowers IRR at every threshold, consistent with Table 2.

What you learn

Stronger internal recovery is associated with higher HRR and FRR, improved jailbreak robustness, and diffusion rather than AR sampling.

Jailbreak generations look different in SRI space.

Harmless and refused generations produce smooth, low-variance trajectories, while jailbreak generations are volatile and often move toward refusal without converging, mostly under autoregressive sampling. Under Linear Discriminant Analysis the groups occupy distinct regions of SRI space in most cases, imperfectly separable yet stable enough for a detector.

Explore the SRI space, one model at a timeFigures 12 and 13 · left: SRI signal under a jailbreak prompt · right: LDA projection of the SRI space
SRI signal over 32 generation steps for LLaMA-3 under a jailbreak prompt. The harmless and refusal reference trajectories stay flat near 1 and 0, while the jailbreak trajectory drops into the shaded refusal region within the first steps and climbs back into the compliance region.
Two-dimensional LDA projection of the SRI space for LLaMA-3: harmless and harmful responses form neighboring but distinct clusters on the right, and refusal responses cluster far to the left.
Autoregressive

How the score is computed.

Cosine distances → calibrated log-ratio → SRI dtk = 1 − ⟨φt, μtk⟩ / (‖φt‖ ‖μtk‖),   k ∈ {harmless, harmful}
ℓt = [ log(dtharmful + ε) − log(dtharmless + ε) ] / τσt = sigmoid(ℓt)
σt = 1 is an internally compliant state, σt = 0 a refusal-aligned state, intermediate values transitional. τ = 0.1 keeps signals calibrated and sensitive to transitions, ε ensures stability, and T = 32 steps form one signal.

SRI Guard holds up model by model.

Defenses compared under one protocolTable 13 · pick a model · overhead relative to base generation time

The guard holds under diffusion-specific attacks.

Under PAD and DIJA, two white-box attacks built for diffusion models, SRI Guard on LLaDA lowers Attack Success Rate from 0.53 to 0.29 and from 0.76 to 0.29, where DiffuGuard, a diffusion-specific defense, reaches 0.45 and 0.59 (Table 14). SRI Guard also assumes less, since DiffuGuard requires the original pre-manipulated prompt.

SRI Guard does not replace alignment training.

SRI Guard is an inference-time detection mechanism and does not replace alignment training. It relies on internal refusal signals, so it can miss a jailbreak when the model shows no hesitation toward refusal at any step. Its negligible overhead and independence from sampling make it a natural complement to other defenses, a combination the paper leaves to future work.

Four questions the results leave open.

  • Other layers. Middle-layer SRI separates substantially better than first-layer SRI and slightly below the last layer, and the last layer wins in four of six models, so the choice of layer could be revisited (Section 4.4 and Appendix E.1).
  • Other model families. The hybrid model SDAR recovers and LLaDA-2 at 16B preserves the SRI geometry, which suggests hybrid and larger models as next targets (Appendices A and E.4).
  • Combined defenses. SRI Guard leaves sampling untouched and adds negligible overhead, so stacking it with other defenses is untested and open (Limitations).
  • Black-box targets. SRI from a white-box surrogate (Dream, 7B) already defends a black-box LLaDA-2 (Table 24), which leaves the choice of surrogate and the limits of transfer to future work (Appendix E.6).