Adaptive Steering and Remasking for Safe Generation in Diffusion Language Models

2026-05-13

AI Summary

This work addresses the vulnerability of diffusion language models to harmful intermediate tokens during iterative denoising, which can compromise the safety of final outputs. The authors propose a plug-and-play defense framework operating at inference time that introduces a contrastive Safe Gradient Direction (SGD) as a semantic harmlessness criterion in the latent space. At each denoising step, the method dynamically identifies and re-masks harmful tokens while employing an adaptive guidance mechanism to continuously steer the generation trajectory toward safer outcomes. This approach achieves stepwise intervention without requiring model fine-tuning, significantly enhancing safety by reducing jailbreak attack success rates to 0.64%, all while preserving generation quality on par with the original model.

Abstract

Diffusion Language Models (DLMs) provide a promising alternative to autoregressive language models by generating text through iterative denoising and bidirectional refinement. However, this iterative generation paradigm also introduces unique safety vulnerabilities when harmful tokens generated at intermediate denoising steps propagate through subsequent refinement processes and eventually induce unsafe outputs. While there are a few attempts to remedy this issue, they either fail to generate safe outputs or generate safe yet low-quality outputs. This motivates us to propose an inference-time defense framework based on the step-wise intervention during the denoising process, which then improves the safety without compromising the output quality. The key component of our framework is a contrastive safety direction (SGD), a latent direction that captures the semantic boundary between harmful and safe generations. We leverage SGD to assess the alignment of generated tokens with harmful semantics at each denoising step. When harmful alignment is detected, our method remasks the corresponding tokens and resumes the denoising process with adaptive steering, where the steering strength is modulated according to the estimated degree of harmfulness. As a plug-and-play module, our method circumvents the need for additional fine-tuning and can be directly incorporated into off-the-shelf diffusion models. The experimental results show that our approaches reduce jailbreak success rates to 0.64% while preserving generation quality close to the original model performance. This confirms the effectiveness of step-wise intervention for safe diffusion language model generation. Our code is available at https://github.com/leeyejin1231/DLM_Steering_Remasking.

Problem

• Diffusion Language Models
• Safety
• Harmful Tokens
• Iterative Denoising
• Unsafe Generation

Innovation

• Diffusion Language Models
• Safety Steering
• Remasking
• Contrastive Safety Direction
• Inference-time Intervention