Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge

Xutao Ma, Yixiao Huang, Hanlin Zhu, Somayeh Sojoudi

TL;DR

This work addresses the reversal curse in autoregressive language models, where learning a forward relation $r_{+}$ fails to infer its inverse $r_{-}$. It introduces Identity Bridge regularization by adding identity samples with $r_ ext{id}$ to the training data, reformulating the problem without changing the model or loss. The authors prove, via implicit-bias analysis, that even a one-layer transformer can break reversal when augmented with identity-bridge data, and show an OCR-equivalence interpretation that clarifies the mechanism. Empirically, finetuning a 1B-scale real LLM with the proposed data recipe yields as much as ~40% reversal accuracy on real reversal tasks, confirming a practical, low-cost path to enabling LLMs to learn higher-level rules from data.

Abstract

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.

Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge

TL;DR

This work addresses the reversal curse in autoregressive language models, where learning a forward relation fails to infer its inverse . It introduces Identity Bridge regularization by adding identity samples with to the training data, reformulating the problem without changing the model or loss. The authors prove, via implicit-bias analysis, that even a one-layer transformer can break reversal when augmented with identity-bridge data, and show an OCR-equivalence interpretation that clarifies the mechanism. Empirically, finetuning a 1B-scale real LLM with the proposed data recipe yields as much as ~40% reversal accuracy on real reversal tasks, confirming a practical, low-cost path to enabling LLMs to learn higher-level rules from data.

Abstract

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 "" (e.g., Alice's husband is Bob), the model is unable to deduce the reversal knowledge "" (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 "" (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.
Paper Structure (24 sections, 5 theorems, 46 equations, 10 figures)

This paper contains 24 sections, 5 theorems, 46 equations, 10 figures.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Reversal reasoning task.
  4. Dataset.
  5. One-layer transformer.
  6. Loss function.
  7. Theoretical Results
  8. Implicit Bias Explains the Reversal Curse
  9. Breaking the Reversal Curse via Identity Bridge
  10. Relation to Out-of-Context Reasoning
  11. Experiments
  12. One-Layer Transformer Experiments
  13. Real Large Language Model Experiments
  14. Ablation Experiments on Dataset
  15. Ablation Experiments on Entity Token Length
  16. ...and 9 more sections

Key Result

Lemma 3.2

Suppose assum:fix_kq holds. Consider running gradient descent with a small enough learning rate or gradient flow on loss eq:loss where the model is defined by eq:ntp and eq:transformer_logit. If there exists a time $t_0$, such that $\mathcal{L}_{\mathcal{D}}(\boldsymbol{\theta}(t_0)) < 1$, then any where $\| \cdot \|_{\star}$ denotes the nuclear norm.

Figures (10)

  • Figure 1: $\boldsymbol{W}_{OV}^{\textup{+}}$ solution in \ref{['thm:reversal_happen']} with forward relation dataset. In $\boldsymbol{W}_{OV}^{\textup{+}}$, the diagonal weight of the upper right block is equal to the off-diagonal weight of it, which means when tested with the reversal data $[b_i,r_{-}]$, the trained model will output equal logits over all $a\in \mathcal{A}$. Thus, only training with the forward relation dataset will lead to the reversal curse.
  • Figure 2: $\boldsymbol{W}_{OV}^\ast$ solution in \ref{['thm:idn_bridge']} with identity bridge regularized dataset. In $\boldsymbol{W}_{OV}^{\ast}$, the diagonal weight of the upper right block is larger than the off-diagonal weight of it, so the logits of the correct reversal answer $a_i$ will be larger than other answers when tested with $[b_i,r_{-}]$. Thus, the model can break the reversal curse with identity bridge regularized dataset.
  • Figure 3: Illustration of \ref{['prop:relation_OCR']} when $\boldsymbol{z}_{r_{+}} = - \boldsymbol{z}_{r_{-}}$. Identity form refers to the form of the identity regularized dataset, and OCR form refers to its OCR form given in \ref{['prop:relation_OCR']}. In the concrete "Husband-Wife" example in the Identity form, we set $a$ to "Alice", $b$ to "Bob", $r_{+}$ to "husband", $r_{-}$ to "wife", and $r_\text{id}$ to "name". In the corresponding OCR form, $s_a = a+ r_{+}$ is "Alice's husband", $s_b = b$ is Bob, $r_1 = r_\text{id}$ is "name", and $r_2 = r_{-}$ is "wife". The left Identity form and the right OCR form are semantically equivalent and share the same test dataset.
  • Figure 4: Reversal test loss and mean reciprocal rank (MRR) of forward relation only training data vs. identity bridge regularized data. Without the identity bridge, the model stays around the initialization level, while the model can generalize to all the reversal tests after adding the identity bridge regularization.
  • Figure 5: (a). $\boldsymbol{W}_\text{O}\boldsymbol{W}_\text{V}^{\mathrm{T}}$ weight after training with the forward relation dataset; (b) $\boldsymbol{W}_\text{O}\boldsymbol{W}_\text{V}^{\mathrm{T}}$ weight after training with the identity bridge regularized dataset. The trained weights match the SVM solutions in \ref{['thm:reversal_happen']} and \ref{['thm:idn_bridge']}.
  • ...and 5 more figures

Theorems & Definitions (10)

  • Lemma 3.2: SVM form, Theorem 1 in huang2025generalization
  • Theorem 3.3
  • proof
  • Theorem 3.4
  • proof
  • Proposition 3.5
  • proof
  • Lemma A.1
  • proof
  • proof