Step-Wise Refusal Dynamics in Autoregressive and Diffusion Language Models

TL;DR

This work addresses safety in text generation by comparing autoregressive decoding with diffusion remasking, revealing that sampling mechanics significantly shape refusal and jailbreak robustness beyond learned representations. It introduces the Step-Wise Refusal Internal Dynamics (SRI) signal, an internal trajectory-based representation that captures step-wise safety states and enables efficient, model-agnostic detection of unsafe generations via SRI Guard. The authors show that harmful generations exhibit incomplete internal recovery in SRI space, and that diffusion sampling affords higher internal recoverability and stronger text-level refusals than AR decoding. The proposed framework and detectors achieve comparable or superior defense performance with over 100x lower inference overhead, offering practical, scalable tools for reinforcing safety across AR and diffusion-based language models.

Abstract

Diffusion language models (DLMs) have recently emerged as a promising alternative to autoregressive (AR) models, offering parallel decoding and controllable sampling dynamics while achieving competitive generation quality at scale. Despite this progress, the role of sampling mechanisms in shaping refusal behavior and jailbreak robustness remains poorly understood. In this work, we present a fundamental analytical framework for step-wise refusal dynamics, enabling comparison between AR and diffusion sampling. Our analysis reveals that the sampling strategy itself plays a central role in safety behavior, as a factor distinct from the underlying learned representations. Motivated by this analysis, we introduce the Step-Wise Refusal Internal Dynamics (SRI) signal, which supports interpretability and improved safety for both AR and DLMs. We demonstrate that the geometric structure of SRI captures internal recovery dynamics, and identifies anomalous behavior in harmful generations as cases of \emph{incomplete internal recovery} that are not observable at the text level. This structure enables lightweight inference-time detectors that generalize to unseen attacks while matching or outperforming existing defenses with over $100\times$ lower inference overhead.

Figures (13)

• Figure 1: Recovery from harmful intermediate content during diffusion generation in LLaDA. Harmful tokens produced at intermediate steps are iteratively revised across diffusion steps, enabling recovery to a safe final output without committing to a fixed prefix.
• Figure 2: Example of incomplete internal recovery captured by the Step-Wise Refusal Internal Dynamics (SRI) signal in Qwen (AR model). The SRI signal for a harmful generation (blue) is shown alongside reference SRI signals for harmless (green) and refused (red) responses. Shaded regions indicate compliance-aligned (green) and refusal-aligned (red) states, with the dashed line denoting the text-level refusal signal; although the harmful response is not flagged at the text level, the SRI signal reveals anomalous behavior indicative of incomplete internal recovery.
• Figure 3: Per-model IRR with compliance threshold $\lambda_c=0.5$ and refusal thresholds $\lambda_r \in \{0.5, 0.3, 0.1\}$. AR models are shown in blue and diffusion models in red.
• Figure 4: IRR under different sampling strategies for the same model weights. Compliance threshold $\lambda_c=0.5$ and refusal thresholds $\lambda_r \in \{0.5, 0.3, 0.1\}$.
• Figure 5: LDA projection of the learned SRI latent space for a representative diffusion model (LLaDA) and AR model (LLaMA).

Theorems & Definitions (8)

• Proposition 3.1: Structural asymmetry between AR and diffusion sampling
• Definition 3.2: Harmful Remasking Rate (HRR)
• Definition 3.3: Full Recovery Rate (FRR)
• Definition 4.1: Internal Recovery Rate (IRR)
• Proposition 1.1: No recovery under AR decoding
• proof
• Proposition 1.2: Recovery is possible under remasking diffusion
• proof