Forget the Fact, Keep the Form: Distribution Redistribution for Retain-Free Unlearning in Masked Diffusion Language Models
Abstract
Masked Diffusion Language Models (MDLMs) have emerged as a promising alternative to autoregressive models, enabling parallel generation and flexible token ordering. Language-model unlearning aims to remove unwanted information while preserving the model's behavior on retained knowledge. Retain-free methods are increasingly desirable because retain data is often unavailable in practice. However, existing retain-free approaches face two fundamental challenges. Overly aggressive suppression can degrade generation fluency and coherence, whereas suppressing only ground-truth tokens may still allow the model to express the same information through paraphrases. We introduce **Di**stribution **R**edistribution **U**nlearning (**DiRU**), a retain-free framework that suppresses the model's reliance on forgotten information while preserving the structure and quality of its responses. **DiRU** uses the original model as a frozen teacher and the model being unlearned as a student. It masks query and response tokens associated with the forgotten information to construct an unconditional teacher. The contrast between the conditional and unconditional teachers identifies condition-dependent probability mass, including mass assigned to paraphrases of the targeted fact. The student is then distilled toward the unconditional teacher while this condition-dependent probability mass is suppressed. Through this redistribution process, **DiRU** reduces both direct recall and paraphrased re-expression of targeted information while maintaining overall generation quality. Experiments across diverse benchmarks and MDLM families show that **DiRU** outperforms existing retain-free methods, achieving strong forgetting while substantially improving model utility and preserving linguistic quality close to that of the original model.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.