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

TL;DR

This paper provides a theoretical explanation for why masked diffusion language models better mitigate the reversal curse compared to autoregressive models, focusing on parameter-level coupling and shared evidence storage.

Contribution

It introduces a theoretical analysis demonstrating how shared Transformer parameters enable evidence transfer between forward and reverse queries in masked diffusion models.

Findings

Shared parameters store token-pair evidence facilitating transfer.

Forward training strengthens reversible evidence and aligns attention routes.

Experiments confirm improved reverse prediction in masked diffusion models.

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.