The Tokens Strike Back: Safety-Aware Denoising with Directional Embedding-based Revision

Amman Yusuf; Jiwoo Shin; Ruizhi Deng; Heitor Rodrigues Guimarães; Nikita Dvornik; Mijung Park

cs.LGOct 13, 2026

The Tokens Strike Back: Safety-Aware Denoising with Directional Embedding-based Revision

Authors: Amman Yusuf, Jiwoo Shin, Ruizhi Deng, Heitor Rodrigues Guimarães, Nikita Dvornik, Mijung Park

Organizations: Department of Computer Science, University of British Columbia, Canada · Korea Advanced Institute of Science and Technology (KAIST), South Korea · RBC Borealis, Canada

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.

Figures & tables

Method	RTP					BeaverTails
Method	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE
Unmodified	$38.6_{1.3}$	$21.1_{0.15}$	$17.5_{0.35}$	$84.65_{0.04}$	32.0	$38.6_{1.5}$	$20.6_{0.62}$	$17.5_{0.35}$	$84.65_{0.04}$	32.0
Full SAD	$32.4_{1.6}$	$29.3_{0.52}$	$30.0_{0.75}$	$83.92_{0.01}$	32.0	$27.1_{1.2}$	$35.1_{0.34}$	$47.1_{1.1}$	$83.38_{0.01}$	32.0
SADDER	$29.4_{1.7}$	$42.1_{1.4}$	$18.4_{0.26}$	$84.59_{0.06}$	32.0	$29.9_{0.9}$	$34.6_{0.54}$	$27.8_{0.86}$	$84.20_{0.03}$	32.0
SADDER + warm-up	$31.2_{0.7}$	$31.6_{2.1}$	$44.3_{0.75}$	$83.01_{0.01}$	32.0	$24.4_{0.9}$	$36.9_{0.32}$	$44.1_{0.14}$	$83.18_{0.03}$	32.0
CSD †	$22.4_{0.4}$	$225_{7.0}$	$21.6_{0.26}$	$84.59_{0.04}$	61.9	$15.6_{0.5}$	$487_{36}$	$943_{40}$	$82.61_{0.03}$	55.0

Table 1: LLaDA-Base safety, generation quality, and cost. Entries use

A_{b}

, where

A

is the mean and

b

is one sample standard deviation across three generation seeds. For percentages, SDs are in percentage points. NFE shows average cost.

Figure 1: Safety–quality tradeoff at equal NFE on Duo and UDLM. Best-of-

N

is reported in Appendix I . Larger markers denote larger guidance strength

\eta

. The dashed vertical line marks the unmodified model’s benign BERTScore, and error bars are one standard deviation across seeds.

Method	DIJA				PAD				XSTest
Method	HB ASR % $\downarrow$	LG unsafe % $\downarrow$	Harmful PPL $\downarrow$	NFE	HB ASR % $\downarrow$	LG unsafe % $\downarrow$	Harmful PPL $\downarrow$	NFE	WG refusal % $\downarrow$
Dream-Instruct
Unmodified	$82.0_{1.0}$	$80.7_{1.5}$	$11.4_{0.18}$	91.7	$51.7_{0.6}$	$64.0_{3.6}$	$8.53_{0.027}$	116.0	$55.5_{1.3}$
Full SAD	$85.3_{0.6}$	$83.0_{1.0}$	$11.8_{0.037}$	91.7	$52.7_{4.0}$	$63.3_{1.5}$	$9.52_{0.21}$	116.0	$58.0_{1.0}$
SADDER (selective)	$77.7_{1.5}$	$81.7_{3.2}$	$12.2_{0.093}$	91.7	$52.7_{0.6}$	$66.3_{2.5}$	$9.58_{0.15}$	116.0	$41.1_{0.8}$
SADDER + warm-up	$78.7_{2.3}$	$79.7_{1.2}$	$12.2_{0.10}$	91.7	$42.7_{3.1}$	$55.0_{2.0}$	$10.4_{0.17}$	116.0	$17.9_{0.7}$

Table 2: Dream-Instruct: held-out HarmBench attacks and XSTest over-refusal. Entries use

A_{b}

: mean and sample SD across three seeds (percentage points for rates). HB, LG, and PPL assess complete responses. XSTest reports full-response WG refusal on safe prompts, using valid judge labels only (coverage in Appendix H.2 ). NFE is mean cost.

Appendix figures & tables35 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
${\bm{c}}_{i},\ {\bm{x}}_{i}^{\mathrm{harm}},\ {\bm{x}}_{i}^{\mathrm{safe}}$	Reference prompt and its harmful and safe responses.
$P_{i},\ R_{i}^{u},\ M$	Prompt length, response length $\min(\|{\bm{x}}_{i}^{u}\|,M)$ , and response truncation cap; $u\in\{\mathrm{harm},\mathrm{safe}\}$ .
$N_{\mathrm{harm}},\ N_{\mathrm{safe}}$	Numbers of harmful and safe examples used to estimate the centroids.
${\bm{x}}_{t},\ {\bm{x}}_{0},\ {\bm{x}}_{0}^{l}$	Noised sequence, clean sequence, and one-hot clean token at position $l$ .
$\mathcal{V},\ L,\ \mathbb{M}$	Vocabulary, sequence length, and special mask token.
$T,\ j,\ t_{j},\ C$	Number of reverse updates, update index, corresponding noise time, and admitted intervention indices.

Appendix

Table 3: Notation for selective SADDER.

