Step-Wise Refusal Dynamics in Autoregressive and Diffusion Language Models

AI summary

Overview

• Research area: AI safety and language model generation — specifically jailbreak robustness in autoregressive (AR) versus diffusion language models (DLMs).
• Technical level: Advanced. The paper assumes familiarity with transformer hidden states, autoregressive decoding, diffusion remasking, and jailbreak attack/defense terminology.
• Scope: A study of how step-by-step sampling mechanics shape refusal and recovery from harmful intermediate generations, plus an internal signal ("SRI") and a detector ("SRI Guard") built on it.

What This Paper Is About

Diffusion language models have been reported to resist jailbreaks better than autoregressive models, but it was unclear whether that advantage comes from the model's learned internal representations or from the diffusion sampling mechanism itself. The authors isolate sampling by holding model weights fixed and switching only the decoding strategy, showing that remasking lets models revise harmful intermediate tokens before they are committed. They then build an internal, step-wise representation of refusal dynamics to capture safety behavior that final text alone cannot reveal.

Key Contributions

1. Recovery-by-revision analysis: A step-level account, with empirical support, of why diffusion remasking promotes recovery from harmful intermediate generations — harmful predictions must both win the position-level argmax/top-k and be selected for commitment among all uncommitted positions, and committed harmless tokens reduce the likelihood of harmful continuations.
2. Fixed-weight jailbreak comparison: Evidence that switching from AR sampling to diffusion remasking improves jailbreak robustness under identical model weights for LLaDA and LLaDA-1.5, measured by Refusal Rate (RR) and Attack Success Rate (ASR), with additional confirmation via Harmfulness Score (HS) and StrongREJECT Score (SRS).
3. The SRI signal: The Step-Wise Refusal Internal Dynamics representation, a bounded per-step score in [0,1] built from last-layer hidden activations, mean pooling, cosine distance to harmless/harmful prototype centers, a temperature-scaled log-ratio, and a sigmoid.
4. SRI Guard: An inference-time anomaly detector trained only on benign SRI signals using an MLP autoencoder whose reconstruction error flags jailbreaks, reported to match or outperform jailbreak detection baselines (LlamaGuard, perplexity-based filtering, Self-Examine) with negligible overhead.

Main Findings

• Recovery is frequent and persistent in DLMs: Harmful Remasking Rate (HRR) ranges from 0.81 to 0.96 and Full Recovery Rate (FRR) from 0.63 to 0.73 across LLaDA (0.81 / 0.63), Dream (0.96 / 0.73), and LLaDA-1.5 (0.92 / 0.65).
• Sampling, not weights, drives the safety gain: With weights fixed, switching from AR sampling to diffusion remasking produced positive changes in both RR and ASR for LLaDA and LLaDA-1.5 across five attacks. LLaDA-1.5 gains ranged from +0.0 ΔRR (PAIR) to +52.0 ΔRR (Refusal Suppression), and from +2.0 ΔASR (PAIR) to +60.0 ΔASR (Flip Attack). LLaDA gains ranged from +7.0 ΔRR (PAIR) to +55.0 ΔRR (Refusal Suppression), and from +5.0 ΔASR (PAIR) to +36.0 ΔASR (Flip Attack).
• The gain is not a general quality advantage: Perplexity showed no consistent advantage for diffusion sampling; on average AR sampling achieved better perplexity (14.88 vs. 17.38).
• DLMs are more robust at comparable scale: Aggregated over all jailbreak attacks, DLM ASR stayed between 9% and 21% while AR models reached 48%–62%. DLM refusal rates reached 44%–67% versus 11%–46% for AR models.
• Aggregate refusal rate alone is misleading: Gemma and Dream had comparable aggregate refusal rates (46.2% vs. 44.4%), yet their attack success rates differed sharply (48.2% vs. 9.4%).
• Internal recovery tracks text-level robustness: With compliance threshold λ_c = 0.5 and refusal thresholds λ_r ∈ {0.5, 0.3, 0.1}, DLMs achieved consistently higher Internal Recovery Rate (IRR) than AR models, and switching from diffusion remasking to AR sampling produced a pronounced IRR drop under fixed weights.
• SRI separates safety-relevant structure better than alternatives: Average AUROC across models was 0.920 ± 0.047 for SRI (last layer), 0.864 ± 0.062 for SRI (middle layer), 0.566 ± 0.101 for SRI (first layer), 0.584 ± 0.138 for static activations at the first step, and 0.573 ± 0.221 for a text-based signal.
• SRI Guard is cheap: It operates between 150× and 300× faster than external moderation or self-reflection-based defenses, corresponding to roughly 0.01% overhead relative to the base model's generation time.
• Robustness to decoding choices: Both greedy and random remasking improved safety over AR sampling. Dynamic decoding (confidence threshold τ = 0.90) committed more harmful tokens earlier and reduced recovery slightly (LLaDA 0.73 HRR / 0.55 FRR; LLaDA-1.5 0.81 / 0.57) but still improved over AR generation.
• Recovery persists in hybrids and at larger scale: The hybrid AR–diffusion model SDAR reached HRR = 0.59 and FRR = 0.49; LLaDA-2 (16B) reached HRR = 0.78 and FRR = 0.64.

