Knowledge Editing for Masked Diffusion Language Models
TL;DR
A simple correction is introduced that optimizes the edit including such states in which the target is partially unmasked, substantially restoring multi-token performance in masked diffusion models.
Abstract
Knowledge editing aims to update or correct factual knowledge in a language model. A widely used approach, locate-then-edit, first localizes a fact within the model and then edits the weights there. To date, such methods have been developed exclusively for autoregressive models (ARMs). Whether they work for masked diffusion models (MDMs), which model text bidirectionally and generate by iterative denoising rather than next-token prediction, remains an open question. We address it by transferring locate-then-edit to MDMs and comparing multiple MDMs with their matched ARMs. Our central finding has two parts. First, where an edit should be applied transfers between them: the same early-to-mid-layer MLP at the last subject token is most effective for both. Second, this shared location does not guarantee a shared outcome. Single-token edits succeed in both, but as targets grow longer, editing degrades far more sharply in the MDMs than in the ARMs. The failure stems from how the edited fact is generated: producing a multi-token target passes through intermediate states in which the target is partially unmasked, for which the edit was never optimized. Guided by this diagnosis, we introduce a simple correction that optimizes the edit including such states, substantially restoring multi-token performance. Our code is available at https://github.com/holi-lab/MDM-KE.
