💬 NLP

Adaptive Steering and Remasking for Safe Generation in Diffusion Language Models

This paper proposes a plug-and-play inference-time defense framework for Diffusion Language Models that combines contrastive safety direction-based detection with adaptive steering and token remasking to effectively mitigate safety vulnerabilities during iterative denoising while preserving generation quality.

Original authors: Yejin Lee, Yo-Sub Han

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yejin Lee, Yo-Sub Han

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Diffusion Language Model (DLM) as a team of artists working together to paint a picture based on a prompt. Unlike a traditional writer who paints one brushstroke at a time in a straight line, this team starts with a canvas completely covered in static (noise). They work in rounds, gradually clearing away the static to reveal the image.

The problem, as the paper explains, is that if a "bad" idea (like a harmful instruction) sneaks into the picture during the early rounds of clearing static, it gets locked in. Because the artists keep refining the image based on what they see, that early bad idea spreads and eventually ruins the whole painting, even if the artists try to fix it later.

Here is how the authors' new method, Adaptive Steering and Remasking, fixes this, explained through simple analogies:

1. The Problem: The "Bad Seed" in the Garden

In traditional AI, if you ask a dangerous question, the AI might refuse immediately. But in these "diffusion" models, the AI starts with a blank slate and fills it in slowly.

The Vulnerability: If a harmful word (like "bomb") appears early in the process, the model treats it as a fact and builds the rest of the sentence around it. By the time the model realizes it's doing something wrong, the "bad seed" has grown into a full tree, and the model can't easily cut it down.
The Old Fix: Previous methods tried to be like a strict gardener who, upon seeing any weed, just stopped watering the whole garden. This kept the weeds away, but it also killed the flowers, resulting in empty or broken sentences.

2. The Solution: A Smart Guide and a Precision Pruner

The authors propose a two-step safety system that acts like a smart guide and a precision tool, working while the painting is being made, not after.

Step 1: The "Compass" (Adaptive Steering)

The team first creates a Contrastive Safety Direction (CSD). Think of this as a compass that points specifically toward "Safe" and away from "Harmful."

How it works: They show the model examples of safe answers and harmful answers. They calculate the mathematical difference between them to create this "compass."
The Action: During the early rounds of the painting process (when the image is still mostly static), the system checks the model's direction. If the model starts leaning toward the "Harmful" side of the compass, the system gently pushes it back toward the "Safe" side.
The Analogy: Imagine you are walking through a foggy forest. If you start heading toward a cliff (harmful), a guide gently nudges your shoulder to steer you back toward the path (safe) before you get too close to the edge. This happens early, so you don't get lost.

Step 2: The "Precision Pruner" (Remasking)

Even with the compass, sometimes a bad word slips through. This is where the second step comes in.

How it works: The system scans the words that have been revealed so far. If it spots a word that is strongly aligned with the "Harmful" direction, it doesn't delete the whole sentence. Instead, it puts a "mask" (a blank cover) over just that specific bad word.
The Action: The model is then asked to repaint only that specific spot, trying to find a safer word to replace it.
The Analogy: Imagine a sculptor is carving a statue. If they accidentally carve a jagged, ugly piece of stone, they don't smash the whole statue. They just chip away that specific ugly piece and smooth it over with fresh clay. This keeps the rest of the beautiful statue intact.

3. Why This is Better

No Heavy Lifting: The authors didn't need to retrain the model (teach it new lessons from scratch). They just added this "compass" and "pruner" as a plug-in tool that works while the model is thinking.
Quality vs. Safety: Old methods were like using a sledgehammer to kill a fly; they stopped the danger but broke the furniture (the quality of the text). This new method is like using a fly swatter; it stops the danger but leaves the furniture (the story or code) perfectly fine.
The Results: In their tests, this method stopped "jailbreak" attacks (attempts to trick the AI into being unsafe) from succeeding almost entirely (dropping success rates to less than 1% in some cases) while keeping the AI's ability to write good stories and solve math problems almost the same as before.

Summary

The paper argues that to keep these "iterative" AI models safe, you can't just wait until the end to check for mistakes. You have to guide the process gently at the very beginning (Steering) and fix specific bad parts as they appear (Remasking). This ensures the final output is both safe and high-quality, without needing to rebuild the AI from the ground up.

Technical Summary: Adaptive Steering and Remasking for Safe Generation in Diffusion Language Models