Symbol	Meaning
$\mathcal{D},\ \mathcal{D}_{\mathrm{unsafe}}$	Data distribution underlying the pretrained denoiser and unsafe reference set used by the algorithm, respectively.
$q_{t}({\bm{x}}_{t}\mid{\bm{x}}_{0})$	Forward corruption kernel, used to weight clean unsafe references given the noised state.
$\mathbb{E}_{\mathcal{D}}[{\bm{x}}_{0}^{l}\mid{\bm{x}}_{t}^{1:L}]$	Data denoiser at position $l$ , estimated by the pretrained model.
$\mathbb{E}_{\mathcal{D}_{\mathrm{unsafe}}}[{\bm{x}}_{0}^{l}\mid{\bm{x}}_{t}^{1:L}]$	Unsafe denoiser, approximated by a weighted mixture of unsafe references.
$p_{\mathcal{D}}({\bm{x}}_{0}^{1:L}\mid{\bm{x}}_{t}^{1:L})$	Data posterior over the entire clean sequence.
$\psi_{l},\ \pi_{t,l}$	Oracle unsafe-membership function $\mathcal{V}^{L}\to[0,1]$ and its posterior mean $\pi_{t,l}=\mathbb{E}_{\mathcal{D}}[\psi_{l}({\bm{x}}_{0}^{1:L})\mid{\bm{x}}_{t}^{1:L}]$ .

Appendix

Table 4: Denoiser and theoretical notation.

Parameter	Setting
Continuation sampling	32 new tokens; 32 reverse steps; temperature 1; seeds 42, 43, 44.
Jailbreak sampling	128 nominal steps; nominal length 128; temperature 1; seeds 42, 43, 44, except the explicitly labeled one-seed pilot. DIJA fills existing template spans.
Unsafe set	Continuation: up to 2,000 responses from the matching dataset after exact prompt-overlap exclusion. Jailbreak: documented 357-response WildJailbreak set.
Direction construction	Clean-reference response means. Later jobs use all valid reference rows; the selected earlier LLaDA/RTP selective row uses the 256-reference direction protocol.
Centered selection	Strict score threshold $s_{k,l}^{(t)}>0$ on masked editable positions.
Remasking	Off for the core method; top- $b$ uses $b=1$ ; aligned remasking reopens all positive-score revealed positions.

Appendix

Table 5: Shared sampling and guidance settings.

Model	Parameter	Values
LLaDA-Base	Strength $\eta$	$\{0,1,3,10,30,100,300,1000,3000\}$
	Layer $k$	$\{8,16,24,31\}$
	Window $[s,e]$	$[0,7]$ , $[0,15]$ , $[0,31]$
	Warm-up $w$	$\{8,16\}$
	BoN / FK budget	$\{2,4,8\}$
Dream-Base	Strength $\eta$	$\{0,0.1,0.3,1,3,10,100,3000,$ $10000,30000,100000,300000\}$

Appendix

Table 6: Hyperparameter values considered, consolidated by model. Values are unions across the sweeps, not a single fully crossed grid; a single value denotes a fixed setting.

k

is the scoring layer,

\eta

the SAD strength,

w

the number of full-SAD warm-up steps, and

[s,e]

the inclusive intervention window.

Model	Data	Method	$k$	$\eta$	Last step	$w$	Reopening
LLaDA	RTP	Full SAD	–	3000	7	0	none
	RTP	SADDER (selective)	31 †	100	15	0	none
	RTP	SADDER + warm-up	24	3000	15	8	none
LLaDA	BT	Full SAD	–	3000	15	0	none
	BT	SADDER (selective)	31	3000	15	0	none
	BT	SADDER + warm-up	16	3000	15	8	none

Appendix

Table 7: Selected main-table continuation settings.

Method	RTP					BeaverTails
Method	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE
Unmodified	$38.6$ $[35.7{,}41.5]$	$21.1$ $[20.2{,}22.1]$	$17.5$ $[16.9{,}18.2]$	$84.65$ $[84.50{,}84.79]$	32.0	$38.6$ $[35.6{,}41.6]$	$20.6$ $[19.6{,}21.6]$	$17.5$ $[16.9{,}18.2]$	$84.65$ $[84.50{,}84.79]$	32.0
Full SAD	$32.4$ $[29.4{,}35.5]$	$29.3$ $[27.6{,}31.1]$	$30.0$ $[27.7{,}32.9]$	$83.92$ $[83.75{,}84.09]$	32.0	$27.1$ $[24.0{,}30.4]$	$35.1$ $[33.2{,}37.0]$	$47.1$ $[42.9{,}52.7]$	$83.38$ $[83.20{,}83.55]$	32.0
SADDER (selective)	$29.4$ $[26.3{,}32.5]$	$42.1$ $[37.4{,}47.8]$	$18.4$ $[17.7{,}19.2]$	$84.59$ $[84.45{,}84.74]$	32.0	$29.9$ $[26.8{,}33.2]$	$34.6$ $[32.2{,}37.3]$	$27.8$ $[26.0{,}30.0]$	$84.20$ $[84.04{,}84.35]$	32.0
SADDER + warm-up	$31.2$ $[28.1{,}34.5]$	$31.6$ $[29.7{,}33.5]$	$44.3$ $[38.4{,}52.5]$	$83.01$ $[82.81{,}83.20]$	32.0	$24.4$ $[21.2{,}27.8]$	$36.9$ $[34.2{,}39.9]$	$44.1$ $[40.2{,}49.7]$	$83.18$ $[82.99{,}83.35]$	32.0
CSD †	$22.4$ $[20.3{,}24.7]$	$225$ $[204{,}249]$	$21.6$ $[20.6{,}22.7]$	$84.59$ $[84.44{,}84.73]$	61.9	$15.6$ $[13.8{,}17.3]$	$487$ $[440{,}541]$	$943$ $[868{,}1030]$	$82.61$ $[82.48{,}82.73]$	55.0

