💬 NLP

A2D: Any-Order, Any-Step Safety Alignment for Diffusion Language Models

A2D introduces a token-level safety alignment method for diffusion language models that leverages randomized masking to emit immediate [EOS] refusal signals, effectively neutralizing any-order and any-step attacks like DIJA while enabling real-time monitoring and significantly faster safe termination.

Original authors: Wonje Jeung, Sangyeon Yoon, Yoonjun Cho, Dongjae Jeon, Sangwoo Shin, Hyesoo Hong, Albert No

Published 2026-02-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Wonje Jeung, Sangyeon Yoon, Yoonjun Cho, Dongjae Jeon, Sangwoo Shin, Hyesoo Hong, Albert No

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

The Big Picture: A New Kind of AI Writer

Imagine most AI chatbots (like the ones you talk to every day) as assembly line workers. They build a sentence one word at a time, strictly from left to right. If you ask them something dangerous, they usually check the first few words, decide "No, that's bad," and stop immediately.

But there is a new type of AI called a Diffusion Language Model (dLLM). Think of this AI not as an assembly line, but as a sculptor working on a block of marble. Instead of writing word-by-word, it starts with a blank page full of question marks (masks) and gradually fills them in. It can fill in the end of a sentence before the beginning, or the middle before the start. It can work on the whole story all at once.

The Problem: The "Backdoor" Attack

Because this sculptor can fill in words in any order, it has a new weakness.

Imagine a bad actor wants to trick the sculptor into writing a guide on how to steal a car.

The Old Trap: If the sculptor is an assembly line worker, the bad actor tries to trick them at the very start.
The New Trap (DIJA Attack): With the sculptor, the bad actor can fill in the first half of the story with harmless text, then sneak in a "trap" in the middle, or fill in the end first. They use a template that looks like a normal request but hides the dangerous parts in the blanks.

The paper found that current safety training for these sculptors is "shallow." It's like a security guard who only checks your ID when you walk through the front door. If you manage to sneak past the door and start walking down the hallway, the guard stops paying attention. The AI might say "No" to the first word, but if you trick it into filling in the 10th word with something dangerous, it forgets to say "No" again.

The Solution: A2D (Any-Order, Any-Step Defense)

The authors created a new safety method called A2D. Here is how it works, using a simple analogy:

The "Red Stop Sign" Token
Instead of just training the AI to say "I can't do that" at the very beginning, A2D teaches the AI a special "Stop Sign" token (called [EOS]).

The Training: They take the AI and show it dangerous sentences. But instead of letting it finish the sentence, they cover the dangerous parts with question marks and tell the AI: "If you see a question mark here, and the sentence is dangerous, you must replace that question mark with a Stop Sign immediately."
The Result: The AI learns that at any point in the sentence, if it senses something harmful, it should instantly slap a Stop Sign on that spot.

Why is this special?

Any-Order: It doesn't matter if the AI is writing the first word or the last word. If a dangerous idea pops up, the Stop Sign appears.
Any-Step: It doesn't matter if the danger appears at step 1 or step 50. The safety guard is awake the whole time, not just at the door.

The "Fill-in-the-Blank" Test

To prove their method works, the researchers invented a super-hard test called FITS (Fill-in-the-Sentence).

The Scenario: Imagine a bad guy writes a 99% complete guide on how to make a bomb. They leave out just one sentence (the final step) and ask the AI to fill it in.
The Result: Old safety methods failed miserably here. They were so focused on the beginning of the text that they let the AI finish the bomb guide.
A2D's Performance: A2D looked at that single missing sentence, realized it was dangerous, and immediately put up the Stop Sign. It blocked the attack almost 100% of the time.

The Bonus: Real-Time Safety Radar

Because the AI is trained to put up a Stop Sign whenever it sees danger, the probability of that Stop Sign appearing acts like a safety radar.

Early Rejection: The researchers found that if they check the AI's "Stop Sign probability" right at the very first step, they can tell if the whole request is bad before the AI even writes a single word.
Speed: This allows the AI to reject bad requests 19.3 times faster than before. It's like a bouncer at a club who can tell you're not on the list before you even reach the door, saving everyone time.

Does it ruin the AI's ability to be helpful?

The paper tested this extensively. They asked the AI to do math, write code, and answer general questions.

The Verdict: A2D kept the AI smart and helpful. It didn't make the AI refuse to answer normal questions (like "How do I kill a python process?" which sounds dangerous but is actually about coding).
Comparison: Other safety methods were too sensitive and refused to answer harmless questions. A2D was precise: it only stopped the bad stuff.

Summary

The paper introduces A2D, a safety upgrade for a new type of AI that writes in any order.

Old Safety: A guard at the front door who gets tired after a few minutes.
A2D Safety: A security system that puts up a "Stop" sign instantly, anywhere, anytime, if it smells danger.
Outcome: It stops sneaky attacks that slip past the front door, works 19x faster at saying "no," and keeps the AI helpful for everyone else.

Technical Summary: A2D – Any-Order, Any-Step Safety Alignment for Diffusion Language Models

1. Problem Statement

Diffusion Large Language Models (dLLMs) represent a significant departure from autoregressive (AR) models by enabling any-order decoding, where tokens are generated iteratively in arbitrary positions rather than a fixed left-to-right sequence. While this flexibility supports complex reasoning and structured content synthesis, it introduces unique safety vulnerabilities that render conventional alignment methods ineffective.

The core problem identified is shallow alignment. Similar to AR models, current dLLMs exhibit strong safety signals only at the initial generation steps, which rapidly decay as decoding progresses. This limitation creates a large attack surface:

