A Theoretical Analysis of Why Masked Diffusion Models Mitigate the Reversal Curse

Moongyu Jeon; Sangwoo Shin; BumJun Kim; Kyelim Lee; Albert No

cs.AIFeb 2, 2026

A Theoretical Analysis of Why Masked Diffusion Models Mitigate the Reversal Curse

Authors: Moongyu Jeon, Sangwoo Shin, BumJun Kim, Kyelim Lee, Albert No

Organizations: Department of Artificial Intelligence, Yonsei University

Abstract

Autoregressive language models (ARMs) suffer from the reversal curse: after learning '' $A$ is $B$ ,'' they often fail on the reverse query '' $B$ is $A$ .'' Masked diffusion language models (MDMs) exhibit this failure in a much weaker form, but the underlying reason has remained unclear. A common explanation attributes this mitigation to their any-order masked training objective. However, observing '' $[\mathbf{M}]$ is $B$ '' during training teaches recovery of $A$ from $B$ in one positional configuration, and does not by itself explain why the learned evidence should transfer to the reverse prompt '' $B$ is $[\mathbf{M}]$ .'' We provide a theoretical analysis showing that this transfer arises from a parameter-level coupling between forward and reverse positional conditionals: shared Transformer parameters store token-pair evidence, while relative positional encodings route attention through queries and keys without changing the value-side evidence being retrieved. In a one-layer MDM, we prove that forward masked training strengthens evidence that is reusable in reverse queries, induces correlated forward--reverse attention routes, and yields a positively aligned shared-storage gradient component that decreases the reverse loss to first order. Controlled one-layer experiments and large-scale LLaDA/Dream experiments verify these signatures and show that they translate into improved reverse prediction.

Figures & tables

Figure 1: Masked training teaches one conditional, but the objective alone does not explain its link to the reverse conditional.

Figure 2: Storage-routing decomposition under positional reversal. For the illustrative statement “Power is Knowledge,” random masking can produce “

[\textup{{M}}]

is Knowledge,” where evidence token Knowledge contributes to the target Power logit through

\alpha^{\mathrm{fwd}}_{1,3}Y_{P,K}

. The reverse query “Knowledge is

[\textup{{M}}]

” has the corresponding contribution

\alpha^{\mathrm{rev}}_{3,1}Y_{P,K}

. Thus, both configurations share the position-invariant storage term

Y_{P,K}

, while their attention weights determine how strongly this evidence is retrieved. Under the stated assumptions, the corresponding forward and reverse attention scores are positively correlated in expectation, supporting reverse access to the same evidence. Actual retrieval also depends on competing context (see App. C.1 ).

Figure 3: One-layer validation of storage-routing signatures. Left: Pearson correlation between forward and reverse attention scores for the evidence token across total relative distance

\Delta_{1}+\Delta_{2}

. Middle: Forward–reverse gradient cosine similarity over all trainable parameters during training. Right: Target-token probabilities under forward and reverse masked queries during training.

	MDM		ARM
$L$	Fwd.	Rev.	Fwd.	Rev.
10	99.31	43.10	99.83	0.00
20	97.36	55.70	99.80	0.00
30	96.91	33.89	99.93	0.00
40	97.27	38.37	99.93	0.00

Table 1: In the one-layer setting, the MDM retains substantial reverse accuracy, whereas the ARM collapses.

$L$	Dir.	Absolute	RoPE	ALiBi
$10$	Fwd.	98.3	99.3	99.7
$10$	Rev.	11.5	43.1	99.8
$20$	Fwd.	97.8	97.4	99.5
$20$	Rev.	27.8	55.7	99.5

Table 2: Relative encodings transfer better than learned absolute embeddings in this experiment, with ALiBi giving the strongest transfer. This ordering is consistent with the storage-routing theory.

Figure 4: Large-scale validation of storage-routing signatures in LLaDA. Left: Pearson correlation between forward and reverse attention scores from the mask to the evidence entity across relative-distance placements. Middle: Forward–reverse gradient cosine similarity over trainable parameters during forward-only fine-tuning. Right: Correct target-entity probabilities under forward and reverse masked queries throughout forward-only fine-tuning on the Parent–Child dataset.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

$D$	64	128	256
Minimum bound	0.337	0.384	0.410

Appendix

Table 3: Minimum of

\operatorname{Tr}(R(2\Delta))/D

over integer

0\leq\Delta\leq 50

	$D/2=64$			$D/2=128$
$\Delta$	$c=0.10$	$c=0.20$	$c=0.30$	$c=0.10$	$c=0.20$	$c=0.30$
10	0.978	0.966	0.941	0.989	0.983	0.970
20	0.959	0.930	0.859	0.982	0.969	0.938
30	0.947	0.905	0.791	0.977	0.958	0.906
40	0.942	0.893	0.753	0.974	0.952	0.885
50	0.956	0.922	0.828	0.973	0.949	0.874

Appendix

Table 4: Cantelli lower bounds for

\Pr[\operatorname{Corr}(S_{\mathrm{fwd}},S_{\mathrm{rev}}|\mathbf{q})\geq c]

under standard RoPE frequencies.

$D$	64	128	256
Estimated minimum	0.220	0.363	0.397
$99\%$ simultaneous lower bound	0.197	0.340	0.374

Appendix

Table 5: Minimum estimated expected correlation lower bound over integer

0\leq\Delta_{1},\Delta_{2}\leq 30

, with a sign-corrected surrogate. The simultaneous lower confidence bounds are rounded downward.

