Enhanced Masked Diffusion Unlearning

Tomer Garber; Tom Tirer

cs.LGOct 18, 2026

Enhanced Masked Diffusion Unlearning

Authors: Tomer Garber, Tom Tirer

Organizations: Faculty of Engineering, Bar-Ilan University, Ramat Gan, Israel · Red Hat

Abstract

Masked diffusion language models (MDLMs) generate text by iteratively unmasking tokens with bidirectional context. As these models reach real-world deployment, machine unlearning becomes essential. This process aims to remove the influence of a designated "forget set" without retraining from scratch, so that sensitive, copyrighted, or restricted content can no longer be elicited while preserving retained capabilities. Masked Diffusion Unlearning (MDU) is the first publicly available unlearning method customized for these models: it tunes the model $p_θ$ to match a frozen prompt-masked anchor $p_{\mathrm{anchor}}$ on the forget set by minimizing the anchor-last divergence $\mathrm{KL}(p_θ\Vert p_{\mathrm{anchor}})$ . We identify two failure modes of MDU. First, random answer masking can leave key forget tokens visible to the anchor even when the question is fully masked. Second, at a memorized (near-one-hot) prediction, the anchor-last logit gradient vanishes while the anchor-first divergence $\mathrm{KL}(p_{\mathrm{anchor}} \Vert p_θ)$ logit gradient does not. As a result, MDU training stalls on the tokens we most need to forget. Backed by theoretical and empirical analyses, we propose Enhanced Masked Diffusion Unlearning (EMDU): anchor-first KL divergence toward a frozen anchor whose entire answer span is masked. On TOFU and RWKU benchmarks using LLaDA-8B and Dream-7B, our EMDU matches or exceeds the forgetting of MDU and adapted AR baselines, while better preserving retain utility.

Figures & tables

Panel of Figure 1 — Figure 1: TOFU forget-set generations at $\tau{=}0$ with unlearning following the MDU recipe. Highlighted spans mark leaked facts in the prediction and the matching tokens in GT; the ellipsis (…) denotes elided repeated tokens in these incoherent generations.

Panel of Figure 2 — Figure 2: TOFU forget-set generations (LLaDA-8B-Instruct) at $\tau{=}1$ with unlearning following the MDU recipe: fluent and confident, yet either completely wrong (top) or a full leak of the correct answer (bottom), both indistinguishable from a real answer to a downstream user.

Panel of Figure 3 — Figure 3: Even with the prompt fully masked, MDU’s random answer masking leaves forget tokens as context for the prompt-masked anchor. Those tokens are not trained on, but they bias $p_{\mathrm{anchor}}$ at the remaining masked positions (where the KL is applied). FAM sets $t{=}1$ on the answer span of the anchor input. Highlighted spans are surviving context tokens on a TOFU forget pair.

	Anchor-last				Anchor-first
Epoch	mean $q_{y}$	${>}0.9$	${>}0.95$	${>}0.99$	mean $q_{y}$	${>}0.9$	${>}0.95$	${>}0.99$
$0$ (SFT)	$0.960$	$92.2\%$	$88.6\%$	$77.2\%$	$0.960$	$92.2\%$	$88.6\%$	$77.2\%$
$1$	$0.951$	$90.2\%$	$86.4\%$	$75.0\%$	$0.714$	$50.1\%$	$41.9\%$	$26.1\%$
$2$	$0.820$	$70.8\%$	$66.5\%$	$55.8\%$	$0.442$	$15.9\%$	$8.9\%$	$2.4\%$
$3$	$0.705$	$55.9\%$	$50.8\%$	$40.0\%$	$0.412$	$11.2\%$	$6.0\%$	$1.4\%$
$4$	$0.633$	$46.3\%$	$40.6\%$	$29.1\%$	$0.399$	$10.0\%$	$5.5\%$	$1.2\%$

Table 1: Forget-token mass

q_{y}

over training (LLaDA on TOFU).

	LLaDA-8B-Instruct								Dream-7B-Instruct
	Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$		Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$
Method	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$
Base	0.822	0.337	0.808	0.301	0.647	0.044	0.840	0.133	0.931	0.837	0.944	0.872	0.651	0.137	0.853	0.134
MDU (AL w/o FAM)	0.113	0.153	0.467	0.360	0.701	0.102	0.858	0.185	0.258	0.326	0.394	0.452	0.413	0.120	0.673	0.118
AF w/o FAM	0.031	0.074	0.644	0.460	0.747	0.126	0.838	0.158	0.000	0.107	0.512	0.441	0.365	0.215	0.464	0.249
AL with FAM	0.032	0.016	0.409	0.356	0.571	0.081	0.829	0.131	0.000	0.003	0.636	0.408	0.148	0.040	0.267	0.047