Methodology in Plain English

The authors first attack the problem at the text level. They assemble 600 prompts from WildJailbreak, JailbreakBench, and HarmBench, apply jailbreak attacks (Flip Attack, PAIR, Refusal Suppression, Random Search), and score outputs with an LLM-based judge plus standard metrics (RR, ASR, and later HS and SRS). To separate the effect of sampling from the effect of model weights, they take DLMs that can also run in an AR mode — LLaDA and LLaDA-1.5 — and change only the decoding strategy.

To see inside the generation process, they define a per-step score. At each step they mean-pool the last-layer hidden activations of the tokens generated so far, then compare that vector by cosine distance to two prototype centers: one averaged over harmless prompts and one over harmful prompts. The relative distances are combined into a log-ratio, divided by a temperature, and squashed through a sigmoid to produce a value between 0 and 1, where 1 means internally compliant and 0 means refusal-aligned. This sequence is the SRI signal. Signals were built using 400 harmless prompts from Alpaca and 400 harmful prompts from AdvBench, with T = 32 steps and τ = 0.1.

From this signal they derive the Internal Recovery Rate: how often a response spikes above a compliance threshold λ_c before the final step but ends below a refusal threshold λ_r. They then train SRI Guard, a small MLP autoencoder, only on benign SRI signals; at inference, a high reconstruction error means the generation looks unlike anything benign and is flagged as a jailbreak.

Why This Matters

Impact on research. The paper separates two explanations for diffusion robustness that were previously conflated — learned representations versus sampling dynamics — and shows that sampling alone accounts for a substantial part of the gap. It also argues that safety evaluation should consider generation trajectories, not only final outputs, and introduces a signal that makes those trajectories measurable and comparable across architectures.

Real-world applications.

• Content moderation and guardrail layers for deployed assistants, especially where latency budgets rule out external moderation models.
• Red-teaming and safety evaluation tooling that can score intermediate generation behavior rather than relying solely on output judges.
• Deployment decisions for open-weight models, where an operator needs to know whether switching a model's decoding mode materially changes its safety profile.
• Monitoring of emerging diffusion-based generation systems whose safety properties are not yet well characterized.

Industry relevance. SRI Guard does not modify the sampling process, adds roughly 0.01% overhead, and is trained only on benign signals, so it can be dropped into existing serving stacks without retuning. The fixed-weight results also mean a provider could change safety behavior on an already-shipped checkpoint by changing the decoder, which is a meaningful lever for deployment.

Future Directions

• Empirically testing integration of SRI Guard with other defenses. The authors explicitly note that although integration potential is strong, such empirical testing is not the focus of this work and is left for future work.
• Addressing the failure mode where SRI Guard cannot detect jailbreaks if the model shows no hesitation toward refusal at any point during generation.
• Extending the scaling analysis. The paper reports a LLaDA-2 (16B) ablation alongside a prior-work claim that LLaDA-2 extends results to the 100B-parameter regime; broader scaling evidence is raised but not fully resolved here.
• Further study of hybrid AR–diffusion architectures, which the appendix reports at HRR = 0.59 and FRR = 0.49, below full diffusion models.
• Reducing reliance on discrete, judge-based text-level evaluation, which the authors identify as low-resolution and computationally inefficient compared to operating in the hidden space.

Target Audience

AI safety and alignment researchers studying refusal behavior and jailbreak robustness; diffusion language model researchers; and machine learning engineers responsible for deploying or monitoring AR and diffusion-based generation systems in latency-sensitive settings. Readers will get the most from this paper with prior background in transformer internals, decoding strategies, and jailbreak terminology, though the text-level results in Sections 2.2 through 2.4 are accessible without it.

Authors’ abstract

Diffusion language models (DLMs) have recently emerged as a competitive alternative to autoregressive (AR) models, offering parallel decoding, competitive generation quality, and initial evidence of improved jailbreak robustness. Despite this progress, the role of sampling mechanisms in shaping refusal behavior remains poorly understood. To address this gap, we present a comprehensive study of step-wise refusal dynamics. We show that diffusion remasking can promote recovery from harmful intermediate generations, provide evidence that this behavior is tied to the sampling mechanism, and demonstrate that switching from AR to diffusion sampling improves jailbreak robustness, including under fixed model weights. To capture generation dynamics not observable at the text level, we propose the Step-Wise Refusal Internal Dynamics (SRI) signal. Consistent with our text-level findings, SRI shows that recovery fails primarily under AR sampling, with these failures often appearing anomalous relative to harmless generations in the SRI space. Based on this observation, we show that SRI enables a simple jailbreak detector that does not modify inference and generalizes to unseen attacks by training only on benign SRI signals. Our evaluation shows that this detector matches or outperforms existing jailbreak detection baselines while adding negligible overhead.