Machine Unlearning for Masked Diffusion Language Models

TL;DR

This paper introduces Masked Diffusion Unlearning (MDU), a novel framework for unlearning specific knowledge in masked diffusion language models, demonstrating superior performance on benchmarks.

Contribution

The paper presents the first unlearning method tailored for MDLMs, utilizing diffusion process insights to effectively unlearn knowledge while balancing privacy and utility.

Findings

MDU achieves high unlearning performance compared to existing methods.

Empirical results on benchmarks validate MDU's effectiveness.

Code is publicly available at the provided GitHub link.

Abstract

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.