Table 2: Ablation on TOFU with LLaDA-8B and Dream-7B. Forget

\downarrow

; Retain, Real Authors (RA), and World Facts (WF)

\uparrow

. Bold / underline : best / second-best.

	LLaDA-8B-Instruct								Dream-7B-Instruct
	Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$		Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$
Method	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$
Base	0.822	0.337	0.808	0.301	0.647	0.044	0.840	0.133	0.931	0.837	0.944	0.872	0.651	0.137	0.853	0.134
GA	0.479	0.029	0.501	0.033	0.492	0.044	0.832	0.128	0.351	0.012	0.356	0.017	0.500	0.035	0.666	0.101
GD	0.407	0.012	0.510	0.024	0.499	0.085	0.849	0.171	0.717	0.320	0.838	0.548	0.518	0.053	0.805	0.077
NPO	0.231	0.016	0.622	0.174	0.689	0.040	0.875	0.173	0.039	0.003	0.797	0.585	0.590	0.124	0.832	0.123

Table 3: TOFU results on LLaDA-8B and Dream-7B. Forget

\downarrow

; Retain, Real Authors (RA), and World Facts (WF)

\uparrow

. Bold / underline : best / second-best.

	LLaDA-8B-Instruct								Dream-7B-Instruct
	Forget $\downarrow$			Neighbor $\uparrow$		Utility $\uparrow$			Forget $\downarrow$			Neighbor $\uparrow$		Utility $\uparrow$
Method	F-L1	F-L2	F-L3	N-L1	N-L2	MMLU	TruQA	TriQA	F-L1	F-L2	F-L3	N-L1	N-L2	MMLU	TruQA	TriQA
Base	0.484	0.393	0.475	0.535	0.344	0.680	0.400	0.376	0.425	0.434	0.460	0.502	0.361	0.760	0.326	0.366
GA	0.135	0.149	0.224	0.456	0.225	0.667	0.334	0.374	0.136	0.126	0.243	0.377	0.268	0.711	0.298	0.365
NPO	0.160	0.160	0.270	0.463	0.241	0.665	0.328	0.372	0.136	0.213	0.237	0.341	0.310	0.690	0.300	0.416
SimNPO	0.135	0.170	0.251	0.470	0.237	0.669	0.320	0.378	0.170	0.209	0.248	0.355	0.300	0.689	0.286	0.412

Table 4: RWKU results on LLaDA-8B and Dream-7B (10 subjects). Forget / Neighbor: mean of per-subject means (macro over all 10 subjects). Forget

\downarrow

; Neighbor and Utility

\uparrow

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

$q_{y}$	$\mathcal{L}_{\mathrm{AL}}$	$g_{\mathrm{AL}}(y)$	$g_{\mathrm{AF}}(y)$	$g_{\mathrm{AF}}/g_{\mathrm{AL}}$
$0.90$	$10.25$	$1.26$	$0.900$	$0.72\times$
$0.95$	$10.96$	$0.70$	$0.950$	$1.36\times$
$0.99$	$11.57$	$0.16$	$0.990$	$6.12\times$
$0.999$	$11.73$	$0.019$	$0.999$	$53.6\times$
$0.9999$	$11.75$	$0.002$	$0.9999$	$477\times$

Appendix

Table 5: Two-level

y

