💬 NLP

Knowledge Editing in Masked Diffusion Language Models

This paper investigates the transferability of locate-and-edit knowledge editing methods from autoregressive models to masked diffusion models, revealing that while the optimal editing location remains consistent, diffusion models suffer from degraded multi-token editing performance due to unoptimized intermediate states—a limitation effectively addressed by a proposed correction that optimizes edits for these states.

Original authors: Haewon Park, Yohan Jo

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haewon Park, Yohan Jo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Fixing a Robot's Memory

Imagine you have a very smart robot (a Large Language Model) that knows millions of facts. Sometimes, the robot gets things wrong—maybe it thinks Lionel Messi plays for FC Barcelona when he actually plays for Inter Miami.

Usually, to fix this, you have to retrain the whole robot, which is like rebuilding a house just to change a lightbulb. It's expensive and risky. A better method, called "Locate-then-Edit," is like finding the exact lightbulb that's broken and swapping it out without touching the rest of the house.

This method worked perfectly for robots that read and write like humans do: left-to-right, one word at a time (called Autoregressive Models). But a new type of robot has emerged that works differently: it reads both ways (left and right) and guesses the whole sentence at once, then slowly reveals the words like peeling back layers of an onion (called Masked Diffusion Models).

The big question this paper asks is: Does the "lightbulb swap" trick work on this new type of robot?

Part 1: The Good News (The "Where" is the Same)

The researchers first asked: Where exactly is the fact stored in the robot's brain?

The Old Robot (Left-to-Right): They found that facts are stored in a specific "filing cabinet" (a specific layer of the brain) right after the robot reads the subject (e.g., "Lionel Messi").
The New Robot (Peeling the Onion): Surprisingly, even though this robot works differently, the fact is stored in the exact same filing cabinet.

The Analogy: Imagine two different types of chefs. One cooks by adding ingredients one by one into a pot. The other mixes everything in a bowl and then slowly removes the lid to see what's cooking. Even though their cooking styles are totally different, they both store the secret recipe for "Spaghetti" in the exact same spice jar on the same shelf.

Conclusion: The "Locate-then-Edit" method works great for finding where to make the change in both types of robots.

Part 2: The Bad News (The "How" Breaks)

Here is where things get tricky. While they found the right filing cabinet, the result of the fix was different.

Short Answers (1 word): If you want to change "Messi" to "Miami" (one word), the new robot works perfectly.
Long Answers (Many words): If you want to change a fact to a long phrase (e.g., "The capital of France is Paris" -> "The capital of France is Lyon"), the new robot starts to fail. The longer the answer, the worse it gets.

Why does this happen?
The researchers realized the problem isn't where they edited, but how the robot generates the answer.

The Old Robot: It generates words one by one. When it writes "Lyon," it is in the exact same mental state that the "fix" was designed for.
The New Robot: It starts with a blank, masked screen. It reveals words step-by-step.
- Step 1: It reveals "L".
- Step 2: It sees "L" and tries to guess the next letter.
- Step 3: It sees "Ly" and tries to guess the next.

The Analogy: Imagine you are teaching a student to write a sentence.

The Old Robot is like a student writing on a blank page. You give them a hint at the start, and they write the whole sentence.
The New Robot is like a student writing on a page where you cover the words with sticky notes. You tell them, "Write 'Lyon'."
- First, you lift the note on "L". The student writes "L".
- Now, the student sees "L" and has to guess the rest.
- The Problem: The "fix" you gave the student was only designed for the moment when the page was completely covered. It didn't teach them what to do when they see just "L" or "Ly" and have to keep going. Because the robot has to pass through these "half-written" states, the fix falls apart.

Part 3: The Solution (Teaching the Robot the "In-Between" Steps)

The researchers realized the fix failed because it only practiced on the "fully covered" state. They needed to practice on the "half-covered" states too.

The Fix:
They tweaked the training process. Instead of just teaching the robot to fix the fact when the answer is completely hidden, they taught it to fix the fact at every single stage of the reveal process.

Practice fixing it when the answer is 100% hidden.
Practice fixing it when 1 letter is revealed.
Practice fixing it when 2 letters are revealed.

The Result:
By teaching the robot how to handle these "in-between" steps, the long answers suddenly worked again. The robot went from failing miserably on long phrases to succeeding almost as well as the old robot.

Summary

Location: The "memory" for facts is in the same spot for both old (left-to-right) and new (diffusion) robots.
The Glitch: The new robot fails on long answers because the "fix" was only trained for the start of the process, not the middle steps where the robot is guessing based on partial information.
The Cure: By training the fix to work at every stage of the guessing process (not just the start), the new robot can successfully update long, complex facts.

The Takeaway: When you upgrade a robot's brain, you can't just swap the part where the memory lives; you also have to make sure the new robot knows how to use that memory while it's still "thinking" and "guessing" its way through a long answer.

