Understanding the Reversal Curse Mitigation in Masked Diffusion Models through Attention and Training Dynamics

TL;DR

This work addresses the reversal curse observed in autoregressive language models, where learning a forward relation $A$ is $B$ does not guarantee the inverse $B$ is $A$. It demonstrates that masked diffusion models (MDMs) mitigate this failure via a Transformer encoder architecture with full attention and a random masking objective, enabling bidirectional conditioning. The authors show that the mitigation cannot be explained by the objective alone; they prove a structural coupling of attention under RoPE and a gradient-alignment mechanism that makes forward-loss updates also reduce reverse loss. Large-scale experiments on 7–8B parameter models corroborate the theory, showing robust reverse-inference performance for MDMs where ARMs fail, and toy-scale analyses validate the proposed mechanism. Overall, the study highlights how architecture and optimization dynamics jointly enable bidirectional reasoning in diffusion-based approaches, with implications for designing more robust, reversible language systems.

Abstract

Autoregressive language models (ARMs) suffer from the reversal curse: after learning that "$A$ is $B$", they often fail on the reverse query "$B$ is $A$". Masked diffusion-based 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 the any-order training objective. However, observing "[MASK] is $B$" during training does not necessarily teach the model to handle the reverse prompt "$B$ is [MASK]". We show that the mitigation arises from architectural structure and its interaction with training. In a one-layer Transformer encoder, weight sharing couples the two directions by making forward and reverse attention scores positively correlated. In the same setting, we further show that the corresponding gradients are aligned, so minimizing the forward loss also reduces the reverse loss. Experiments on both controlled toy tasks and large-scale diffusion language models support these mechanisms, explaining why MDMs partially overcome a failure mode that persists in strong ARMs.

Figures (14)

• Figure 1: Why the training objective of MDMs does not directly enable reverse inference. When $A$ is masked in “$A$ is $B$,” the model only learns to restore $A$ from "$[\textbf{M}]$ is $B$," i.e., $p(x\!=\!A| y\!=\!B)$. True reversal instead requires $p(y\!=\!A| x\!=\!B)$, restoring $A$ from "$B$ is $[\textbf{M}]$," which is mathematically unrelated to $p(x\!=\!A| y\!=\!B)$ under the MDM objective. Thus, training with random masking cannot by itself explain reversal capability; additional architectural factors must account for the observed success.
• Figure 2: Illustration of the evaluation setup on the Parent–Child. Each model is trained only in forward direction (e.g., parent$\to$child or child$\to$parent) and then evaluated on both forward and reverse queries. Forward queries follow the trained mapping and reverse queries require the unseen inverse mapping.
• Figure 3: Empirical validation of the attention correlation mechanism. Correlation of attention scores as a function of total relative distance $\Delta_1+\Delta_2$ in a one-layer RADD shown for sequence lengths $L=20,30,40$. The consistent positive correlation provides strong empirical support for \ref{['thm:attention_correlation']}.
• Figure 4: Probability of the ground-truth token under forward and reverse queries. In MDMs, both probabilities increase in tandem, with cosine similarity remaining near 0.999 throughout training.
• Figure 5: Gradient alignment and probability transfer in one-layer RADD. Left: Cosine similarity between forward and reverse gradients during training. Right: Probability of the ground-truth token under forward and reverse queries. For all lengths, the cosine similarity remains above 0.6, indicating that the forward and reverse optimization directions remain aligned beyond the minimal setting.

Theorems & Definitions (23)

• Theorem 4.1
• Theorem 4.2: Forward-to-reverse transfer via gradient alignment
• Theorem 2.1: Closed-form decomposition
• Corollary 2.2: Positive forward--reverse gradient alignment
• Lemma 2.3
• proof
• Lemma 2.4
• Lemma 2.5
• proof
• Lemma 2.6