-coordinate (

V{=}126{,}464

p_{y}{=}1/V

g_{\mathrm{AL}}(y):=(\nabla_{z}\mathcal{L}_{\mathrm{AL}})_{y}

and likewise for AF.

	LLaDA-8B-Instruct								Dream-7B-Instruct
	Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$		Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$
Method	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$	rL	$p$
MDU (official code, aligned dataset)	0.346	0.223	0.586	0.391	0.668	0.066	0.815	0.115	0.429	0.401	0.550	0.519	0.499	0.090	0.679	0.086
MDU (chosen config, Table 3 )	0.113	0.153	0.467	0.360	0.701	0.102	0.858	0.185	0.258	0.326	0.394	0.452	0.413	0.120	0.673	0.118

Appendix

Table 6: TOFU results from running Lee et al. (2026) ’s official code, after aligning the dataset, under their suggested configuration, versus the config we report as “MDU” in Table 3 . Forget

\downarrow

; Retain, Real Authors (RA), and World Facts (WF)

\uparrow

Figure 4: MDU’s retain-loss weight ( $\alpha\in\{1,1.1,1.2,1.3,1.4,1.5,2,5\}$ ) sweep on TOFU, both backbones: Forget vs. Retain, in ROUGE-L (left) and probability (right). For EMDU, the configs that are reported in Table 3 are marked with a star.

Figure 5: TOFU Forget vs. Retain on LLaDA-8B, all eight methods, ROUGE-L (left) and probability (right). For each method, the config that is reported in Table 3 is marked with a star.

Figure 6: TOFU Forget vs. Retain on Dream-7B, all eight methods, ROUGE-L (left) and probability (right). For each method, the config that is reported in Table 3 is marked with a star.

Figure 7: $\tau$ sweep on TOFU, both backbones and both methods: Forget vs. Retain, ROUGE-L (left) and probability (right).

$L$	$S$	Method	LLaDA-8B-Instruct				Dream-7B-Instruct
			Forget $\downarrow$	Retain $\uparrow$	RA $\uparrow$	WF $\uparrow$	Forget $\downarrow$	Retain $\uparrow$	RA $\uparrow$	WF $\uparrow$
$32$	$32$	SFT	0.776	0.819	0.591	0.847	0.821	0.869	0.696	0.813
		MDU	0.171	0.554	0.653	0.834	0.210	0.344	0.390	0.648
		EMDU	0.037	0.795	0.623	0.866	0.000	0.786	0.583	0.799
$64$	$32$	SFT	0.882	0.882	0.635	0.849	0.927	0.933	0.637	0.845
		MDU	0.166	0.510	0.647	0.879	0.266	0.393	0.222	0.516

Appendix

Table 7: TOFU under varying generation output length (

L

) and denoising steps (

S

) (Rouge-L recall). Forget

\downarrow

; Retain, Real Authors (RA), and World Facts (WF)

\uparrow

		LLaDA-8B-Instruct		Dream-7B-Instruct
Method	Prompt	rL $\downarrow$	$p$ $\downarrow$	rL $\downarrow$	$p$ $\downarrow$
SFT	original	0.822	0.337	0.931	0.837
	paraphrased	0.439	0.037	0.607	0.358
EMDU	original	0.030	0.008	0.000	0.001
	paraphrased	0.090	0.016	0.008	0.004

Appendix

Table 8: TOFU forget10 under original vs. paraphrased questions (Rouge-L recall and Eq. ( 13 ) answer probability). Paraphrased prompts use the TOFU paraphrased_question field from forget10_perturbed ; gold answers remain the original forget answers. SFT is the Base checkpoint from Table 3 ; EMDU is the same

\tau{=}1

checkpoint.

128

generation steps. Forget

\downarrow

	Forget $\downarrow$		Retain $\uparrow$		RA $\uparrow$		WF $\uparrow$
Backbone	rL	$p$	rL	$p$	rL	$p$	rL	$p$
LLaDA-8B	$0.027\pm 0.005$	$0.009\pm 0.001$	$0.794\pm 0.008$	$0.491\pm 0.010$	$0.664\pm 0.014$	$0.098\pm 0.005$	$0.858\pm 0.018$	$0.171\pm 0.006$
Dream-7B	$0.000\pm 0.000$	$0.001\pm 0.000$	$0.857\pm 0.018$	$0.723\pm 0.015$	$0.603\pm 0.008$	$0.193\pm 0.007$	$0.760\pm 0.031$	$0.185\pm 0.001$

Appendix

Table 9: EMDU seed variance on TOFU: mean

{\pm}

std over

3

independent unlearning seeds, same metrics as Table 3 , one row per backbone. Forget

\downarrow

; Retain, Real Authors (RA), and World Facts (WF)

\uparrow

	Forget $\downarrow$			Neighbor $\uparrow$		Utility $\uparrow$
Backbone	F-L1	F-L2	F-L3	N-L1	N-L2	MMLU	TruQA	TriQA
LLaDA-8B	$0.119\pm 0.013$	$0.047\pm 0.004$	$0.144\pm 0.007$	$0.475\pm 0.016$	$0.239\pm 0.003$	$0.711\pm 0.007$	$0.331\pm 0.009$	$0.361\pm 0.002$
Dream-7B	$0.103\pm 0.003$	$0.097\pm 0.010$	$0.142\pm 0.006$	$0.376\pm 0.006$	$0.154\pm 0.002$	$0.774\pm 0.004$	$0.309\pm 0.009$	$0.271\pm 0.010$

Appendix

Table 10: EMDU seed variance on RWKU: mean

{\pm}

std over

3

independent unlearning seeds, same metrics as Table 4 , one row per backbone.

Panel of Figure 8 — Figure 8: RWKU/LLaDA-8B: all six evaluated baselines leak the forgotten province; EMDU abstains.

Panel of Figure 9 — Figure 9: TOFU/Dream-7B: six of the seven evaluated baselines leak at least part of the forgotten name or birthplace (NPO leaks part of the name); EMDU abstains.

Panel of Figure 10 — Figure 10: TOFU/LLaDA-8B, Retain: EMDU names the correct book; most baselines list fabricated titles.

Panel of Figure 11 — Figure 11: TOFU/LLaDA-8B, Real Authors: EMDU answers correctly; most baselines name a fabricated author.

Panel of Figure 12 — Figure 12: TOFU/Dream-7B, Real Authors: EMDU answers correctly; several baselines name the wrong author or degenerate.

Panel of Figure 13 — Figure 13: TOFU/LLaDA-8B, World Facts: EMDU answers correctly; three baselines name the wrong country.

Panel of Figure 14 — Figure 14: RWKU/LLaDA-8B, Neighbor (Level 1), target Bruce Lee: EMDU answers correctly; all baselines fill the blank wrongly.

Panel of Figure 15 — Figure 15: RWKU/Dream-7B, Neighbor (Level 1), target Christina Aguilera: EMDU answers correctly; three baselines give a different Pitbull single.

Machine Unlearning

Machine unlearning removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.

See all 349 papers

May 18, 2026

Machine Unlearning for Masked Diffusion Language Models

Recent masked diffusion language models (MDLMs), such as LLaDA and Dream, have achieved performance comparable to autoregressive large language models. Unlike autoregressive models, which generate text sequentially, MDLMs generate text by iteratively denoising masked positions in parallel. During fine-tuning, MDLMs learn to recover responses from masked response states conditioned on a prompt, thereby shifting their predictions from a prompt-masked unconditional distribution toward a prompt-conditional distribution. Despite this distinct generative and fine-tuning mechanism, machine unlearning for MDLMs remains largely unexplored. In this paper, we propose Masked Diffusion Unlearning (MDU), the first unlearning framework for MDLMs, by revisiting the process of learning specific knowledge in terms of diffusion. Specifically, MDU minimizes a forward KL divergence from the prompt-conditional prediction to a prompt-masked unconditional anchor at every masked response position, with a temperature scaling parameter to control the privacy-utility trade-off. Our empirical results on standard benchmarks and MDLM backbones show that MDU achieves high unlearning performance compared to existing LLM unlearning methods. Code is available at https://github.com/leegeoru/MDU.

Georu Lee, Seungwon Jeong, Hoki Kim +2

1Dongguk University-Seoul · 2Chung-Ang University · 3Korea Institute for Advanced Study

Masked Diffusion Models Diffusion Model Unlearning Diffusion Language Models Machine Unlearning LLM Unlearning

Ask PDF

Jul 16, 2026

Mask-Aware Policy Gradients for Diffusion Language Models

Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.

Haran Raajesh, Kulin Shah, Adam Klivans +1

Department of Computer Science The University of Texas at Austin

Masked Diffusion Models RL for Language Model Reasoning Policy Gradient Methods RL for Diffusion Models

Ask PDF

Apr 20, 2026

Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Large Language Models

Diffusion large language models (dLLMs) gain speed by committing multiple tokens in parallel at each denoising step, but any erroneous commitment persists as conditioning context and biases every subsequent prediction. LLaDA2.1 repairs such errors with Token-to-Token (T2T) editing, which re-examines previously unmasked tokens and overwrites them when an alternative becomes sufficiently confident. We argue that this replacement action is itself the limiting factor: under polluted context, a confident replacement can propagate the error, while under a multimodal posterior no alternative may be confident enough to trigger an edit. We propose Token-to-Mask (T2M) remasking, a training-free rule that revokes suspicious commitments by resetting them to [M] and lets the subsequent mask-filling steps re-predict them from a cleaner context. T2M improves accuracy by +13.33 points on AIME 2025 and +8.56 points on CMATH. These results suggest that, for parallel discrete generators, remasking suspect tokens rather than overwriting them is a more reliable self-correction primitive.

Lin Yao

School of Computer Science, Shanghai Jiao Tong University, Shanghai, 200240, China · Zhongguancun Academy, Beijing, 100097, China

LLM Self-Correction Multi-Token Prediction LLM Reliability Diffusion Language Model Decoding Diffusion Language Models

Ask PDF

Abstract

Machine Unlearning

Explore similar work

Machine Unlearning for Masked Diffusion Language Models

Mask-Aware Policy Gradients for Diffusion Language Models

Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Large Language Models