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

Problem Statement

Autoregressive language models (ARMs) suffer from the reversal curse: after learning a relational fact in one order (e.g., "A is B"), they frequently fail to answer the corresponding order-reversed query (e.g., "B is A"). While recent Masked Diffusion Models (MDMs) exhibit a significantly weaker reversal curse, the underlying mechanism has remained unclear.

A common explanation attributes this mitigation to the any-order masked training objective of MDMs, which allows the model to see masked examples in both forward and reverse directions. However, the authors argue this explanation is incomplete. As illustrated in Figure 1 of the paper, observing a forward masked example like "[M] is B" (where the target is A) provides supervision for the conditional $p_\theta(A|B)$ in a specific positional configuration ($i < j$). This does not inherently explain why the learned evidence should transfer to the reverse query "B is [M]" (where $j' < i'$), which corresponds to a distinct conditional $p_\theta(A|B)$ in a different positional configuration.

Methodology

The paper provides a theoretical analysis of a simplified one-layer Masked Diffusion Model to isolate the mechanism enabling reverse transfer. The core of the methodology is a Storage-Routing Decomposition:

1. Storage-Routing Decomposition: The authors decompose the contribution of an input token $B$ to the logit of a target token $A$ into two components:

   • Shared Storage ($Y_{A,B}$): A position-invariant term representing the evidence stored in the model parameters. Under relative positional encodings (like RoPE or ALiBi), the value vector of token $B$ is independent of its position, meaning $Y_{A,B}$ is shared regardless of whether $B$ appears before or after the mask.
   • Attention Routing ($\alpha_{i,j}$): A position-dependent attention weight that determines how strongly the masked position attends to token $B$.

2. Theoretical Analysis:

   • Storage Dynamics: The authors prove that forward masked training (recovering A from B) increases the shared storage term $Y_{A,B}$, making the evidence reusable for reverse queries.
   • Routing Correlation: They analyze whether the reverse query can still access this stored evidence. Using RoPE (Rotary Position Embedding) and ALiBi, they prove that attention scores for forward and reverse configurations remain positively correlated. Specifically, for RoPE, the expected correlation is bounded below by a positive value dependent on the trace of the rotation matrix; for ALiBi, the correlation is perfect (1.0) because the content term is identical.
   • Gradient Alignment: They demonstrate that the gradients of the forward and reverse losses with respect to the shared storage column $Y_{:,B}$ are positively aligned. Consequently, a gradient step taken to minimize the forward loss decreases the reverse loss to the first order.

3. Experimental Validation:

   • Controlled One-Layer Experiments: Using a synthetic dataset with lowercase-uppercase pairs, the authors train one-layer MDMs and ARMs. They measure attention-score correlation, gradient cosine similarity, and reverse prediction probability.
   • Positional Encoding Interventions: They compare learned absolute embeddings, RoPE, and ALiBi to verify that position-invariant storage and specific routing geometries directly impact reverse transfer.
   • Large-Scale Validation: They fine-tune multi-layer diffusion LLMs (LLaDA and Dream) and autoregressive baselines (LLaMA, Qwen) on relation-reversal benchmarks (Parent-Child, Person-Description, T-REx) to check if the theoretical signatures persist in deep models.

Key Contributions

1. Theoretical Mechanism: The paper identifies that the mitigation of the reversal curse in MDMs is not solely due to the any-order objective, but specifically due to the coupling of forward and reverse conditionals via:
   • Position-invariant storage of input-target evidence.
   • Positively correlated attention routing under relative positional encodings.
   • Positively aligned gradients on shared storage parameters.
2. Formal Proofs: The authors provide proofs showing that:
   • Forward training strengthens the shared evidence term $Y_{A,B}$.
   • Relative positional encodings (RoPE, ALiBi) preserve positive correlation between forward and reverse attention scores.
   • Forward updates decrease reverse loss to the first order due to gradient alignment.
3. Empirical Verification: The study validates these signatures in both controlled one-layer settings and large-scale multi-layer diffusion models, showing that the theoretical mechanisms translate into improved reverse prediction probabilities.

Results

• One-Layer Validation:
  • Attention Correlation: Forward and reverse attention scores assigned to the evidence token remain strongly positively correlated across different relative distances (Fig. 3, left).
  • Gradient Alignment: The cosine similarity between forward and reverse gradients remains positive (approx. >0.6), indicating aligned optimization directions (Fig. 3, middle).
  • Prediction Transfer: Under forward-only training, the probability of the correct target token in reverse queries increases significantly (Fig. 3, right).
  • Baseline Failure: In contrast, one-layer ARMs learn the forward direction but fail to transfer to the reverse direction (Table 1).
• Positional Encoding Impact:
  • Absolute Embeddings: Result in poor reverse transfer because they make value vectors position-dependent, breaking the shared storage term.
  • RoPE: Enables nontrivial but imperfect reverse transfer due to partial attention correlation.
  • ALiBi: Yields the strongest transfer (near-perfect) because it preserves the content score perfectly, maintaining the strongest routing correlation (Table 2).
• Large-Scale Validation:
  • Multi-layer diffusion models (LLaDA, Dream) exhibit the same qualitative signatures: positive attention correlation, positive gradient alignment, and improved reverse probabilities during forward-only fine-tuning (Fig. 4).
  • ARMs (including reasoning-oriented models like Qwen-3.5) show high forward accuracy but collapse in reverse accuracy (Table 7).

Significance and Claims

The paper claims to provide a mechanistic explanation for why MDMs mitigate the reversal curse, moving beyond the observation that "any-order training helps." The significance lies in identifying the specific architectural and training dynamics required for reverse transfer:

• Storage: Evidence must be stored in a position-invariant manner (achieved by relative positional encodings affecting only attention scores, not value vectors).
• Routing: The model must maintain a correlated path to retrieve this evidence in reverse configurations.
• Dynamics: The training process must naturally align forward and reverse gradients on the shared evidence.

The authors are modest about their claims, noting that:

• The theory explains why reverse transfer can occur and predicts the sign of the transfer effect, but it does not guarantee perfect reverse accuracy (probability approaching 1.0) in complex, multi-token settings.
• The "reversal curse" mitigation is specific to the storage-routing mechanism and does not necessarily imply that MDMs possess superior general logical reasoning or abstract inversion capabilities.
• The one-layer theory serves as a minimal mechanistic explanation; the persistence of these signatures in deep models is an empirical finding, not a theoretical guarantee for all layers.

In summary, the paper establishes that the reversal curse mitigation in MDMs arises from a parameter-level coupling where shared storage and correlated routing allow forward-learned evidence to be accessible under order-reversed queries, a property fundamentally broken in standard autoregressive training.