Any-Order Vulnerability: Harmful content can be injected at arbitrary positions during the decoding trajectory, bypassing initial refusals.
Any-Step Vulnerability: Template-based attacks, such as DIJA (Wen et al., 2025), exploit this decay by interleaving adversarial text with [MASK] tokens in partially prefilled prompts. These attacks target the model after several decoding steps when safety signals have already faded, inducing unsafe completions that response-level alignment methods fail to prevent.

Existing baselines, including Refusal Training (RT), Supervised Fine-Tuning (SFT), and Variance-Reduced Preference Optimization (VRPO), rely on response-level supervision. Consequently, they struggle to maintain safety under diverse decoding strategies (e.g., random, confidence-based) and fail against attacks that bypass the initial tokens.

2. Methodology: A2D (Any-Order, Any-Step Defense)

The authors propose A2D, a token-level alignment method designed specifically for the flexible decoding process of dLLMs. Unlike prior approaches that supervise refusals only at the full-response level, A2D aligns the model to emit a specific [EOS] (End of Sequence) token whenever harmful content is encountered at any masked position.

Core Mechanism

Token-Level Supervision: During training, if a harmful span is present in the target sequence, the model is supervised to replace masked tokens within that span with [EOS] rather than the original harmful token.
Training Objective: The method utilizes a standard masked diffusion objective with a modified target distribution:
- Harmful Set ( $D_{harm}$ ): Masked tokens in harmful responses are supervised to predict [EOS]. This teaches the model to halt harmful continuations regardless of the partial context.
- Retain Set ( $D_{retain}$ ): Masked tokens in safe responses (or safe completions of harmful prompts) are supervised to reconstruct the original tokens, preserving general utility.
Randomized Masking: The training process samples a corruption rate $\lambda$ uniformly, exposing the model to both early and late decoding stages. This ensures the alignment remains robust across the entire generation trajectory, preventing the "shallow alignment" phenomenon.
Refusal Token: The [EOS] token is chosen because it is already a standard padding and end-marker in dLLMs, ensuring the model is familiar with its usage without requiring the introduction of out-of-distribution (OOD) tokens that might degrade performance.

Real-Time Monitoring and Early Rejection

A2D inherently provides a real-time safety signal. Since the model is trained to assign high probability to [EOS] when harmful content is present, this probability mass acts as an online indicator of safety.

Dynamic Intervention: The model can begin a response but automatically terminate if unsafe continuations emerge mid-generation.
Early Rejection: By monitoring the [EOS] probability at the leftmost masked position of the first decoding step, the system can implement a threshold-based early rejection mechanism. If the probability exceeds a threshold $\tau$ , the model halts immediately without generating any tokens, significantly reducing inference latency for harmful prompts.

3. Key Contributions

Identification of Any-Order/Any-Step Vulnerabilities: The paper demonstrates that dLLMs suffer from shallow alignment that decays rapidly across decoding steps and is non-robust to arbitrary decoding orders, making them susceptible to attacks like DIJA.
Token-Level Alignment Framework: A2D introduces a novel alignment paradigm that shifts from response-level refusals to token-level suppression using [EOS], enabling robust defense at any step of the generation process.
Real-Time Safety Monitoring: The method leverages the [EOS] probability as an internal safety signal, enabling dynamic intervention and early rejection of harmful prompts.
Comprehensive Evaluation: The authors introduce the Fill-in-the-Sentence (FITS) attack to stress-test alignment under extreme conditions where a single masked sentence is embedded within a harmful response.

4. Experimental Results

The authors evaluated A2D on three instruction-tuned dLLMs: LLaDA-8B-Instruct, LLaDA-1.5, and Dream-v0-Instruct-7B, comparing it against RT, SFT, and VRPO baselines.

Defense Against Jailbreaks: A2D consistently achieved the lowest Average Attack Success Rate (ASR) across all benchmarks (Zeroshot, PAIR, ReNeLLM, Prefilling, and DIJA).
- DIJA Performance: A2D reduced DIJA success rates from over 80% in baseline models to near-zero: 1.3% on LLaDA-8B, 3.5% on LLaDA-1.5, and 0.0% on Dream.
- FITS Attack: Under the extreme FITS attack (where only one sentence is missing from a harmful response), A2D reduced ASR to nearly 0% across all models, whereas other methods failed with ASRs exceeding 60%.
Robustness to Decoding Strategies: A2D maintained low ASR across five decoding strategies (Left-to-Right, Right-to-Left, Confidence, Entropy, Random), whereas baseline models showed high vulnerability to specific orders.
Preservation of Utility: A2D avoided over-refusal, achieving 0% false positives on the XSTest benchmark (benign prompts resembling unsafe ones). It maintained competitive performance on general capability benchmarks (MMLU, GSM8K, HumanEval), outperforming other alignment methods that suffered from significant capability degradation.
Efficiency Gains: The early rejection mechanism yielded up to 19.3× faster safe termination on harmful prompts (at $\tau=0.8$ ) while maintaining high alignment fidelity.

5. Significance and Claims

The paper claims that A2D establishes a new standard for safety in diffusion-based text generation by addressing the fundamental mismatch between the flexible decoding of dLLMs and the rigid, response-level nature of traditional alignment.

Deep Alignment: A2D achieves "deep alignment," ensuring safety signals persist throughout the entire generation trajectory, not just at the start.
Order Independence: It provides a defense that is robust regardless of the decoding strategy, a critical requirement for the practical deployment of dLLMs.
Practical Deployment: By enabling real-time monitoring and early rejection, A2D offers a pathway to deploy dLLMs safely in real-world applications where latency and safety are paramount.

The authors conclude that this approach bridges the gap between the generative flexibility of dLLMs and the safety constraints required for real-world use, paving the way for reliable diffusion-based language systems.