Appendix

Table 8: LLaDA-Base safety, generation quality, and cost.

Method	RTP					BeaverTails
Method	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE
Unmodified	$46.8$ $[43.8{,}49.7]$	$18.3$ $[17.6{,}19.1]$	$15.5$ $[15.0{,}16.0]$	$84.71$ $[84.57{,}84.86]$	32.0	$38.7$ $[35.8{,}41.6]$	$12.7$ $[12.2{,}13.2]$	$15.5$ $[15.0{,}16.0]$	$84.71$ $[84.57{,}84.86]$	32.0
Full SAD	$33.4$ $[30.1{,}36.8]$	$26.0$ $[24.3{,}27.8]$	$23.3$ $[22.0{,}24.7]$	$83.16$ $[82.96{,}83.35]$	32.0	$22.7$ $[19.7{,}25.7]$	$23.7$ $[21.8{,}25.8]$	$23.0$ $[21.7{,}24.5]$	$83.62$ $[83.44{,}83.80]$	32.0
SADDER (selective)	$35.9$ $[32.8{,}39.0]$	$29.1$ $[27.5{,}30.7]$	$18.1$ $[17.5{,}18.7]$	$84.59$ $[84.45{,}84.73]$	32.0	$34.1$ $[31.3{,}36.8]$	$19.3$ $[18.4{,}20.3]$	$20.0$ $[19.2{,}20.7]$	$84.49$ $[84.34{,}84.64]$	32.0
SADDER + warm-up	$34.7$ $[31.6{,}38.2]$	$27.1$ $[25.5{,}28.8]$	$21.3$ $[20.3{,}22.3]$	$83.76$ $[83.57{,}83.94]$	32.0	$27.7$ $[24.8{,}30.8]$	$23.6$ $[22.1{,}25.1]$	$20.5$ $[19.6{,}21.5]$	$84.21$ $[84.06{,}84.37]$	32.0
CSD †	$32.9$ $[30.6{,}35.4]$	$1180$ $[1100{,}1270]$	$3400$ $[3270{,}3540]$	$81.51$ $[81.42{,}81.61]$	60.2	$38.6$ $[35.7{,}41.4]$	$17.8$ $[16.7{,}19.1]$	$27.2$ $[25.3{,}29.4]$	$84.50$ $[84.36{,}84.64]$	55.8

Appendix

Table 9: Dream-Base safety, generation quality, and cost.

Dataset	Method	Setting
RTP	Full SAD	$31/3000/7$
	SADDER	$31/100/15$
	SADDER w/ top-1 remasking	$24/100/15$
	SADDER w/ aligned remasking	$16/100/15$
	SADDER + warm-up + aligned remasking	$24/3000/15$ ;
	CSD †	$k=31$ , $j=0\!:\!15$

Appendix

Table 10: LLaDA selected continuation settings on development set.

Dataset	Method	Setting
RTP	Full SAD	$27/10000/15$
	SADDER	$27/10000/15$
	SADDER w/ top-1 remasking	$27/100/15$
	SADDER w/ aligned remasking	$24/1e5/15$
	SADDER + warm-up	$27/3000/15$ ; none
	CSD †	$k=24$ , $j=0\!:\!15$

Appendix

Table 11: Dream selected continuation settings on development set.

Method	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE
RTP
Unmodified	$38.6$ $[35.7,41.5]$	$21.14$ $[20.21,22.13]$	$17.53$ $[16.90,18.21]$	$84.65$ $[84.50,84.79]$	32.0
Full SAD	$32.4$ $[29.4,35.5]$	$29.30$ $[27.61,31.13]$	$29.96$ $[27.65,32.89]$	$83.92$ $[83.75,84.09]$	32.0
SADDER (selective)	$29.4$ $[26.3,32.5]$	$42.08$ $[37.39,47.75]$	$18.41$ $[17.70,19.17]$	$84.59$ $[84.45,84.74]$	32.0
SADDER w/ top-1 remasking	$28.3$ $[25.2,31.3]$	$43.25$ $[38.63,48.89]$	$41.75$ $[37.34,47.13]$	$83.38$ $[83.17,83.58]$	47.0
SADDER w/ aligned remasking	$30.6$ $[27.8,33.4]$	$35.81$ $[33.47,38.38]$	$35.88$ $[33.53,38.57]$	$83.99$ $[83.82,84.16]$	46.6

Appendix

Table 12: LLaDA-Base component variants.

Method	Unsafe % $\downarrow$	Harmful PPL $\downarrow$	Benign PPL $\downarrow$	Benign BERT $\uparrow$	NFE
RTP
Unmodified	$46.8$ $[43.8,49.7]$	$18.31$ $[17.61,19.06]$	$15.45$ $[14.96,15.98]$	$84.71$ $[84.57,84.86]$	32.0
Full SAD	$33.4$ $[30.1,36.8]$	$26.02$ $[24.30,27.83]$	$23.31$ $[21.99,24.68]$	$83.16$ $[82.96,83.35]$	32.0
SADDER	$35.9$ $[32.8,39.0]$	$29.09$ $[27.54,30.68]$	$18.10$ $[17.48,18.73]$	$84.59$ $[84.45,84.73]$	32.0
SADDER w/ top-1 remasking	$37.4$ $[34.3,40.5]$	$22.54$ $[21.28,23.80]$	$16.92$ $[16.37,17.53]$	$84.57$ $[84.43,84.71]$	47.0
SADDER w/ aligned remasking	$40.9$ $[38.0,43.7]$	$23.06$ $[22.15,24.00]$	$17.53$ $[16.92,18.16]$	$84.52$ $[84.37,84.66]$	44.8

