Source-linked AI summary
Beyond Token Positions: Safety Alignment Across Denoising Steps in Diffusion Language Models
Guoli Wang, Haonan Shi, Tu Ouyang, An Wang
TL;DR
dLLMs raise a safety-alignment question because denoising time and response position are separate generation dimensions. The paper traces refusal distributions and commitment decisions, finds that early persistent refusal commitments shape safety, and introduces RAEC, which reduces attack success rates while largely preserving utility. The study is primarily scoped to vanilla dLLM architectures and does not examine safety dynamics during fine-tuning.
Problem
Because dLLMs generate across denoising steps and token positions, it is unclear how safety-alignment dynamics such as shallow alignment transfer from autoregressive models.
Method
The paper measures refusal behavior across denoising steps and response positions, analyzes commitment persistence, and proposes training-free Refusal-Aware Early Commitment (RAEC).
Results
Early denoising commitments strongly affect final safety, while refusal signals in unsafe trajectories often fail to persist until commitment; RAEC reduces attack success rates while largely preserving utility.
Takeaways & Limitations
Denoising step and refusal-token persistence are important dimensions for understanding and improving dLLM safety.
Takeaways & Limitations
The experiments largely use vanilla dLLM architectures, and the study does not cover safety dynamics during dLLM fine-tuning.
Abstract
from arXiv · showhide
Diffusion large language models (dLLMs) generate text through iterative denoising rather than left-to-right decoding. This generation paradigm introduces two axes that can influence safety alignment: when tokens are generated during denoising and where they appear in the response. In this paper, we measure dLLM safety behavior under harmful prompts by tracing intermediate token distributions and commitment decisions throughout denoising. Our analysis shows that refusal signals are concentrated in early denoising steps and leading response positions, and the tokens committed early can strongly shape the final safety outcome. Our measurements further show that the denoising step and persistence of refusal-token commitment are important for understanding dLLM safety. Based on these findings, we propose Refusal-Aware Early Commitment (RAEC), a simple training-free decoding method that commits persistent refusal signals from early steps. Experiments on LLaDA and Dream show that RAEC reduces attack success rates while largely preserving utility. The code is available at https://github.com/Glresearch1/RAEC.
1 Introduction
dLLMs separate denoising-time generation from response position, creating distinct safety-sensitive dimensions. The paper finds that early commitments and persistent refusal signals strongly shape safety, motivating RAEC.
- dLLMs generate tokens through iterative denoising, so generation time and response position become distinct dimensions.Unlike autoregressive decoding, early denoising can commit tokens at any response position.
- The paper identifies shallow-step alignment: early refusal commitments improve safety, while early compliance commitments can steer aligned models toward unsafe outputs.The effect is strongest at leading positions but remains visible at later positions.
- Unsafe trajectories often contain refusal evidence, but those signals are weak, transient, or overwritten before commitment.Safe trajectories show stronger and more persistent refusal mass, with refusal tokens committed more often.
- RAEC is a training-free decoding method that commits sufficiently strong, persistent early refusal signals without changing model parameters.It modifies commitment decisions rather than model weights.
- The paper’s contributions combine step-wise and position-wise measurement, shallow-step alignment analysis, refusal-persistence analysis, and RAEC.
2 Preliminaries
Masked dLLMs iteratively refine a fully masked sequence by predicting tokens across positions and committing only selected masked positions. Their decoding order is therefore independent of left-to-right response order, and safety is evaluated on harmful-instruction benchmarks.
- Masked dLLMs refine an initially masked sequence into a clean sequence across T denoising steps.At each step, the model conditions on the current sequence and predicts clean-token distributions.
- The model predicts tokens for all positions, but decoding updates typically apply only to positions that remain masked.
- At step t, Mt denotes masked positions and Ct is the subset selected for commitment during the transition from x_t to x_(t−1).
- Unlike autoregressive models, dLLMs can update multiple positions per step, allowing different response positions to be generated at different denoising times.
- Safety measurements use harmful-instruction benchmarks including JailbreakBench and StrongREJECT to assess unsafe compliance.
3 Measuring Safety Alignment Dynamics in dLLMs
Safety alignment in dLLMs depends on both denoising step and response position. Measurements show that refusal signals emerge early, early commitments shape later safety, and persistence distinguishes safe from unsafe trajectories.
- 3.1 Understanding Refusal Behavior in dLLM Generation: Refusal behavior concentrates within the first 8% of denoising steps and the first 2%–8% of response positions.The analysis tracks representative refusal words as top-1 commitment candidates across both dimensions.
- 3.2 Examining Early Steps and Leading Token Positions for Safety: Early refusal or compliance commitments substantially change safety across response positions, while leading positions amplify the effect.Refusal prefilling reduces ASR for base models, whereas compliance prefilling increases ASR for instruct models.
- 3.2 Examining Early Steps and Leading Token Positions for Safety: Shallow-step alignment describes dLLM safety being strongly influenced by early denoising commitments in addition to leading response positions.Unlike autoregressive decoding, dLLMs can commit early-step tokens at arbitrary response positions.
- 3.3 A Closer Look at Shallow-Step Alignment: The safety influence of refusal-token commitment is strongest at the earliest denoising steps and weakens as intervention moves later.After step 8, ASR gradually approaches the initial decoding level in the reported intervention experiments.
- 3.4.2 Persistence of Refusal Signals: Unsafe trajectories often contain refusal signals, but those signals are weaker and less persistent than in safe trajectories.For LLaDA-8B-Instruct, refusal signals appear in 98% of unsafe traces and 100% of safe traces, while average persistence is 7.41 versus 21.28 steps.
4 Dynamic Early Commitment of Refusal Tokens to Improve dLLM Safety
RAEC is a training-free decoding strategy that preserves sufficiently strong, persistent early refusal signals by changing commitment decisions rather than model parameters. Evaluations across utility and safety benchmarks find improved safety for LLaDA and Dream while largely preserving utility.
- RAEC changes denoising commitment decisions using model-produced logits, without modifying parameters or requiring additional training.
- RAEC operates during an early denoising window and inspects still-masked positions among the first generated answer positions.
- RAEC requires refusal evidence to exceed thresholds, avoid compliance dominance, include an explicit refusal candidate, and persist for L = 3 consecutive steps.
- When the criteria are satisfied, RAEC commits the most likely refusal token and augments the standard commitment set, while other positions follow standard decoding.
- RAEC reduces attack success rates across most harmful and jailbreak settings for LLaDA-8B-Instruct and Dream-v0-Instruct-7B while preserving utility without retraining.
5 Related Work
Related work covers diffusion language-model formulations and safety-specific attacks and defenses arising from iterative, parallel denoising. These studies motivate analyzing alignment dynamics beyond autoregressive decoding assumptions.
- Diffusion language models include continuous-latent and discrete-token formulations, with masked diffusion models iteratively reconstructing masked tokens.
- Prior safety research identifies dLLM-specific attack surfaces, including adversarial interleaved mask-text prompts and parallel generation that steers multiple positions toward unsafe outputs.
6 Conclusion
The paper finds that refusal signals concentrate early in denoising and near leading response positions, while persistent early refusal commitments can improve safety. RAEC demonstrates this principle by reducing attack success rates while largely preserving utility.
- Refusal signals concentrate in early denoising steps and leading response positions, and early commitments can strongly affect final safety outputs.
- RAEC reduces attack success rates while largely preserving utility by committing persistent early refusal signals.
7 Limitations
The experiments and conclusions primarily cover vanilla dLLM architectures, with Dream evaluated only within the setting available for LLaDA. The study does not examine safety dynamics during dLLM fine-tuning.
- The evidence is mainly limited to vanilla dLLM architectures, and Dream experiments are constrained to what is available in LLaDA.
- The study does not cover how adversarial fine-tuning shifts token distributions across denoising steps or response positions.
A.1 Compliance token set construction
The compliance-token set Vcmp is built from unsafe responses and used to measure competing compliance mass in RAEC.
- Vcmp is constructed from model responses classified as unsafe by the Llama-Guard-3-8B judge used for ASR evaluation.The responses come from four harmful-prompt datasets: StrongREJECT, DIJA-SR, JailbreakBench, and DIJA-JailbreakBench.
- Each unsafe response is tokenized with its corresponding model tokenizer to identify compliance cues.
- The resulting domain-agnostic set measures competing compliance mass in RAEC.
A.2 Refusal token set examples
The appendix lists the refusal-token sets identified for Dream-v0-Instruct-7B and LLaDA-8B-Instruct.
- The refusal-token examples are the top-10 refusal tokens identified for Dream-v0-Instruct-7B and LLaDA-8B-Instruct.The token lists are derived from the analysis in Section 3.4.1.
A.3 Experiment setting
RAEC evaluates early refusal signals within localized denoising-step and response-position windows, then tests safety, utility, evaluator robustness, and attack generalization. The reported settings improve safety while preserving general-task performance, with refusal-specific evidence outperforming substitute token sets.
- Experiment setting: RAEC uses the first 8 denoising steps and first 8 response positions as its early safety-sensitive windows.The default threshold is τ = 0.00589 for LLaDA-8B-Instruct and τ = 0.00221 for Dream-v0-Instruct-7B, with δ = 0.001 and at most 2 commitments.
- Experiment setting: RAEC computes token distributions at masked positions, tracks refusal evidence persistence, and commits a likely refusal token when the persistence condition is met.The procedure augments the standard commitment set while retaining ordinary decoding updates elsewhere.
- Evaluation: 2–7 percentage points: RAEC reduces DIJA-JBB ASR across evaluators and reduces JBB ASR to zero.The consistent improvement indicates that the result is not specific to one automatic judge or token-level refusal matching.
- Additional attacks: 12 percentage points: under adaptive PAIR, RAEC reduces ASR from 48% to 36%, while reducing the three ReNeLLM variants from 0.33% to zero.These evaluations require neither retraining nor additional target-model queries during generation.
- Ablation and sensitivity: 6% to 9% and 12%: replacing refusal tokens with random or compliance tokens increases ASR, while the 50% refusal-token baseline reaches 7%.Across parameter sweeps, ASR remains between 5% and 9%; the default window preserves 82.20% AGNews accuracy versus 82.30% for vanilla decoding.