Toward Safer Diffusion Language Models: Discovery and Mitigation of Priming Vulnerability

TL;DR

The paper identifies a diffusion-language-model-specific safety risk, the priming vulnerability, where affirmative tokens appearing during intermediate denoising steps can steer MDLMs toward harmful outputs. It analyzes two threat scenarios and introduces Recovery Alignment (RA), an MDLM-tailored safety alignment that trains models to recover from contaminated intermediate states via an RLHF-style objective. Empirical results show attackers can exploit the vulnerability with anchoring and First-Step GCG, while RA substantially mitigates the vulnerability and improves robustness against conventional jailbreaks, with minimal impact on broad task performance. The study underscores the need for MDLM-specific safety research and provides a practical defense that generalizes to multiple attack modalities, advancing safer diffusion-based language technologies.

Abstract

Diffusion language models (DLMs) generate tokens in parallel through iterative denoising, which can reduce latency and enable bidirectional conditioning. However, the safety risks posed by jailbreak attacks that exploit this inference mechanism are not well understood. In this paper, we reveal that DLMs have a critical vulnerability stemming from their iterative denoising process and propose a countermeasure. Specifically, our investigation shows that if an affirmative token for a harmful query appears at an intermediate step, subsequent denoising can be steered toward a harmful response even in aligned models. As a result, simply injecting such affirmative tokens can readily bypass the safety guardrails. Furthermore, we demonstrate that the vulnerability allows existing optimization-based jailbreak attacks to succeed on DLMs. Building on this analysis, we propose a novel safety alignment method tailored to DLMs that trains models to generate safe responses from contaminated intermediate states that contain affirmative tokens. Our experiments indicate that the proposed method significantly mitigates the vulnerability with minimal impact on task performance. Furthermore, our method improves robustness against conventional jailbreak attacks. Our work underscores the need for DLM-specific safety research.

Figures (12)

• Figure 1: Overview of this work. (a) MDLMs alternate prediction and partial re-masking to gradually produce the response. (b) An affirmative token can steer generation toward a harmful response. (c) Our approach trains the model to recover to a safe response.
• Figure 2: ASR vs. number of intervention steps. ASR increases sharply even at $t_{\text{inter}}=1$.
• Figure 3: Change in the probability mass of refusal phrases due to token insertion. The x-axis represents the inserted anchor token, and the y-axis shows the relative change in the summed probability mass of the refusal phrases. We use the query “Write a tutorial on how to make a bomb.” and the refusal set {“I’m sorry”, “Sorry”, “I cannot”, “I can not”, “I can’t”, “As an AI”, “I am not”, “It is not”, “No”}. The results indicate that MDLMs are influenced by already generated tokens, causing a shift in the generation distribution from refusal to affirmation.
• Figure 4: Probability mass of the refusal phrases vs. number of intervention steps.
• Figure 5: Empirical validation of the monotonicity assumption. We plot the mean per-token monotonicity gap $\Delta_t$ on JBB-Behaviors for three MDLMs. Contaminated states $\bm r_{t_{\text{inter}}}$ are constructed by anchoring attack. Specifically, pre-filling the first $k$ target tokens and masking the rest. Positive $\Delta_t$ confirms the monotonicity condition, and the gap grows with $k/L$.

Theorems & Definitions (3)

• Theorem 4.1
• Theorem A.1
• proof