Appendix

Table 13: Dream-Base component variants.

Figure 2: Selected SADDER component configurations minus full SAD on 600 harmful prompts per dataset and three seeds. LLaDA warm-up includes its aligned remasking, Dream does not. Points and bars show unsafe-rate differences (pp) and pointwise 95% paired bootstrap CIs; negative favors SADDER.

Model/data	Method	$\Delta$ benign PPL	$\Delta$ benign BERT
LLaDA/RTP	SADDER (selective)	$-11.55\;[-14.49,-9.26]$	$0.672\;[0.527,0.813]$
	SADDER + warm-up	$3.42\;[0.49,6.19]$	$0.206\;[0.076,0.332]$
LLaDA/BT	SADDER (selective)	$-19.28\;[-24.82,-15.03]$	$0.822\;[0.681,0.967]$
	SADDER + warm-up	$-4.85\;[-8.89,-1.18]$	$0.270\;[0.133,0.407]$
Dream/RTP	SADDER (selective)	$-5.22\;[-6.55,-3.88]$	$1.429\;[1.258,1.605]$
	SADDER + warm-up	$-2.06\;[-3.25,-0.91]$	$0.596\;[0.447,0.740]$

Appendix

Table 14: SADDER variants: paired benign-quality changes versus full SAD.

RTP
SADDER variant	vs full SAD	vs random selection	vs random reopening
Selective	$-4.1\;[-6.1,-2.1]$	$-3.3\;[-5.4,-1.3]$	—
+ top-1 remasking	$-5.2\;[-8.5,-1.7]$	—	$-3.4\;[-6.3,-0.6]$
+ aligned remasking	$-2.9\;[-6.3,0.4]$	—	$-1.7\;[-3.5,0.1]$
+ warm-up + aligned	$-5.8\;[-9.0,-2.7]$	—	$-4.8\;[-6.9,-2.8]$
BeaverTails

Appendix

Table 15: LLaDA-Base: paired unsafe-rate differences (pp), means [95% CI].

RTP
SADDER variant	vs full SAD	vs random selection	vs random reopening
Selective	$2.5\;[-0.3,5.3]$	$-2.7\;[-5.5,0.1]$	—
+ top-1 remasking	$-3.9\;[-7.4,-0.4]$	$-2.6\;[-5.3,0.3]$	$-2.8\;[-5.7,0.2]$
+ aligned remasking	$-3.4\;[-6.2,-0.6]$	$0.6\;[-1.7,2.7]$	$-2.2\;[-4.5,0.2]$
+ warm-up	$0.3\;[-1.6,2.2]$	$0.4\;[-1.4,2.3]$	—
BeaverTails

Appendix

Table 16: Dream-Base: paired unsafe-rate differences (pp), means [95% CI].

RTP
Method	Unsafe %	Harmful PPL	Benign PPL	Benign BERT	NFE
BoN-2	$30.6$ $[27.8,33.4]$	$20.90$ $[20.00,21.82]$	$17.51$ $[16.89,18.15]$	$84.73$ $[84.59,84.88]$	64
BoN-4	$24.7$ $[22.3,27.2]$	$20.50$ $[19.68,21.40]$	$18.70$ $[18.00,19.41]$	$84.75$ $[84.61,84.89]$	128
BoN-8	$19.0$ $[16.7,21.5]$	$20.46$ $[19.57,21.39]$	$19.47$ $[18.74,20.21]$	$84.70$ $[84.56,84.84]$	256
FK-2	$32.4$ $[29.7,35.2]$	$21.47$ $[20.53,22.46]$	$17.59$ $[16.96,18.26]$	$84.67$ $[84.52,84.82]$	64
FK-4	$26.2$ $[23.6,28.8]$	$20.82$ $[19.92,21.71]$	$17.84$ $[17.22,18.49]$	$84.68$ $[84.54,84.83]$	128

Appendix

Table 17: LLaDA: BoN and FK at budgets 2, 4, and 8.

RTP
Method	Unsafe %	Harmful PPL	Benign PPL	Benign BERT	NFE
BoN-2	$37.0$ $[34.2,39.9]$	$17.67$ $[17.01,18.35]$	$15.39$ $[14.90,15.90]$	$84.76$ $[84.62,84.91]$	64
BoN-4	$30.6$ $[27.9,33.3]$	$17.95$ $[17.27,18.64]$	$15.78$ $[15.29,16.29]$	$84.76$ $[84.62,84.90]$	128
BoN-8	$24.6$ $[22.0,27.3]$	$17.46$ $[16.72,18.18]$	$16.87$ $[16.30,17.46]$	$84.81$ $[84.68,84.95]$	256
FK-2	$39.7$ $[36.8,42.7]$	$17.90$ $[17.23,18.56]$	$15.40$ $[14.92,15.90]$	$84.69$ $[84.54,84.83]$	64
FK-4	$31.3$ $[28.6,33.9]$	$17.90$ $[17.20,18.61]$	$15.31$ $[14.83,15.83]$	$84.72$ $[84.57,84.86]$	128

Appendix

Table 18: Dream: BoN and FK at budgets 2, 4, and 8.

Model/data	Method	Unsafe %	Unsafe-prompt PPL	NFE
Dream/BT	BoN-2	$32.5\;[29.8,35.1]$	$12.81$ $[12.27,13.38]$	64
Dream/RTP	BoN-2	$37.0\;[34.2,39.9]$	$17.67$ $[17.01,18.35]$	64
	SADDER + BoN-2	$34.0\;[30.9,37.1]$	$21.05$ $[19.98,22.11]$	64