Technical Summary: Knowledge Editing in Masked Diffusion Language Models

Problem Statement
Knowledge editing aims to update or correct specific factual knowledge within Large Language Models (LLMs) without retraining the entire model. The dominant paradigm, "locate-then-edit" (e.g., ROME, MEMIT), was developed exclusively for Autoregressive Models (ARMs). These methods rely on causal tracing to identify that factual recall is localized to the early-to-mid-layer Multi-Layer Perceptrons (MLPs) at the last token of the subject. They then modify the weights at this location to associate a new object with the subject.

However, Masked Diffusion Models (MDMs) represent a distinct paradigm. Unlike ARMs, which predict tokens sequentially from left to right, MDMs model text bidirectionally and generate text through iterative denoising, recovering masked tokens from a fully masked state. It remains an open question whether the assumptions underlying locate-then-edit—specifically the localization of facts and the efficacy of editing at that location—hold for MDMs, given their fundamentally different training objectives and generation processes.

Methodology
The authors transfer the locate-then-edit framework (specifically the MEMIT algorithm) to MDMs without altering its core algorithm, applying only minimal adjustments to accommodate the MDM's masked prediction mechanism.

Experimental Setup: The study compares two matched pairs of models:
- ARMs: LLaMA-3-8B and Qwen2.5-7B.
- MDMs: LLaDA-8B (trained from scratch) and Dream-v0-7B (initialized from Qwen2.5 weights).
- Datasets: COUNTERFACT (standard benchmark, mostly single-token targets) and KAMEL (adapted for multi-token target analysis).
Causal Tracing: The authors perform causal tracing to determine where facts are causally localized in MDMs. They corrupt subject token embeddings and measure the recovery of correct predictions by restoring hidden states at specific layers and token positions.
Editing Evaluation: They measure editing performance (efficacy, generalization, specificity) across different token positions and layer ranges. Crucially, they evaluate performance across varying target lengths (single-token vs. multi-token) and generation settings to identify breakdowns specific to the MDM generation process.
Proposed Correction: Upon identifying a performance degradation in MDMs for multi-token targets, the authors propose a correction. The standard edit optimizes the weight update ( $\delta$ ) for a single conditioning state (the fully masked answer). The proposed correction broadens the optimization to include intermediate partially unmasked states that the model traverses during iterative denoising.

Key Findings

Shared Localization: The core assumption of locate-then-edit transfers across paradigms. Despite bidirectional training, MDMs localize factual recall to the same site as ARMs: the early-to-mid-layer MLP at the last subject token. Causal tracing confirms that the "early site" (fact recall) and "late site" (prediction readout) exist in both model types.
Divergent Outcomes (The Multi-Token Breakdown): While the location is shared, the outcome is not.
- Single-Token Targets: Editing performs comparably in ARMs and MDMs.
- Multi-Token Targets: Editing performance in MDMs degrades systematically as the target length increases. For example, on 4-token targets, LLaDA's efficacy drops from ~~0.92 (1-token) to ~0.27, whereas LLaMA remains high (~~0.89).
- Root Cause: The failure stems from a mismatch between the optimization state and the generation process. MEMIT optimizes the edit for the fully masked state. In ARMs, generation proceeds left-to-right, matching this state. In MDMs, generation proceeds through a chain of intermediate states where some tokens are unmasked and others remain masked. The edit, optimized only for the fully masked state, fails to account for the shifting context created by these intermediate unmasked tokens.
Restoration via Intermediate State Optimization: By modifying the optimization objective to fit the residual $\delta$ across all intermediate partially unmasked states (cycling through counts of revealed tokens), the authors substantially restore multi-token editing performance. For LLaDA, efficacy on 4-token targets rises from 0.27 to 0.73.

Significance and Claims
The paper claims to provide the first analysis of locate-then-edit knowledge editing in MDMs. Its primary contributions are:

Validation of Localization: Demonstrating that the causal localization of facts to the subject-token MLP is a paradigm-agnostic property, holding true for both autoregressive and diffusion-based language models.
Identification of a Generation-Induced Failure: Revealing that standard benchmarks (dominated by single-token targets) mask a systematic failure of locate-then-edit in MDMs for multi-token facts. This failure is attributed to the mismatch between single-state optimization and iterative denoising.
A Simple Correction: Proposing a method to optimize edits for intermediate states, which effectively bridges the performance gap for multi-token editing in MDMs without changing the underlying edit location or algorithm structure.

The authors conclude that porting editing methods across architectures requires matching not only the site of the edit but also the generation process through which the edited fact is realized. They note that while their correction addresses the primary cause of failure, a performance gap remains at the longest targets, suggesting future work may need to account for the specific order in which MDMs unmask tokens.