Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics

Hanlin Zhu, Baihe Huang, Shaolun Zhang, Michael Jordan, Jiantao Jiao +2 more
5/7/2024
cs.LGcs.CL

Official code: marlo-z/reversal_curse_analysis (5 stars, MIT, last commit 2025-01-17).

Abstract

Auto-regressive large language models (LLMs) show impressive capacities to solve many complex reasoning tasks while struggling with some simple logical reasoning tasks such as inverse search: when trained on '$A \to B$' (e.g., 'Tom is the parent of John'), LLM fails to directly conclude '$B \gets A$' (e.g., 'John is the child of Tom') during inference even if the two sentences are semantically identical, which is known as the 'reversal curse'. In this paper, we theoretically analyze the reversal curse via the training dynamics of (stochastic) gradient descent for two auto-regressive models: (1) a bilinear model that can be viewed as a simplification of a one-layer transformer; (2) one-layer transformers under certain assumptions. Our analysis reveals that for both models, the reversal curse is a consequence of the (effective) model weights 'asymmetry', i.e., the increase of weights from a token $A$ to token $B$ during training does not necessarily cause the increase of the weights from $B$ to $A$, which is caused by the training dynamics under certain choice of loss function and the optimization space of model parameters. Moreover, our analysis can be naturally applied to other logical reasoning tasks such as chain-of-thought (COT), which provides a new perspective different from previous work that focuses on expressivity. Finally, we conduct experiments to validate our theory on multi-layer transformers under different settings. Our code is available at https://github.com/marlo-z/reversal_curse_analysis/.

View on arXivView PDF

Code Implementations(1)

MIT

Code for 'Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics'

50PythonOct 27, 20241 years agoMIT
gpt-2large-language-modelstraining-dynamicstransformers

Similar Papers

Cite this paper

@article{zhu2024towards,
  title  = {Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics},
  author = {Hanlin Zhu and Baihe Huang and Shaolun Zhang and Michael Jordan and Jiantao Jiao and Yuandong Tian and Stuart Russell},
  year   = {2024},
  eprint = {2405.04669},
  archivePrefix = {arXiv},
  doi    = {10.48550/arXiv.2405.04669},
  url    = {http://arxiv.org/abs/2405.04669v2},
  journal = {NEURIPS 2024 2024}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2405.04669.svg)](https://codeofpaper.com/paper/2405.04669)

Discussion