Factual recall is a lookup: some representation inside the network serves as the key, and what that key is made of decides which queries can retrieve the fact, what an edit must cover, and what a retriever can read. We call it the factual access state. In autoregressive (AR) models the subject position enriches attributes largely independently of the relation, and relation-specific selection happens later, at the readout position. In masked diffusion language models (DLMs) we find the state is assembled differently, and much earlier: in the first attention layer the wording of the relation is routed into the position of whichever entity is visible and the early feed-forward layers look the fact up on this fused representation. The state therefore varies with the relation's surface form and with the query direction, so the queries about one fact occupy a region of state space rather than a point, and an operation on the fact succeeds only for queries that enter it. We treat the state as an interface to a DLM's factual knowledge, and put it to two uses. Writing to the state gives a knowledge editor: from one declarative statement, Multi-View Edit fits one shared value per direction across that region and installs it in the early-MLP block, preserving source-prompt reliability while substantially improving paraphrase and reverse question answering, where AR editors gain near zero. Reading it suggests a complementary use: the state at each entity position can serve as a relation-specific retrieval vector with no training, returning both pieces of evidence for two-entity queries on which pooled, prompted and lexical retrievers almost never return both.
Figures & tables
Figure 1: Factual access states. A masked DLM fuses an entity with the queried relation early; the resulting state supports knowledge editing and evidence retrieval.
Figure 2: Window-tracing on a masked DLM (LLaDA-8B). Sliding-window causal tracing with w=10 , showing the indirect effect across layers at subject, relation, and other token positions. Left: full block; middle: attention only; right: FF/MLP only. The subject signal is bimodal and MLP-dominated.
Figure 3: Single-layer causal effect. Per-layer IE inside the early ( 0 – 8 ) and late ( 18 – 25 ) peaks, split by attention vs. MLP. Only layer- 0 attention stands out; no single MLP layer exceeds ∼0.03 .
Layer- 0 restriction on subject tokens
IE (mean ± sem) ↑
% of full L0 attn ↑
sink removed ( sink_off )
+0.187±0.047
79%
subject → relation only
+0.199±0.044
84%
subject → subject-internal only
+0.049±0.035
21%
subject → self only
+0.008±0.010
3%
subject → [MASK] only
−0.008±0.010
0%
Table 1: Layer- 0 attention surgery on LLaDA-8B. Each row keeps only the listed class of edges out of the subject tokens; % is normalised by the full layer- 0 attention IE of 0.238 .
Model
Method
Rel. ↑
Para. ↑
Neigh. ↑
Loc-Pres. ↑
Dream-7B
FT+L
99.67
86.47
18.95
31.33
Single-View
94.83
72.05
74.48
89.91
Multi-View
99.88
99.24
73.04
88.29
LLaDA-8B
FT+L
99.27
51.94
47.28
65.04
Single-View
96.87
53.65
57.13
68.31
Multi-View
98.14
67.92
73.73
89.92
Table 2: CounterFact. ROME-based rows use a last-subject-token write; FT+L is a constrained fine-tuning baseline. Single-View uses Vobj∪Vctx ; Multi-View adds Vrel . Published AR references are collected in App. D.2 .
Model
Method
Rel. ↑
Para. ↑
R-QA ↑
R-Judge ↑
Δ (R-QA / R-Judge) ↑
LLaDA-8B
no edit (baseline)
–
–
22.35
30.59
0.00/0.00
FT+L
74.76
45.73
19.94
28.79
−2.41 / −1.80
Single-Forward
95.53
73.03
22.52
38.88
+0.17 / +8.29
Single-Binary
97.04
65.48
42.03
41.03
+19.68 / +10.44
Multi-Binary
98.49
84.56
63.24
55.60
+40.89 / +25.01
Dream-7B
no edit (baseline)
–
–
20.26
18.15
0.00/0.00
Table 3: BAKE. R-QA and R-Judge are reverse probes; Δ is measured from each DLM’s no-edit baseline. Single/Multi denotes one view versus the masked family; Forward/Binary denotes whether the explicit reverse group is installed. FT+L is a constrained fine-tuning baseline. Published AR references are collected in App. D.2 .
Rel.
Gen.
Model
Method
=1.0↑
≥0.5↑
Weighted s↑
=1.0↑
≥0.5↑
Weighted s↑
LLaDA-8B
FT+L
15.0%
26.3%
0.206
11.9%
22.9%
0.174
Single-View
14.9%
26.2%
0.205
11.7%
23.3%
0.175
Multi-View
17.6%
32.5%
0.250
12.6%
25.5%
0.187
Dream-7B
FT+L
16.2%
32.4%
0.243
13.4%
26.8%
0.201
Single-View
20.1%
38.1%
0.291
16.5%
33.1%
0.248
Table 4: QAEdit on pre-edit-wrong examples. Fraction of answers scored exact ( =1.0 ) or at least partially correct ( ≥0.5 ), and weighted score s=(n1.0+0.5n0.5)/n . Rel. denotes source prompts and Gen. paraphrased questions. Target-length results are in App. D.6 .
Figure 4: Multi-View trigger and editing effect. Left: mean trigger strength αfire ; right: mean actual Δlogit . M1 = forward single-view, M2 = forward + reverse single-view, M3 = full Multi-View (ours). LLaDA-8B. Query difficulty increases from Paraphrase (left) to Reverse-Judge (right).
Global corpus
Entity-conditioned
Time ↓
Method
Single
Adj.
Cross.
Single
Adj.
Cross.
Global
Entity
R@1 ↑
All@2 ↑
All@2 ↑
R@1 ↑
All@2 ↑
All@2 ↑
ms/query
ms/query
Shared baselines
BM25
92.82
1.38
1.01
95.35
11.83
11.83
0.51
0.31
Qwen2.5 pooling
69.06
1.28
0.00
81.19
8.83
8.14
12.76
0.43
PromptReps (20/20), hybrid
92.53
1.58
1.56
93.47
16.35
16.26
264.98
263.20
Table 5: Controlled benchmark: global search and entity-conditioned retrieval (%). All methods search the same 5,528 cards for the same 10,800 test queries. Single is R@1; Adj./Cross. are dual-target All@2. Time is measured in ms/query. Shaded cells mark the best value in each column.
Appendix figures & tables24 assets
Supplementary material from the paper’s appendix.
Appendix
Setting
Backbone
Update mechanism
Variation explicitly covered
Reverse edit
ROME / MEMIT
AR
Backbone MLP weights
Context templates around one forward conditional
Not native
BIRD / BAKE
AR
Bidirectional objectives / evaluation
Forward–reverse objectives under AR factorisation
Explicit
Park & Jo
Masked DLM
Early–middle MLP weights
Partially unmasked target states during decoding
Not central
TimeROME-DLM
Masked DLM
Inference-time residual memory
Diffusion-time / temporal retrieval states
Not central
Multi-View Edit
Masked DLM
Early-MLP down-projections
Relation form, bidirectional context, and forward–reverse views from one statement
Explicit
Appendix
Table 6: Positioning of closely related editing settings. Existing masked-DLM editors primarily address robustness across diffusion states or introduce inference-time memories. Multi-View Edit instead expands one declarative statement into a family of conditional retrieval views.
Figure 5: Early-block decomposition (referenced in § 2.2 ). IE recovered by restoring different subsets of the early block. Layer- 0 attention combined with MLPs 1 – 8 ( 0.943 ) matches the full early block ( 0.936 ), while either sub-component alone falls short.
Figure 6: Selective peak ablation with channel decomposition ( n=550 ). Restoring the early peak alone already recovers most of the causal effect on the MLP channel ( 0.917 vs. 0.967 for all layers), while restoring the late peak alone leaves the attention channel almost silent ( 0.073 ).
Figure 7: Cross-peak conditional patching ( n=550 ). Late + Block Early preserves the late peak’s IE ( 1.012→1.012 on Full Block), whereas Early + Block Late destroys the early peak’s effect ( 0.946→0.019 ). The asymmetry is strongest on the MLP channel ( 0.917→0.148 ).
Figure 8: Subject position (Exp 3) vs. [MASK] position (Exp 7), n=194 . Left: full-block IE per layer. Middle / Right: attention-only and MLP-only splits. The two positions are almost mutually exclusive as causal loci: subject-side dominance gives way to [MASK] -side dominance only in the last few layers.
Figure 9: Dream-7B 1-D causal trace, n=432 . Left: hidden states. Middle: attention only. Right: FF/MLP only. Four curves correspond to the four role-typed positions. Dream-7B reproduces every qualitative feature reported on LLaDA-8B: a subject-anchored early peak in Hidden and MLP, a very large low-layer spike on the Attention channel at the subject positions (layers 0 – 5 , AIE≈0.25 ), a Relation-side attention peak around layers 18 – 23 , and a [MASK] -side takeover in the top layers.
Figure 10: LLaDA-8B subject-side 1-D causal trace, n=156 . Setting: the object is kept clean; the subject span is replaced by [MASK_subj] . Curves show AIE at each role-typed position across layers. The Object position now plays the causal role that the subject played in the tail-mask setting: it dominates Hidden / MLP in the early-to-mid layers and its attention spikes at the lowest layers. The [MASK_subj] position takes over in the top layers, exactly mirroring the [MASK] takeover of App. B.3 .
Category
Parameter
Value
z∗ optimisation
Optimiser
Adam
Learning rate
5×10−2
Max steps
100 (early stop at loss <0.09 )
Weight decay vwd
0.5
Clamp factor κ
3.0 ( ∥δ∥≤κ∥z0∥ )
Weight update
Cov. regularisation λ
10−5
Appendix
Table 7: Editing hyperparameters.
#
Template (prefix || suffix)
Prefix templates (prepended before prompt):
1
"" (identity, no wrapper)
2
"In one context, "
3
"According to records, "
4
"In this case, "
5
"For reference, "
Appendix
Table 8: Context templates. Each prompt is wrapped as [prefix] + prompt + [MASK] + [suffix] . Suffix templates exploit the DLM’s bidirectional attention to condition on context after the prediction slot.
CounterFact
BAKE
Model
Method
Rel. ↑
Para. ↑
Neigh. ↑
Rel. ↑
Para. ↑
R-QA ↑
R-Judge ↑
GPT-J (6B)
ROME
100.00
97.95
81.46
98.01
94.60
2.80
33.89
MEMIT
100.00
97.10
81.96
96.21
86.99
0.16
13.47
BIRD
–
–
–
99.33
96.00
2.90
34.96
LLaMA-2 (7B)
ROME
99.68
98.10
77.59
99.70
94.74
0.26
40.78
MEMIT
99.58
97.99
77.80
99.85
96.92
0.14
30.10
Appendix
Table 9: Published autoregressive references. CounterFact and BAKE results reported by the original papers for GPT-J (6B) and LLaMA-2 (7B). BAKE reports a pre-edit score of 0.00 on both reverse probes for every base LLM, so the published Δ equals the post-edit value. BIRD does not report CounterFact numbers.
Split
Position
Rel. ↑
Para. ↑
R-QA ↑
R-Judge ↑
Development ( n=4694 )
first
97.23
67.22
60.94
52.53
mean
98.27
73.63
61.02
50.79
last
98.49
84.56
63.24
55.60
spread
1.26
17.34
2.30
4.81
Held-out ( n=1000 )
first
97.50
67.60
58.60
56.60
mean
97.90
75.40
58.60
55.90
Appendix
Table 10: Subject-position sweep on LLaDA-8B (BAKE, ROME base editor). first / mean / last write the update at the first subject token, at the mean over the subject span, and at the last subject token respectively. The top block is the development pool on which the position was selected; the bottom block is a disjoint held-out split evaluated once after the choice was fixed. Reliability and both reverse probes are insensitive to the choice on both splits; Paraphrase is not. last is best on all four columns on both splits, and the three-way ordering is identical column by column. The last setting is the configuration reported throughout the main text.
Variant
#Views
Rel. ↑
Para. ↑
R-QA ↑
R-Judge ↑
Single-Binary
14
97.04
65.48
42.03
41.03
Multi-Binary (matched)
14
97.81
82.24
63.81
55.92
Multi-Binary
20
98.49
84.56
63.24
55.60
Appendix
Table 11: Matched view budget on BAKE (LLaDA-8B). #Views counts the supervised views per direction. Multi-Binary (matched) uses the relation-masked and reverse construction with six context views removed, giving it the same budget as Single-Binary . All other components are held fixed.
Variant
#Views
Rel. ↑
Para. ↑
Neigh. ↑
Loc-Pres. ↑
Single-View
14
96.87
53.65
57.13
68.31
Multi-View (matched)
14
98.03
66.85
71.05
89.64
Multi-View
20
98.14
67.92
73.73
89.92
Appendix
Table 12: Matched view budget on CounterFact (LLaDA-8B). Same three configurations as Table 11 , evaluated on the CounterFact edit pool.
Model
Variant
Rel. ↑
Para. ↑
Neigh. ↑
Loc-Pres. ↑
Dream-7B
Multi-View
99.88
99.24
73.04
88.29
Multi-View + KL
98.27
92.47
74.89
89.21
LLaDA-8B
Multi-View
98.14
67.92
73.73
89.92
Multi-View + KL
96.83
66.38
74.03
90.18
Appendix
Table 13: KL regularisation ablation on CounterFact. All rows use the ROME base editor.
Whitening statistics
Rel. ↑
Para. ↑
Neigh. ↑
Loc-Pres. ↑
No whitening
97.20
59.55
70.94
83.41
Clean-text covariance
97.20
65.32
72.85
87.96
Random-mask covariance (ours)
98.14
67.92
73.73
89.92
Appendix
Table 14: Whitening covariance ablation on CounterFact (LLaDA-8B). The model, base editor, data split, and all non-covariance hyperparameters are held fixed.
Model
Context views
Rel. ↑
Para. ↑
Neigh. ↑
Loc-Pres. ↑
LLaDA-8B
None
91.60
38.20
76.08
92.48
Prefix
95.40
49.40
75.10
91.90
Prefix + suffix
98.14
67.92
73.73
89.92
Dream-7B
None
84.72
68.20
76.44
90.50
Prefix
96.98
85.40
74.58
89.84
Prefix + suffix
99.88
99.24
73.04
88.29
Appendix
Table 15: Prefix/suffix context ablation on CounterFact. None uses no contextual templates, Prefix uses only templates before the prompt, and Prefix + suffix combines prefix and suffix templates; the last row of each block is the Multi-View configuration of Table 2 . All other components of the editor are held fixed. Bold indicates the best result within each model and metric; note that Neighborhood and Loc-Pres. are best without contextual views, i.e. the gains in Rel. and Para. are traded against a small locality cost.
Figure 11: ROME sequential-edit stability. Single-Binary and Multi-Binary are applied cumulatively without restoring the model weights. Scores are evaluated after T∈{1,10,25,50,100} sequential edits.
Table 17: Composing decoding-state views with the multi-view family on QAEdit (LLaDA-8B, pre-edit-wrong examples). The Multi-View rows repeat Table 4 . + partial reveal adds the partially revealed target states of Park and Jo (2026) to V as an additional view type; all other components of the editor are held fixed.
Target length
=1.0↑
=0.5
≥0.5↑
Weighted s↑
1 word
34.3%
18.2%
52.5%
0.434
2 words
25.9%
12.0%
38.0%
0.319
3 words
21.3%
21.3%
42.7%
0.320
4 words
21.7%
17.4%
39.1%
0.304
≥5
20.0%
20.0%
40.0%
0.300
Appendix
Table 18: QAEdit LLaDA-8B breakdown by target length for the composed editor (source prompt, Multi-View + partial reveal). Table 16 reports the same breakdown for Multi-View alone.
Domain
Relations
Entities
Film production
director, composer, screenwriter, producer
590
Person attributes
occupation, citizenship, birthplace, language
351
Sports people
sport, team, playing position, citizenship
236
Family relations
father, mother, spouse, sibling
116
Written works
author, genre, publisher, language
89
Appendix
Table 19: Domains of the controlled benchmark.
HotpotQA ( n=1,296 )
2WikiMultiHopQA ( n=2,933 )
Method
Sup. R@2 ↑
Both@2 ↑
EM ↑
F1 ↑
ms ↓
Sup. R@2 ↑
Both@2 ↑
EM ↑
F1 ↑
ms ↓
Shared baselines
BM25
52.10
25.15
52.31
61.23
0.38
56.50
26.93
61.61
62.18
0.38
Qwen2.5 pooling
18.46
3.32
45.06
53.19
31.43
14.31
1.94
38.56
43.20
31.44
PromptReps (20/20), hybrid
41.86
11.34
49.00
56.94
456.18
44.85
13.16
55.74
56.33
456.94
Dream-7B-Instruct
Appendix
Table 20: Natural comparison-question retrieval and end-to-end RAG. Every retriever receives the same per-question distractor corpus and a total budget of two sentences; all retrieved evidence goes to the same frozen Qwen2.5-7B-Instruct reader. Accuracy in percent; time in ms per query, from question to retrieved evidence, excluding answer generation. Shaded cells mark the best non-oracle method in each column.
Figure 12: Relation messages support early evidence selection. Query-side layer- 0 interventions on Dream with a fixed document index. Bars: within-entity relation R@1 averaged over targets; diamonds: global entity identification; error bars: entity/pair-clustered 95% confidence intervals. C: clean; R: relation messages removed; M: mass-matched random messages (three-seed mean); N: norm-matched random messages.
Query-side L0 intervention
Single
Adjacent
Crossed
Clean
54.31
59.11
42.53
Remove relation messages
27.07
25.26
25.35
Native random removal, mean of three
54.73
59.07
42.39
Exact-mass random message, mean of three
57.61
59.51
41.97
Norm-matched random message
54.37
58.56
41.81
Clean relation add-back
54.31
59.11
42.53
Appendix
Table 21: Causal readout results: within-entity relation R@1 (%). The first nine rows read the L0 residual. Random-seed rows are averaged only for display; all thirteen conditions are stored individually. Former/latter removal is a no-op in layouts without these tokens.
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.
Haewon Park, Yohan Jo
Graduate School of Data Science, Seoul National University
This paper shows how diffusion language models (DLMs) can be used as effective and efficient retrievers. Existing DLM-based retrievers (e.g., DiffEmbed) follow BERT-style encoding, representing each query or passage as a single mean-pooled vector. This ignores how DLMs are trained to generate responses through masked-position prediction under bidirectional attention, a capability that can provide stronger retrieval signals. We propose DiffRetriever, which uses the DLM's native masked-position prediction directly for retrieval. For each query or passage, DiffRetriever appends one or more masked positions, using the outputs as retrieval representations in a single forward pass. With one masked position, single-representation DiffRetriever already improves over DiffEmbed on the same backbones. DiffRetriever also naturally extends to multi-representation retrieval: DLMs process multiple masked positions jointly, enabling ColBERT-style fine-grained matching with little additional encoding latency. In autoregressive LLM retrievers, the same multi-representation strategy requires sequential decoding and therefore incurs much higher latency. DiffRetriever obtains the strongest aggregate effectiveness within our matched comparison, outperforming DiffEmbed, PromptReps, and RepLLaMA. Masked-position counts selected on training data transfer well across datasets, while per-query variation suggests headroom for adaptive allocation. Code is available at https://github.com/ielab/diffretriever.
Shuai Wang, Yu Yin, Shengyao Zhuang +2
The University of Queensland, Brisbane, QLD, Australia · CSIRO
Discrete diffusion language models generate text by iteratively denoising an entire response in parallel. At each step, they predict tentative tokens for every masked position, committing the confident predictions to the output and discarding the unconfident ones. We show that the discarded tokens are in fact a useful lookahead signal for retrieval-augmented generation: even low-confidence tokens often surface salient entities early in the denoising trajectory, enabling retrieval of stronger evidence before the output is finalized. We exploit this through Self-Augmenting Retrieval for Diffusion Language Models (SARDI), a dynamic RAG framework that uses these lookahead tokens to guide retrieval during denoising. SARDI is training-free, retriever-agnostic, and applicable to any reasoning-capable discrete diffusion language model. Across five multi-hop QA benchmarks, SARDI outperforms current training-free diffusion and autoregressive retrieval baselines at up to 8× higher throughput.
Paul Jünger, Justin Lovelace, Linxi Zhao +2
Department of Computer Science, Cornell University.