The Tokens Strike Back: Safety-Aware Denoising with Directional Embedding-based Revision
The paper introduces SADDER, a training-free inference-time method for text diffusion models that improves on Safety-Aware Denoising by using token-level safety signals in embedding/activation space to gate which positions get revised toward safer outputs. Its neat contribution is tying this practical gating trick to a theoretical denoising target via position-wise distributions, and experiments suggest a better safety–quality–cost trade-off on MDLMs and USDMs, though its scope still depends on the availability and reliability of the safety signal and is evaluated mainly in unsafe-continuation and jailbreak settings.
Abstract
The Safety-Aware Denoiser (SAD) is a recent, training-free, safety-guidance framework for text diffusion models. During inference, SAD steers text samples toward safe regions of the token probability space. In this paper, we propose $\mathit{SADDER}$ (Safety-Aware Denoising with Directional Embedding-based Revision), which improves the safety and quality trade-offs of SAD by exploiting a safety-guidance signal in the embedding (or activation) space for each token position, then updating the token positions using the SAD update rule multiplied by the guidance signal. This multiplication acts as a gating mechanism: tokens with a high unsafe signal are more likely to be modified during inference following the SAD update rule. However, such an update does not automatically preserve the conditional mean, i.e., the denoiser of a fixed target distribution. We therefore theoretically derive new position-wise target distributions for which the modified denoiser is exact. Furthermore, we derive conditions under which the position-wise target distributions match a global target distribution. We test our $\mathit{SADDER}$ method on popular Masked Diffusion Language Models (MDLMs) and Uniform-State Diffusion Models (USDMs) across unsafe continuation and jailbreak settings, and show that $\mathit{SADDER}$ achieves a better balance between safety, generation quality, and inference cost than existing methods.