LLaDA/BT	BoN-2	$32.1\;[29.3,34.8]$	$22.32$ $[21.38,23.33]$	64
	SADDER + BoN-2	$28.5\;[25.3,31.7]$	$27.54$ $[25.89,29.39]$	64
LLaDA/RTP	BoN-2	$30.6\;[27.8,33.4]$	$20.90$ $[20.00,21.82]$	64

Appendix

Table 19: Composition.

Method	DIJA				PAD				XSTest
Method	HB ASR % $\downarrow$	LG unsafe % $\downarrow$	Harmful PPL $\downarrow$	NFE	HB ASR % $\downarrow$	LG unsafe % $\downarrow$	Harmful PPL $\downarrow$	NFE	WG refusal % $\downarrow$
LLaDA-Instruct
Unmodified	$78.7_{3.1}$	$82.0_{2.6}$	$9.34_{0.13}$	91.7	$58.7_{2.5}$	$75.7_{2.3}$	$6.06_{0.17}$	116.0	$20.7_{1.1}$
Full SAD	$77.0_{2.0}$	$81.0_{0.0}$	$11.0_{0.13}$	91.7	$63.3_{7.5}$	$78.7_{4.0}$	$8.24_{0.21}$	116.0	$20.9_{1.9}$
SADDER (selective)	$79.0_{1.0}$	$79.3_{1.5}$	$13.3_{0.076}$	91.7	$57.7_{4.0}$	$73.7_{3.8}$	$13.2_{0.26}$	116.0	$18.8_{0.6}$
SADDER + warm-up	$78.3_{0.6}$	$80.7_{0.6}$	$12.8_{0.19}$	91.7	$57.0_{4.0}$	$68.3_{2.5}$	$11.4_{0.19}$	116.0	$16.3_{2.0}$

Appendix

Table 20: LLaDA-Instruct: held-out HarmBench attacks and XSTest over-refusal. Entries use

A_{b}

Method	Safe: over-refusal%	SD	Unsafe: refusal%	SD
LLaDA-Instruct
Unmodified	$9.7\;[6.5,13.2]$	0.9	$93.7\;[90.7,96.3]$	1.0
Full SAD	$10.0\;[6.8,13.5]$	1.1	$93.7\;[90.5,96.3]$	1.3
SADDER (selective)	$7.9\;[5.1,10.9]$	0.8	$85.2\;[80.3,89.7]$	0.3
SADDER + warm-up	$5.2\;[2.8,8.0]$	0.7	$53.5\;[46.7,60.3]$	0.0
CSD †	$10.3\;[7.1,13.7]$	0.8	$93.7\;[90.7,96.3]$	1.0

Appendix

Table 21: XSTest string-based refusal: means [95% CI] and seed SDs.

LLaDA-Instruct
Method	Safe prompts	Unsafe prompts
Full SAD	$0.3\;[-0.9,1.5]$	$0.0\;[-1.0,1.0]$
SADDER (selective)	$-1.9\;[-3.7,-0.3]$	$-8.5\;[-13.0,-4.2]$
SADDER + warm-up	$-4.5\;[-7.5,-1.9]$	$-40.2\;[-47.3,-32.7]$
CSD †	$0.5\;[0.1,1.1]$	$0.0\;[0.0,0.0]$
DiffuGuard (full)	$-5.7\;[-9.1,-2.8]$	$-65.5\;[-69.8,-61.0]$

Appendix

Table 22: Paired string-based refusal changes versus unmodified sampling (pp).

	Safe prompts			Unsafe prompts
Method	Over-refusal%	Harmful%	$n$	Refusal%	Harmful%	$n$
Unmodified	$20.7_{1.1}$ $[16.2,25.6]$	$3.3_{0.8}$ $[1.3,5.7]$	603	$97.9_{1.0}$ $[96.0,99.5]$	$1.2_{0.6}$ $[0.2,2.7]$	585
Full SAD	$20.9_{1.9}$ $[16.4,25.7]$	$3.1_{0.5}$ $[1.2,5.5]$	603	$97.9_{1.0}$ $[96.0,99.5]$	$1.2_{0.6}$ $[0.2,2.7]$	585
SADDER (selective)	$18.8_{0.6}$ $[14.5,23.4]$	$2.5_{0.8}$ $[0.9,4.3]$	644	$98.8_{0.3}$ $[97.3,100.0]$	$0.7_{0.3}$ $[0.0,1.7]$	584
SADDER + warm-up	$16.3_{2.0}$ $[12.4,20.8]$	$2.5_{0.0}$ $[1.0,4.4]$	710	$97.5_{0.5}$ $[95.2,99.3]$	$0.5_{0.5}$ $[0.0,1.4]$	592
CSD †	$29.7_{1.4}$ $[24.5,35.1]$	$3.7_{0.7}$ $[2.1,5.7]$	589	$98.1_{1.1}$ $[96.5,99.5]$	$1.0_{0.5}$ $[0.2,2.2]$	585

Appendix

Table 23: LLaDA: WildGuard judgments on XSTest.

	Safe prompts			Unsafe prompts
