Enhanced method improves text model forgetfulness without retraining
Enhanced Masked Diffusion Unlearning
Machine LearningArtificial IntelligenceComputation and Language
Summary
It's important for language models to forget certain sensitive or restricted information without needing to be completely retrained. The authors studied a previous approach that tried to erase such data but found it struggled when the model had very confident answers or when key parts of the data were still visible. They developed an improved method that addresses these problems by changing how the model measures differences from a fixed reference, helping it forget better while keeping its useful skills intact. Tests showed the new method works well on popular benchmarks and models.
What this means in practice
- •For ai deployment teams: Remove sensitive or copyrighted text from deployed language models without retraining them from scratch.
- •For content moderation engineers: Ensure stored language models no longer generate restricted or harmful content by effectively unlearning designated data.
Authors
Tomer Garber, Tom Tirer
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.