Adversarial Reinforcement Learning for Robust Diffusion Large Language Model Unlearning
Abstract
Diffusion language models (DLMs) have recently emerged as an alternative to autoregressive approaches, enabling parallel sequence generation and flexible token generation orders. Machine unlearning plays a critical role in mitigating legal and ethical risks by removing the influence of specific training examples from trained models. While unlearning has been extensively studied for autoregressive language models, its applicability to DLMs remains unexplored. The architectural differences of DLMs raise new challenges for effective and robust unlearning that are not addressed by existing methods. In this paper, we present the first comprehensive study of unlearning for DLMs. Through systematic empirical analysis, we show that unlearning performance in DLMs is highly sensitive to generation hyperparameters, highlighting the need for evaluation across diverse generation settings. We further observe that DLMs tend to reproduce unlearned information when target inputs are embedded within informative contexts, due to their ability to incorporate both prefix and suffix conditioning, which increases vulnerability to elicitation attacks and weakens the robustness of existing unlearning methods. To design a robust unlearning method, we propose an adversarial reinforcement learning framework, where a context generator adversarially produces informative contexts to elicit unlearned knowledge, while the DLM is optimized to suppress undesired recall. We further introduce novel components to address credit assignment and stability issues in this adversarial learning setup. Extensive experiments demonstrate that our method significantly improves unlearning effectiveness while preserving model utility.
Lay Summary
Large language models are often trained on huge collections of text, which may include private, copyrighted, harmful, or otherwise undesirable information. A growing goal in AI safety is “machine unlearning”: teaching a trained model to stop using specific information without having to retrain it from scratch. This paper studies this problem for a new family of language models called diffusion language models, which generate text differently from standard next-word prediction models and are becoming attractive because they can generate text more flexibly and efficiently. We find that unlearning is especially challenging for diffusion language models. Even after a model appears to forget certain information, it may recall that information again when given helpful surrounding context, and its behavior can change significantly under different generation settings. To address this, we propose a training approach where one model tries to create challenging contexts that could make the diffusion language model reveal the forgotten information, while the diffusion language model learns to resist these attempts. Our experiments show that this approach makes unlearning more robust against previously unseen contexts while preserving the model’s general usefulness. Overall, the paper highlights new safety challenges in diffusion language models and provides a step toward making future language models better at reliably forgetting sensitive or unwanted information.