Knowledge Editing in Masked Diffusion Language Models
2026-06-02 • Computation and Language
Computation and Language
AI summaryⓘ
The authors studied how to update facts inside two types of language models: autoregressive models (ARMs) that predict text one token at a time, and masked diffusion models (MDMs) that generate text by gradually refining it. They found that the best place to make edits inside the model is similar for both types, around a specific layer linked to the last word of the fact. However, while short edits worked well in both, longer edits were harder in the MDMs because these models create text in steps that the original edits didn't account for. The authors then fixed this by adjusting the edits to work during these intermediate stages, improving the MDMs’ performance on longer updates.
knowledge editingautoregressive modelsmasked diffusion modelslocate-then-editcausal tracinglanguage modelsmulti-token generationiterative denoisingmodel weights
Authors
Haewon Park, Yohan Jo
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.