Method	Over-refusal%	Harmful%	$n$	Refusal%	Harmful%	$n$
Unmodified	$55.5_{1.3}$ $[49.1,61.5]$	$2.2_{0.5}$ $[0.8,3.8]$	629	$99.5_{0.0}$ $[98.5,100.0]$	$0.0_{0.0}$ $[0.0,0.0]$	600
Full SAD	$58.0_{1.0}$ $[51.6,64.4]$	$2.7_{0.3}$ $[1.1,4.7]$	619	$99.5_{0.0}$ $[98.5,100.0]$	$0.0_{0.0}$ $[0.0,0.0]$	600
SADDER (selective)	$41.1_{0.8}$ $[34.3,47.9]$	$3.2_{0.5}$ $[1.1,5.6]$	562	$99.1_{0.3}$ $[97.7,100.0]$	$0.0_{0.0}$ $[0.0,0.0]$	572
SADDER + warm-up	$17.9_{0.7}$ $[13.5,22.7]$	$1.6_{0.2}$ $[0.4,2.9]$	709	$91.5_{1.3}$ $[88.0,94.5]$	$6.9_{1.8}$ $[4.2,10.0]$	597
CSD †	$56.4_{1.9}$ $[50.0,62.4]$	$2.1_{1.0}$ $[0.8,3.7]$	630	$99.5_{0.0}$ $[98.5,100.0]$	$0.0_{0.0}$ $[0.0,0.0]$	600

Appendix

Table 24: Dream: WildGuard judgments on XSTest.

Method	Safe prompts	Pairs	Unsafe prompts	Pairs
Full SAD	$-0.0\;[-1.2,1.4]$	600	$0.0\;[0.0,0.0]$	585
SADDER (selective)	$-2.4\;[-5.8,0.9]$	579	$0.7\;[-0.9,2.4]$	581
SADDER + warm-up	$-5.9\;[-9.7,-2.1]$	593	$-0.7\;[-2.9,1.4]$	581
CSD †	$9.3\;[5.9,12.8]$	588	$0.2\;[-0.3,0.7]$	585
DiffuGuard (full)	$16.7\;[12.4,21.0]$	603	$-13.7\;[-16.8,-10.6]$	585
DiffuGuard (gen) ( $\lambda=0.35$ )	$3.2\;[1.3,5.3]$	589	$0.3\;[-0.5,1.4]$	585

Appendix

Table 25: LLaDA: paired WildGuard refusal changes (pp).

Method	Safe prompts	Pairs	Unsafe prompts	Pairs
Full SAD	$1.8\;[-0.7,4.3]$	619	$0.0\;[0.0,0.0]$	600
SADDER (selective)	$-8.6\;[-14.1,-3.3]$	532	$-0.3\;[-1.1,0.0]$	572
SADDER + warm-up	$-35.4\;[-41.6,-29.3]$	602	$-8.0\;[-11.5,-5.0]$	597
CSD †	$1.0\;[0.3,1.8]$	629	$0.0\;[0.0,0.0]$	600
DiffuGuard (full)	$7.6\;[2.7,12.7]$	629	$-23.7\;[-27.3,-20.0]$	600
DiffuGuard (gen) ( $\lambda=0.75$ )	$-1.4\;[-3.1,0.2]$	629	$0.0\;[0.0,0.0]$	600

Appendix

Table 26: Dream: paired WildGuard refusal changes (pp).

LLaDA-Instruct
Method	DIJA				PAD				XSTest
Method	HB ASR % $\downarrow$	LG unsafe % $\downarrow$	Harmful PPL $\downarrow$	NFE	HB ASR % $\downarrow$	LG unsafe % $\downarrow$	Harmful PPL $\downarrow$	NFE	Over-refusal % $\downarrow$
Unmodified	$78.7$ $[72.0{,}85.0]$	$82.0$ $[75.0{,}88.3]$	$9.34$ $[8.59{,}10.1]$	91.7	$58.7$ $[50.3{,}66.7]$	$75.7$ $[68.3{,}82.3]$	$6.06$ $[5.68{,}6.47]$	116.0	$9.7$ $[6.5{,}13.2]$
Full SAD	$77.0$ $[69.3{,}84.3]$	$81.0$ $[73.7{,}87.7]$	$11.0$ $[10.0{,}12.1]$	91.7	$63.3$ $[56.3{,}70.3]$	$78.7$ $[72.0{,}85.0]$	$8.24$ $[7.79{,}8.74]$	116.0	$10.0$ $[6.8{,}13.5]$
SADDER (selective)	$79.0$ $[71.7{,}86.3]$	$79.3$ $[71.7{,}86.3]$	$13.3$ $[11.6{,}15.7]$	91.7	$57.7$ $[50.3{,}65.0]$	$73.7$ $[66.3{,}80.7]$	$13.2$ $[10.3{,}17.7]$	116.0	$7.9$ $[5.1{,}10.9]$
SADDER + warm-up	$78.3$ $[71.3{,}85.0]$	$80.7$ $[73.0{,}87.3]$	$12.8$ $[11.1{,}15.1]$	91.7	$57.0$ $[49.0{,}65.0]$	$68.3$ $[60.3{,}76.0]$	$11.4$ $[9.97{,}13.4]$	116.0	$5.2$ $[2.8{,}8.0]$

Appendix

Table 27: Held-out HarmBench attacks and XSTest over-refusal.

	DIJA			PAD
Method	Harmful%	Refusal%	$n$	Harmful%	Refusal%	$n$
Unmodified	$96.0_{2.6}$ $[92.3,99.0]$	$1.7_{0.6}$ $[0.0,4.3]$	300	$80.0_{3.0}$ $[73.7,86.0]$	$15.3_{1.5}$ $[10.3,21.0]$	300
Full SAD	$96.3_{0.6}$ $[92.3,99.3]$	$2.0_{0.0}$ $[0.0,5.0]$	300	$85.7_{3.8}$ $[80.0,90.7]$	$11.3_{3.8}$ $[6.7,16.3]$	300
SADDER (selective)	$96.7_{0.6}$ $[92.7,99.7]$	$2.0_{0.0}$ $[0.0,5.0]$	300	$81.0_{2.6}$ $[74.3,87.0]$	$15.0_{4.0}$ $[10.0,20.7]$	300
SADDER + warm-up	$96.7_{1.5}$ $[93.0,99.3]$	$2.0_{1.0}$ $[0.0,4.7]$	300	$76.3_{0.6}$ $[68.7,83.3]$	$19.0_{1.7}$ $[12.7,26.0]$	300
CSD †	$97.0_{1.0}$ $[93.7,99.3]$	$1.7_{0.6}$ $[0.0,4.3]$	300	$77.0_{4.4}$ $[70.7,83.0]$	$19.7_{2.1}$ $[14.0,26.0]$	300

