Knowledge Editing in Masked Diffusion Language Models

2026-06-02

AI Summary

This work addresses the significant performance degradation of existing knowledge editing methods in masked diffusion language models as the length of target facts increases. By adapting the locate-then-edit paradigm from autoregressive models to masked diffusion architectures, the study employs causal tracing to identify where knowledge is stored and reveals that multi-token editing failures stem from neglecting intermediate denoising states during the diffusion process. To remedy this, the authors propose an editing optimization strategy specifically tailored to these intermediate states. Experiments demonstrate that the proposed approach substantially restores multi-token knowledge editing efficacy on masked diffusion models such as LLaDA and Dream, while also providing evidence for the cross-paradigm transferability of knowledge editing locations between autoregressive and diffusion-based models.

Abstract

Knowledge editing aims to update or correct factual knowledge in a language model. A widely used approach, locate-then-edit, does this in two steps: it first localizes a fact within the model, then edits the weights there. To date, such methods have been developed exclusively on autoregressive models (ARMs). Whether their underlying assumptions hold 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 two MDMs (LLaDA, Dream) with two ARMs (LLaMA, Qwen) at matched scale. Our central finding has two parts. First, where an edit is applied transfers across paradigms: causal tracing highlights the same early-to-mid-layer MLP at the last subject token in both, and editing is most effective there. Second, this shared location does not guarantee a shared outcome. Single-token edits succeed in both, but as targets grow longer, editing degrades systematically in the MDMs but not the ARMs. The failure stems from how the edited fact is generated: producing a multi-token target requires passing through partially unmasked intermediate states for which the edit was never optimized. Guided by this diagnosis, we introduce a simple correction that optimizes the edit for these states, substantially restoring multi-token performance.

Problem

• knowledge editing
• masked diffusion models
• autoregressive models
• multi-token editing
• factual knowledge update

Innovation

• knowledge editing
• masked diffusion models
• locate-then-edit
• iterative denoising
• multi-token editing