The Safety-Aware Denoiser for Text Diffusion Models

Recent work on text diffusion models offers a promising alternative to autoregressive generation, but controlling their safety remains underexplored. Existing safety approaches are geared toward autoregressive models and typically rely on post-hoc filtering or inference-time interventions. These are inadequate for effectively addressing safety risks in text diffusion models. We propose the Safety-Aware Denoiser (SAD), a safety-guidance framework in text diffusion models. The SAD modifies the iterative denoising process such that the text sample at the final denoising step is steered toward provably safe regions of the text space. This inference-time method can integrate safety constraints into the denoiser, avoiding computationally expensive retraining of the underlying diffusion model and enabling flexible, lightweight safety guidance. We evaluate the safety of the generated text using the SAD, with respect to hazard taxonomy, memorization, and jailbreak. Experimental results show that SAD substantially reduces unsafe generations while preserving generation quality, diversity, and fluency, outperforming existing methods. These results demonstrate that our safety guidance during denoising provides an effective and scalable mechanism for enforcing safety in text diffusion models.

Large language models are increasingly built using a new approach called text diffusion, which generates entire responses at once through a process of iterative refinement, somewhat like slowly bringing a blurry image into focus. While these models are becoming very capable, they carry serious safety risks as they can produce harmful content, reproduce sensitive training data, and are vulnerable to adversarial "jailbreak" attacks designed to bypass their safeguards. Existing safety tools were designed for older-style models and do not transfer well to this new paradigm. To account for these shortcomings, we introduce the Safety-Aware Denoiser (SAD), a method that steers text diffusion models toward safe outputs by intervening during the generation process itself rather than filtering outputs after the fact. SAD requires no retraining of the underlying model. Instead, it uses a small collection of known harmful examples to guide generation away from unsafe content at each refinement step. We show that SAD substantially reduces harmful outputs across toxicity, memorization, and jailbreak settings, while preserving fluency and generation quality, offering a practical and principled safety layer for this emerging class of AI systems.