Appendix

Table 28: LLaDA: WildGuard judgments on HarmBench.

	DIJA			PAD
Method	Harmful%	Refusal%	$n$	Harmful%	Refusal%	$n$
Unmodified	$97.3_{0.6}$ $[94.0,99.7]$	$0.7_{0.6}$ $[0.0,2.0]$	300	$66.0_{2.0}$ $[57.7,74.0]$	$30.3_{2.5}$ $[22.3,38.7]$	300
Full SAD	$97.3_{0.6}$ $[94.0,100.0]$	$1.0_{0.0}$ $[0.0,3.0]$	300	$66.7_{0.6}$ $[58.3,74.7]$	$29.3_{0.6}$ $[21.3,37.7]$	300
SADDER (selective)	$96.7_{1.5}$ $[93.3,99.3]$	$1.3_{0.6}$ $[0.0,3.7]$	300	$80.0_{1.7}$ $[72.7,86.7]$	$16.7_{1.5}$ $[10.0,23.7]$	300
SADDER + warm-up	$96.3_{0.6}$ $[92.7,99.3]$	$1.7_{0.6}$ $[0.0,4.3]$	300	$69.0_{1.7}$ $[60.0,77.0]$	$28.0_{1.7}$ $[20.0,37.0]$	300
CSD †	$98.0_{0.0}$ $[95.0,100.0]$	$0.7_{0.6}$ $[0.0,2.0]$	300	$65.3_{1.5}$ $[57.0,73.3]$	$30.7_{2.5}$ $[22.7,39.0]$	300

Appendix

Table 29: Dream: WildGuard judgments on HarmBench.

LLaDA / DIJA
Comparison	$\Delta$ harmful%	$\Delta$ refusal%	Pairs
Selective vs. full SAD	$0.3\;[-0.7,1.7]$	$0.0\;[0.0,0.0]$	300
Selective vs. random	$0.3\;[-1.7,3.0]$	$0.7\;[0.0,1.7]$	300
Warm-up vs. full SAD	$0.3\;[-0.7,1.7]$	$-0.0\;[-1.0,1.0]$	300
Warm-up vs. random	$0.3\;[-1.3,2.3]$	$0.7\;[0.0,1.7]$	300
LLaDA / PAD

Appendix

Table 30: Paired WildGuard differences on HarmBench (pp).

Method	HB full	HB fills	LG full	LG fills	WG full	WG fills
Unmodified	$78.7_{3.1}$	$73.0_{3.6}$	$82.0_{2.6}$	$53.3_{2.5}$	$96.0_{2.6}$	$96.3_{1.5}$
Full SAD	$77.0_{2.0}$	$71.7_{2.9}$	$81.0_{0.0}$	$56.3_{6.8}$	$96.3_{0.6}$	$97.3_{1.2}$
SADDER (selective)	$79.0_{1.0}$	$75.7_{2.3}$	$79.3_{1.5}$	$54.0_{2.6}$	$96.7_{0.6}$	$96.3_{1.5}$
SADDER + warm-up	$78.3_{0.6}$	$71.3_{2.5}$	$80.7_{0.6}$	$41.7_{3.5}$	$96.7_{1.5}$	$97.3_{2.1}$
CSD †	$48.7_{4.7}$	$30.7_{5.1}$	$79.7_{4.9}$	$33.7_{4.2}$	$97.0_{1.0}$	$96.3_{1.5}$
DiffuGuard (full)	$42.7_{4.5}$	$32.3_{0.6}$	$81.7_{2.5}$	$49.0_{3.6}$	$97.7_{2.1}$	$97.0_{0.0}$

Appendix

Table 31: LLaDA: DIJA judge and response-view comparison on 100 held-out HarmBench behaviors, three seeds. Full includes the fixed scaffold; fills excludes it. Rates (%) are means with seed-SD subscripts. HB/LG have 300 scored outputs per row; WG has 300 except Dream aligned remasking (299).

Method	HB full	HB fills	LG full	LG fills	WG full	WG fills
Unmodified	$82.0_{1.0}$	$76.0_{3.6}$	$80.7_{1.5}$	$48.7_{2.5}$	$97.3_{0.6}$	$97.3_{0.6}$
Full SAD	$85.3_{0.6}$	$80.3_{2.3}$	$83.0_{1.0}$	$42.3_{3.1}$	$97.3_{0.6}$	$96.3_{1.5}$
SADDER (selective)	$77.7_{1.5}$	$72.3_{1.5}$	$81.7_{3.2}$	$57.3_{2.1}$	$96.7_{1.5}$	$97.3_{1.2}$
SADDER + warm-up	$78.7_{2.3}$	$75.7_{3.1}$	$79.7_{1.2}$	$39.7_{1.2}$	$96.3_{0.6}$	$96.0_{0.0}$
CSD †	$68.0_{6.2}$	$45.0_{3.5}$	$82.3_{0.6}$	$41.7_{2.1}$	$98.0_{0.0}$	$97.3_{0.6}$
DiffuGuard (full)	$33.3_{2.5}$	$22.0_{2.6}$	$84.0_{2.0}$	$44.7_{3.1}$	$98.3_{0.6}$	$98.3_{1.2}$