$L$	$\rho_{s}\geq 0$	Median $\rho_{s}$
10	56.8%	0.0448
20	55.5%	0.0258
30	57.0%	0.0333

Appendix

Table 6: Per-plane RoPE anisotropy in nine trained one-layer RADD checkpoints.

$L$	Median $\alpha_{i,j}^{\mathrm{rev}}$ on $B$	Uniform $1/L$
10	0.589	0.100
20	0.833	0.050
30	0.567	0.033

Appendix

Table 7: Reverse attention weight on the evidence token in successful one-layer RADD checkpoints.

Hyperparameter	Value
Batch size	$256$
Learning rate	$3\times 10^{-4}$
Training steps	$3{,}000$
Gradient clipping	$1.0$
Weight decay	$0.0$
Dropout	$0.02$

Appendix

Table 8: Hyperparameters for the one-layer experiments.

Figure 5: Additional one-layer validation across sequence lengths

L=10,20,30,40

. Left column: correlation of forward and reverse attention scores as a function of total relative distance

\Delta_{1}+\Delta_{2}

. Right column: gradient alignment and probability transfer during training.

	MDM				ARM
	LLaDA 8B		Dream 7B		LLaMA-3.1 8B		Qwen-2.5 7B		Qwen-3.5 4B
Train Dataset	Fwd.	Rev.	Fwd.	Rev.	Fwd.	Rev.	Fwd.	Rev.	Fwd.	Rev.
Parent $\to$ Child (P2C)	76.7	48.3	62.0	50.0	89.9	15.9	89.9	0.5	99.0	3.0
Child $\to$ Parent (C2P)	87.7	43.7	83.0	54.0	95.9	6.9	89.0	1.4	100.0	1.0
Person $\to$ Description (P2D)	72.7	99.5	54.3	99.8	72.7	3.5	70.7	2.2	81.3	0.0
Description $\to$ Person (D2P)	99.7	41.3	100.0	27.2	83.0	1.8	80.0	1.5	100.0	0.0

Appendix

Table 9: Large-scale reversal evaluation on Parent–Child, Person–Description, and T-REx. Results are exact-match accuracy (%) averaged across three random seeds. Fwd. denotes the trained forward direction, and Rev. denotes the held-out reverse direction. MDMs retain substantially stronger reverse accuracy, whereas ARMs, including the reasoning-oriented Qwen-3.5 model, largely fail under order reversal.

Panel of Figure 6 — Figure 6: Large-scale validation of storage-routing signatures in Dream. Left: Pearson correlation between forward and reverse attention scores for the evidence entity across relative-distance placements. Right: Forward–reverse gradient cosine similarity over all trainable parameters and correct target-entity probabilities during forward-only fine-tuning.

Relative Positional Encoding

See all 16 papers

Mechanism Shift During Post-training from Autoregressive to Masked Diffusion Language Models
Jan 21, 2026Injin Kong, Hyoungjoon Lee, Yohan JoMasked Diffusion Models Autoregressive Diffusion
Graduate School of Data Science, Seoul National University · Department of Biosystems & Biomaterials Science and Engineering, Seoul National University
Understanding and Accelerating the Training of Masked Diffusion Language Models
May 13, 2026Chunsan Hong, Sanghyun Lee, Chieh-Hsin Lai +5Masked Diffusion Models Diffusion Language Model Training
1KAIST · 2Sony AI · University of Tokyo +1
The Illusion of Latent Generalization: Bi-directionality and the Reversal Curse
Mar 13, 2026Julian Coda-Forno, Jane X. Wang, Arslan ChaudhryDecoder-Only Language Models Masked Language Modeling
TUM, Helmholtz Munich · Google DeepMind

Jan 21, 2026J/K move · Enter open · S save

Mechanism Shift During Post-training from Autoregressive to Masked Diffusion Language Models

Injin Kong, Hyoungjoon Lee, Yohan Jo

Graduate School of Data Science, Seoul National University · Department of Biosystems & Biomaterials Science and Engineering, Seoul National University

Masked Diffusion Models Autoregressive Diffusion Autoregressive Language Modeling Mechanistic Interpretability Diffusion Model Fine-Tuning Language Modeling+2

Open paper AskPDF arXiv

Post-training pretrained autoregressive models (ARMs) into masked diffusion models (MDMs) provides an efficient route to diffusion language modeling, but it remains unclear whether the resulting models reuse inherited autoregressive computation or reorganize it for non-autoregressive generation. We compare two 7B ARM-MDM families across four controlled diagnostic tasks and find a task-dependent mechanism shift. On prefix-dominant tasks, MDMs largely preserve inherited high-attribution pathways or exhibit only modest changes in where computation occurs. On globally constrained tasks, the reorganization is substantially stronger, with task-relevant computation shifting toward earlier layers. This depth-wise pattern persists across prompt resampling, circuit budgets, and tested inference budgets, while targeted ablations support the functional importance of the identified structures under the tested intervention protocols. At the component level, diagnostic probes suggest that ARMs rely more strongly on sharply specialized components, whereas MDMs exhibit weaker single-component specialization and more diffuse output-space alignment. Together, these results suggest that diffusion post-training selectively preserves or reorganizes inherited computation according to task structure, rather than uniformly replacing autoregressive mechanisms.

Abstract

Relative Positional Encoding

Explore similar work

Mechanism Shift During Post-training from Autoregressive to Masked Diffusion Language Models

Understanding and Accelerating the Training of Masked Diffusion Language Models

The Illusion of Latent Generalization: Bi-directionality and the Reversal Curse