Problem Statement
Diffusion Language Models (DLMs) generate text through iterative denoising and bidirectional refinement, offering advantages in parallel generation and global context reasoning over autoregressive models. However, this iterative paradigm introduces unique safety vulnerabilities. Unlike autoregressive models where safety is often determined at the token level, DLMs are susceptible to harmful signals emerging and propagating during intermediate denoising steps. Recent studies indicate that early denoising trajectories strongly influence final outputs; once a harmful token is unmasked, it can steer the entire generation trajectory toward unsafe responses. Existing defense methods, primarily adapted from autoregressive frameworks or relying on global remasking, often fail to control the semantic denoising trajectory explicitly. This results in a trade-off where aggressive suppression of harmful tokens leads to low-quality outputs, empty responses, or a collapse in generation utility, while insufficient intervention fails to prevent jailbreaks.

Methodology
The authors propose an inference-time defense framework that operates without additional fine-tuning, combining Adaptive Steering and Token Remasking. The framework functions in two complementary phases:

Contrastive Safety Direction (CSD) Construction:
The method first constructs a latent vector representing the semantic boundary between harmful and safe generations. Using paired samples (identical prompts with harmful vs. safe refusal responses), the authors compute the difference in hidden representations between the two classes. This difference vector, the CSD, captures the "harmful semantic direction" in the model's latent space.
Phase 1: Early-Step Adaptive Steering:
Recognizing that early denoising steps critically determine the final trajectory, the framework applies a steering intervention only during the initial fraction ( $\rho T$ ) of denoising steps.
- Mechanism: The model's hidden states for masked tokens and prompt representations are projected onto the CSD. The component of the hidden state aligned with the harmful direction is suppressed using a scaling factor $\beta$ .
- Adaptivity: The steering strength is modulated based on the estimated degree of harmfulness (the projection magnitude). This guides the denoising trajectory toward safe semantic regions before harmful semantics become stabilized, without over-constraining later refinement stages.
Phase 2: Token Remasking:
Following the initial steering, a refinement stage addresses residual harmful tokens that may have emerged.
- Detection: The framework computes alignment scores for individual tokens against the CSD. Tokens exceeding a threshold are identified as harmful.
- Remasking: Instead of global suppression, the method selectively remasks only the identified harmful tokens. The model then regenerates these positions from the corrected semantic context. This iterative process continues for a fixed number of steps or until no harmful tokens remain, ensuring fluent and safe outputs.

Key Contributions

Inference-Time Defense Framework: The paper introduces a plug-and-play module that integrates semantic steering and selective remasking, requiring no parameter updates or safety fine-tuning of the base DLM.
Identification of Critical Control Points: The work identifies early denoising trajectories as the primary control point for DLM safety. It demonstrates that partial-step steering effectively redirects generation paths before harmful behaviors stabilize.
Balanced Safety and Utility: The proposed method addresses the safety-utility trade-off inherent in prior remasking approaches. By steering the trajectory early and only remasking specific tokens later, it mitigates harmful generations while preserving generation quality and fluency.

Experimental Results
The framework was evaluated on two backbone DLMs (LLaDA-7B and Dream-8B) across multiple benchmarks (JailBreakBench, AdvBench, DIJA, PAP, and Prefix attacks).

Jailbreak Defense: The method significantly reduced Attack Success Rates (ASR). On the Dream model, it achieved an average ASR of 0.64% on AdvBench and 8.00% on JailBreakBench, outperforming existing baselines like DiffuGuard and Self-reminder in terms of balanced robustness across different attack types.
Quality Preservation: Unlike aggressive remasking methods that caused generation collapse (high rates of empty or broken sentences), the proposed approach maintained generation quality close to the original model. On general capability benchmarks (TruthfulQA, MATH-500, MMLU), the method preserved utility better than DiffuGuard, though some degradation compared to the vanilla model was observed, consistent with the trade-off in safety interventions.
Ablation Studies: Experiments confirmed that both phases are necessary. Removing early steering increased ASR significantly (from 18% to 54%), and removing remasking also degraded performance, highlighting that both trajectory guidance and token-level correction are essential.

Significance
The paper claims that DLM safety is not solely a function of the final output distribution but is deeply dependent on the intermediate denoising trajectory. The proposed step-wise intervention framework demonstrates that safety can be maintained throughout the generation process without compromising the model's inherent generative capabilities. By leveraging a contrastive safety direction to steer the latent space and selectively remasking only when necessary, the method offers a robust solution to the unique vulnerabilities of diffusion-based text generation, providing a pathway for safer deployment of DLMs without the need for costly retraining.