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 panchor on the forget set by minimizing the anchor-last divergence KL(pθ∥panchor). 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 KL(panchor∥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
Figure 1: TOFU forget-set generations at τ=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.
Figure 2: TOFU forget-set generations (LLaDA-8B-Instruct) at τ=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.
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 panchor 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 qy
>0.9
>0.95
>0.99
mean qy
>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 qy over training (LLaDA on TOFU).
LLaDA-8B-Instruct
Dream-7B-Instruct
Forget ↓
Retain ↑
RA ↑
WF ↑
Forget ↓
Retain ↑
RA ↑
WF ↑
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 ↓ ; Retain, Real Authors (RA), and World Facts (WF) ↑ . Bold / underline : best / second-best.
LLaDA-8B-Instruct
Dream-7B-Instruct
Forget ↓
Retain ↑
RA ↑
WF ↑
Forget ↓
Retain ↑
RA ↑
WF ↑
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 ↓ ; Retain, Real Authors (RA), and World Facts (WF) ↑ . Bold / underline : best / second-best.
LLaDA-8B-Instruct
Dream-7B-Instruct
Forget ↓
Neighbor ↑
Utility ↑
Forget ↓
Neighbor ↑
Utility ↑
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 ↓ ; Neighbor and Utility ↑ .
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
qy
LAL
gAL(y)
gAF(y)
gAF/gAL
0.90
10.25
1.26
0.900
0.72×
0.95
10.96
0.70
0.950
1.36×
0.99
11.57
0.16
0.990
6.12×
0.999
11.73
0.019
0.999
53.6×
0.9999
11.75
0.002
0.9999
477×
Appendix
Table 5: Two-level y -coordinate ( V=126,464 , py=1/V ). gAL(y):=(∇zLAL)y and likewise for AF.
LLaDA-8B-Instruct
Dream-7B-Instruct
Forget ↓
Retain ↑
RA ↑
WF ↑
Forget ↓
Retain ↑
RA ↑
WF ↑
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 ↓ ; Retain, Real Authors (RA), and World Facts (WF) ↑ .
Figure 4: MDU’s retain-loss weight ( α∈{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: τ 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 ↓
Retain ↑
RA ↑
WF ↑
Forget ↓
Retain ↑
RA ↑
WF ↑
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 ↓ ; Retain, Real Authors (RA), and World Facts (WF) ↑ .
LLaDA-8B-Instruct
Dream-7B-Instruct
Method
Prompt
rL ↓
p↓
rL ↓
p↓
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 τ=1 checkpoint. 128 generation steps. Forget ↓ .
Forget ↓
Retain ↑
RA ↑
WF ↑
Backbone
rL
p
rL
p
rL
p
rL
p
LLaDA-8B
0.027±0.005
0.009±0.001
0.794±0.008
0.491±0.010
0.664±0.014
0.098±0.005
0.858±0.018
0.171±0.006
Dream-7B
0.000±0.000
0.001±0.000
0.857±0.018
0.723±0.015
0.603±0.008
0.193±0.007
0.760±0.031
0.185±0.001
Appendix
Table 9: EMDU seed variance on TOFU: mean ± std over 3 independent unlearning seeds, same metrics as Table 3 , one row per backbone. Forget ↓ ; Retain, Real Authors (RA), and World Facts (WF) ↑ .
Forget ↓
Neighbor ↑
Utility ↑
Backbone
F-L1
F-L2
F-L3
N-L1
N-L2
MMLU
TruQA
TriQA
LLaDA-8B
0.119±0.013
0.047±0.004
0.144±0.007
0.475±0.016
0.239±0.003
0.711±0.007
0.331±0.009
0.361±0.002
Dream-7B
0.103±0.003
0.097±0.010
0.142±0.006
0.376±0.006
0.154±0.002
0.774±0.004
0.309±0.009
0.271±0.010
Appendix
Table 10: EMDU seed variance on RWKU: mean ± std over 3 independent unlearning seeds, same metrics as Table 4 , one row per backbone.
Figure 8: RWKU/LLaDA-8B: all six evaluated baselines leak the forgotten province; EMDU abstains.
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.
Figure 10: TOFU/LLaDA-8B, Retain: EMDU names the correct book; most baselines list fabricated titles.
Figure 11: TOFU/LLaDA-8B, Real Authors: EMDU answers correctly; most baselines name a fabricated author.
Figure 12: TOFU/Dream-7B, Real Authors: EMDU answers correctly; several baselines name the wrong author or degenerate.
Figure 13: TOFU/LLaDA-8B, World Facts: EMDU answers correctly; three baselines name the wrong country.
Figure 14: RWKU/LLaDA-8B, Neighbor (Level 1), target Bruce Lee: EMDU answers correctly; all baselines fill the blank wrongly.
Figure 15: RWKU/Dream-7B, Neighbor (Level 1), target Christina Aguilera: EMDU answers correctly; three baselines give a different Pitbull single.
Machine unlearning removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.
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
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
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