Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge

Autoregressive large language models (LLMs) have achieved remarkable success in many complex tasks, yet they can still fail in very simple logical reasoning such as the “reversal curse” – when trained on forward knowledge data of the form “$A \rightarrow B$” (e.g., Alice’s husband is Bob), the model is unable to deduce the reversal knowledge “$B \leftarrow A$” (e.g., Bob’s wife is Alice) during test. Extensive prior research suggests that this failure is an inherent, fundamental limit of autoregressive causal LLMs, indicating that these models tend to memorize factual-level knowledge rather than capture higher-level rules. In this paper, we challenge this view by showing that this seemingly fundamental limit can be mitigated by slightly tweaking the training data with a simple regularization data recipe called the Identity Bridge of the form “$A \to A$” (e.g., The name of Alice is Alice). Theoretically, we prove that under this recipe, even a one-layer transformer can break the reversal curse by analyzing the implicit bias of gradient descent. Empirically, we show that a 1B pretrained language model finetuned with the proposed data recipe achieves a 40% success rate on reversal tasks, in stark contrast to a near-zero success rate when trained solely on forward-knowledge data. Our work provides a novel theoretical foundation for the reversal curse and offers a principled, low-cost path to encouraging LLMs to learn higher-level rules from data.

💡 Research Summary

The paper tackles a well‑known limitation of autoregressive large language models (LLMs) often called the “reversal curse”: when a model is trained only on forward relational facts of the form “A → B” (e.g., “Alice’s husband is Bob”), it fails to answer the reverse query “B ← A” (e.g., “Who is Bob’s wife?”). Prior work has argued that this asymmetry is an inherent property of causal language modeling, suggesting that LLMs merely memorize factual pairs rather than learning higher‑order relational rules.

The authors propose a minimalist data‑level regularization called the Identity Bridge. The recipe adds sentences of the form “A → A” (e.g., “The name of Alice is Alice”) to the training corpus. These statements convey no new factual information but force the model to treat an entity and an identity relation as a valid input‑output pair. By augmenting the forward‑only dataset with an equal amount of identity‑bridge examples, the model is encouraged to develop symmetric representations that can support reverse inference.

Theoretical contribution
The analysis is conducted on a one‑layer decoder‑only transformer. The key‑query matrix is fixed to zero, which makes the attention weights for the two input tokens (the entity and the relation) equal (½ each). Embeddings are idealized: entity embeddings are one‑hot, the forward relation embedding (z_{r+}) is one‑hot, and the reverse relation embedding is set to its negative, (z_{r-} = -z_{r+}). The identity relation embedding is set to zero, which mathematically corresponds to (z_{r_{id}} = z_{r+} + z_{r-}).

Under these conditions, the model’s logits simplify to a bilinear form ( \frac12 W_O W_V^\top (z_s + z_r) ). Training with cross‑entropy loss and gradient descent leads to an implicit bias toward minimizing the nuclear norm of the weight matrix while satisfying margin constraints. Lemma 3.2 (adapted from prior work) shows that any limit point of the normalized parameters solves a hard‑margin SVM problem.

When only forward data (D_{r+}) is present, the optimal solution (W^{+}{OV}) has a zero upper‑right block, meaning the model never learns a mapping for the reverse relation; the margin for any reverse query is zero (Theorem 3.3). Adding the identity‑bridge dataset (D{idn}) changes the feasible set: the optimal solution (W^{*}_{OV}) now contains positive diagonal entries in the upper‑right block, yielding a positive margin for every reverse query (Theorem 3.4). Thus, the identity bridge directly injects the missing symmetry into the learned weight matrix.

The authors also connect this regularization to Out‑of‑Context Reasoning (OCR). Proposition 3.5 demonstrates that, with a specific choice of identity‑relation embedding (the average of forward and reverse embeddings), the identity‑bridge task is mathematically equivalent to an OCR task where the model must infer a second relation from a shared intermediate fact. This equivalence provides an intuitive explanation: the identity examples act as a bridge that forces the model to compose learned transformations, thereby enabling reverse inference without ever seeing explicit reverse examples.

Empirical validation
To test whether the theory scales beyond the toy setting, the authors fine‑tune a 1 billion‑parameter pretrained LLM on a mixed dataset containing forward relational sentences and identity‑bridge sentences in a 1:1 ratio. They evaluate on a suite of reversal queries (e.g., “Bob’s wife is ?”). The identity‑bridge‑augmented model achieves a 40 % pass rate, whereas a baseline model trained only on forward data scores near zero. This substantial gap confirms that the regularization effect persists in a realistic, large‑scale setting.

Strengths

1. Simplicity – The method requires only a trivial data augmentation; no architectural changes, loss modifications, or additional supervision are needed.
2. Theoretical rigor – By leveraging the implicit bias of gradient descent, the authors provide a clean SVM‑style proof that the identity bridge forces the weight matrix to encode reverse relations.
3. Empirical relevance – Demonstrating a 40 % success rate on a 1 B model shows that the phenomenon is not limited to the idealized one‑layer case.

Limitations and open questions

• The analysis assumes a single transformer layer, zeroed key‑query matrix, and one‑hot embeddings. Real LLMs have many layers, learned attention patterns, and dense embeddings; it remains unclear how closely the theoretical predictions match the dynamics of full‑scale models.
• The identity bridge adds no new factual content, but it does increase the proportion of self‑referential sentences. The impact on overall language modeling perplexity or on unrelated downstream tasks is not reported.
• Experiments focus on binary, bijective relations (husband–wife). It is unknown whether the approach scales to many‑to‑many, hierarchical, or non‑symmetric relations.
• A 40 % success rate, while impressive relative to the baseline, is still far from practical reliability. Combining identity‑bridge augmentation with other techniques (e.g., symmetric loss terms, contrastive training, or explicit rule‑learning modules) may be necessary to reach higher accuracies.

Conclusion
The paper introduces a novel, low‑cost data‑level regularization—Identity Bridge—that theoretically and empirically breaks the reversal curse in autoregressive LLMs. By showing that a simple “A → A” augmentation reshapes the implicit bias of gradient descent to embed reverse mappings, the work challenges the prevailing belief that explicit reverse examples are required. This insight opens a new research direction: leveraging carefully crafted regularization data to coax LLMs into learning higher‑order relational rules without altering their architecture or training objectives. Future work should explore multi‑layer extensions, broader relation families, and hybrid strategies to push reversal performance toward practical levels.