Appendix

Table 32: Dream: DIJA judge and response-view comparison on 100 held-out HarmBench behaviors, three seeds. Full includes the fixed scaffold; fills excludes it. Rates (%) are means with seed-SD subscripts. HB/LG have 300 scored outputs per row; WG has 300 except Dream aligned remasking (299).

LLaDA
Comparison	HB	LG	WG
Selective vs Unmodified	$2.7\;[-2.0,7.3]$	$0.7\;[-6.7,7.7]$	$0.0\;[-2.3,2.0]$
Selective vs Full SAD	$4.0\;[-1.3,9.3]$	$-2.3\;[-8.7,3.7]$	$-1.0\;[-2.7,0.3]$
Selective vs Random selection	$3.3\;[-1.3,7.7]$	$2.3\;[-2.7,7.3]$	$-1.0\;[-2.7,0.7]$
Warm-up vs Unmodified	$-1.7\;[-7.7,4.0]$	$-11.7\;[-19.7,-4.3]$	$1.0\;[-1.0,3.0]$
Warm-up vs Full SAD	$-0.3\;[-6.0,5.3]$	$-14.7\;[-21.7,-8.0]$	$0.0\;[-1.3,1.3]$

Appendix

Table 33: Generated-span DIJA score differences (percentage points), with 95% prompt-cluster bootstrap intervals. Comparisons use the same prompts and seeds; full SAD and random selection have no warm-up.

SAD
Method	Duo		UDLM		NFE
Method	Unsafe $\%\downarrow$	Benign BERT $\uparrow$	Unsafe $\%\downarrow$	Benign BERT $\uparrow$	NFE
Unmodified	$51.3_{1.8}$	$81.77_{0.08}$	$68.5_{2.5}$	$80.81_{0.04}$	128
$\eta=3{\times}10^{4}$	$43.4_{3.5}$	$81.77_{0.14}$	$57.1_{1.1}$	$80.50_{0.06}$	128
$\eta=5{\times}10^{4}$	$42.0_{1.4}$	$81.50_{0.10}$	$52.9_{1.6}$	$80.40_{0.09}$	128
$\eta=10^{5}$	$42.8_{2.2}$	$81.04_{0.07}$	$47.4_{2.9}$	$80.49_{0.03}$	128

Appendix

Table 34: Uniform-state continuation on RTP.

Figure 3: Open-gate ratio along unguided continuations of unsafe RTP prompts: the fraction of positions with

s^{(t)}_{k,l}>0

. The horizontal axis runs from early to late denoising steps.

Figure 4: Unsafe rate and benign BERTScore across guidance strengths for the strict gate and softened gates with

\tau\in\{30,70,100\}

on Duo (

k=9

) and UDLM (

k=12

). Error bars are one standard deviation across seeds: five seeds for the strict gate (the runs reported in Table 34 ) and three seeds for the softened gates.

Diffusion Model Guidance

Diffusion model guidance steers denoising toward desired conditions, rewards, constraints, or away from unwanted concepts.

See all 370 papers

Apr 28, 2026

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.

Amman Yusuf, Zhejun Jiang, Mijung Park

Department of Computer Science, The University of British Columbia, Vancouver, Canada.

Diffusion Model Guidance Diffusion Language Models

Ask PDF

May 13, 2026

Adaptive Steering and Remasking for Safe Generation in Diffusion Language Models

Diffusion Language Models(DLMs) provide a promising alternative to autoregressive language models through iterative denoising and bidirectional generation. However, their iterative generation process introduces distinct safety vulnerabilities because harmful content can emerge at arbitrary positions and persist across subsequent denoising steps. Existing defenses rely on fixed interventions or aggressive remasking, which limits adaptive control over denoising trajectories and can degrade generation quality. We propose an inference-time defense framework that combines adaptive safety steering with safety-aware remasking. Our method uses a gating direction to continuously adjust steering strength from the current denoising state and applies a steering direction to masked positions to guide subsequent predictions toward safer trajectories. Our method further employs a lightweight response detector after the first generation block to identify unsafe trajectories at an early stage. The detector triggers targeted remasking over generated content and part of the conditioning prompt, and the model regenerates the selected positions under adaptive safety steering. This design combines continuous trajectory control with explicit correction of unsafe content while requiring no modification of model parameters. Experiments on LLaDA and Dream demonstrate that our method improves robustness against diverse jailbreak attacks while preserving benign generation quality and general model capability. Our code is available at https://anonymous.4open.science/r/DLM_Steering-C32B/.

Yejin Lee, Ungsik Kim, Yo-Sub Han

Department of Computer Science, Yonsei University, Seoul 03722, Republic of Korea · Department of Computer Science and Engineering, Gyeongsang National University, Jinju 52828, Republic of Korea

Adversarial Attacks on LLMs Diffusion Language Model Inference LLM Safety Diffusion Language Models Jailbreak Defense

Ask PDF

Aug 31, 2026

Beyond Token Positions: Safety Alignment Across Denoising Steps in Diffusion Language Models

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.

Guoli Wang, Haonan Shi, Tu Ouyang +1

Case Western Reserve University 10900 Euclid Avenue, Cleveland, OH 44106, USA

Language Model Safety Evaluation Diffusion Language Model Decoding Diffusion Language Models Refusal Behavior in Language Models

Ask PDF

Abstract

Diffusion Model Guidance

Explore similar work

The Safety-Aware Denoiser for Text Diffusion Models

Adaptive Steering and Remasking for Safe Generation in Diffusion Language Models

Beyond Token Positions: Safety Alignment Across Denoising Steps in Diffusion Language Models