Title: Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics
Venue: NeurIPS 2024
Accepted: yes
Keywords: logical reasoning

reversal curse

training dynamics

autoregressive transformer

chain of thought
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/](https://github.com/marlo-z/reversal_curse_analysis/).

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

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 → B ' (e.g., Tom is the parent of John ), LLM fails to directly conclude ' B ← 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/ .

1 Introductions

Large language models (LLMs) have shown great performance in solving complex reasoning tasks that require multiple reasoning steps through in-context learning (ICL), such as zero-shot learning [1, 2], few-shot learning [3, 4, 5], or via further fine-tuning [6, 7, 8]. However, without the above inferencetime techniques or model fine-tuning (probably combined with data manipulations), an auto-regressive LLM might struggle with simple logical reasoning tasks that require multiple reasoning steps learned during training separately [9], where the reversal curse [10] serves as a well-known example.

The reversal curse refers to the phenomenon that an auto-regressive LLM that learns ' A → B ' (e.g., Tom is the parent of John ) during training fails to generalize to the reverse direction ' B ← A ' (e.g., John is the child of Tom ) even if the pair of relationship ' → ' and ' ← ' are reverse to each other and the two sentences are semantically identical. Although some previous works propose different methods to mitigate the reversal curse, including reversing the training dataset [11, 12] and training on different objectives such as autoregressive blank infilling [13], these methods might negatively affect the model performance on other tasks since they either alter the dataset or the model architecture. Without dataset manipulation or changing the auto-regressive nature (causal structure) of the model, there are two other candidate solutions to mitigate the reversal curse.

First, one might constrain the model parameters to satisfy a higher-level regularity for specific relationships. For example, a reversal-type regularity can be viewed as a pair of relationships ( → , ← ) and two sets A , B such that a model trained on ' A → B ' will also increase its probability of ' B ← A ' for all A ∈ A , B ∈ B , which induces a subspace of model parameters that satisfy this regularity. If one can train the model within this subspace, then training on ' A → B ' can, by definition, help to learn ' B ← A '. However, for a general LLM, it is extremely challenging to find the subspace and manually hard-code the constraint during optimization even for one pair of relationships, not to mention there are numerous relationships. Since it is intractable to manually hard-code the constraints to the model parameter, one can alternatively expect the model to learn the higher-level regularity by training samples under unconstrained optimization. However, this is also hard, according to our analysis, through the popular cross-entropy (CE) loss that aims to maximize the next token prediction probability for the models studied in our paper.

Second, one can use a different loss function which is 'symmetric', rather than the popular CE loss. However, the 'symmetric' loss might drive the model to learn meaningless sentences. For example, when trained on the sentence 'John is tall', a 'symmetric' loss function might drive the model to learn 'tall is John', which is not what we expect. To prevent the model from the above undesired behavior, in practice, CE loss is still widely-used.

Therefore, in this paper, we analyze the reversal curse via training dynamics of the widely-used unconstrained optimization for the CE loss. We summarize our main contributions as follows:

We theoretically analyze the reversal curse where training or test sequences have the form ' A → B ' or ' B ← A ' via training dynamics of (stochastic) gradient descent under two auto-regressive models: a bilinear model (Section 3) and one-layer transformers under certain assumptions similar to [14] (Section 4). The analysis of both models reveals that the widely-used unconstrained optimization for CE loss leads to model weights asymmetry , i.e., the increase of (effective) weights (after reparameterization) from the token A to token B 1 during training does not necessarily cause the increase of the weights from B to A , which further causes the reversal curse. Although the (effective) weights from A to B and from B to A might be related to some extent due to reparameterization, their correlation is weak and thus show asymmetry as empirically verified in Section 5.

The techniques we used to analyze the reversal curse can be applied to other logical reasoning tasks. In particular, we use the above framework to analyze chain-of-thought (COT) [4], and we show that a model trained on ' A → B ' and ' B → C ' separately struggles to directly conclude ' A ❀ C ' without COT even if it is logically true (Section 4.2). Different from the previous work [15] that theoretically studies COT through the expressivity of transformers, our work provides a new perspective through training dynamics.

We also empirically validate our theoretical results on multi-layer transformers (Section 5).

The asymmetry of auto-regressive model weights caused by widely-used unconstrained optimization for CE loss indicates that auto-regressive LLMs might not automatically deduce certain types of conclusions using separate knowledge learned during training under current popular training paradigms: to make a model predicting token B where the input token is A , the model might need to see B following A in the same sequence during the training set. This also highlights the importance of ICL, data augmentation, or planning for LLMs with the current popular causal transformer-based structures to solve complex reasoning tasks.

1.1 Related works

LLMReasoning. The strong performance of LLMs on reasoning tasks [3, 7, 16, 2, 17, 4, 18, 19] has prompted many studies on the reasoning capabilities of LLMs. [20] argues that transformers perform implicit Bayesian inference in ICL. [21] shows that transformers implement a specific type of circuits called 'induction heads' that are key to the ICL abilities of LLMs. [22] proves that causal structures are encoded in transformer layers during the training dynamics. [23] identifies a backward chaining mechanism of transformers in deductive reasoning. Apart from in-context reasoning, LLMs still demonstrate limitations in other types of reasoning tasks [24, 25, 26].

Reversal Curse. [10] identifies the phenomenon of reversal curse. This drawback of LLMs is also demonstrated in [27]. [9] studies a similar phenomenonin which LLMs face difficulty in manipulating already learned knowledge. Several paper studies eliminating the reversal curse by extending causal attention to bidirectional attention [13], training on reversed samples [12], permuting semantic units [11], or introducing reverse logic data [28]. Given all the empirical works, theoretical analysis of the reversal curse phenomenon remains scarce.

Expressivity of LLMs. There is a long line of works [29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 21, 45] studying the behavior of LLMs through the expressivity of transformers. It has been shown that transformers can implement simple functions such as sparse linear functions, two-layer neural networks, and decision trees [42], gradient descent [37, 44, 46], automata [47], Turing machines [48], variational inference [49], and bandit algorithms [50]. Different from [15] that study COT via expressivity, we analyze reversal curse and COT via training dynamics.

Training dynamics of LLMs. There are rich literatures in the optimization of attention layers [51, 52, 53, 54, 55, 56, 57]. [58, 59] study the dynamics of a single linear attention layer in in-context linear regression. [60] proves convergence of one-layer transformers in random feature regime. [61] shows the convergence of gradient descent on one-layer transformers in in-context linear regression with orthogonal data. [14] studies the convergence of one-layer transformers in a class of next-token prediction tasks. [62] studies training dynamics of multi-layer transformers. [22] studies gradient descent on a class of two-layer transformers in in-context learning tasks with latent causal structures. Our paper studies the reversal curse via training dynamics under both bilinear settings and one-layer transformers. For one-layer transformers, we use the same framework as [14] without the need for certain technical assumptions such as long input sequences, different learning rates for different parameters (except for Appendix C.3), or weak correlations that are required for [14]. Besides, we focus on the generalization ability of models for logical reasoning tasks while [14] mainly focus on optimization, and we identify the asymmetry and intransitivity properties of model weights, which are the core reasons for the failure of LLM for certain types of logical reasoning. Moreover, our analysis of the bilinear model only requires the embedding to be almost orthonormal, while [14] essentially assumed the embedding vectors to be fixed and one-hot.

2 Preliminaries

Basic notations. For any integer N > 0 , we use [ N ] to denote the set { 1 , 2 , . . . , N } . Let R , N denote the set of real numbers and natural numbers, respectively. For real variables x 1 , . . . , x n , we use poly( x 1 , . . . , x n ) to denote the polynomial of x 1 , . . . , x n . We use f ( n ) ≲ g ( n ) if there exists a constant C > 0 s.t. f ( n ) ≤ Cg ( n ) , ∀ n ; we say g ( n ) ≳ f ( n ) if f ( n ) ≲ g ( n ) .

We use e i to denote one-hot vectors where only the i -th entry of e i equals one and all other entries are zero. We use 1 to denote all-one vectors, 0 to denote zero vectors or zero matrices, and I to denote the identity matrix. We will also add subscripts when we want to explicitly show the dimension, such as 0 d , I d for d -dimensional zero vector and d × d identity matrix. We use ⊗ to denote tensor product of vectors or matrices and use x ⊗ 2 and A ⊗ 2 to denote x ⊗ x and A ⊗ A for vector x and matrix A .

We use N ( µ , Σ) (or adding subscripts such as N d ( · , · ) if we want to show dimensions explicitly) to denote the (multi-variate) Gaussian distribution with mean µ and covariance Σ . Also, we use ∆( X ) to denote the set of distributions over a set X and use E [ · ] to denote expectation. For any dataset D = { x 1 , x 2 , . . . , x n } where x i ∈ X and a function f : X → R , we define the empirical expectation over the dataset as E D [ f ] = 1 n ∑ n i =1 f ( x i ) . See additional notations in Appendix A.

Auto-regressive models. Define the vocabulary V = [ M ] for a positive integer M > 0 which is the size of the vocabulary. Let x = ( x 1 , x 2 , . . . , x T ) be a sequence of tokens of length T where each token x t ∈ V , ∀ t ∈ [ T ] . See Table 1 for notations of different tokens used in Section 4. We study auto-regressive models p θ ( ·| x ) ∈ ∆( V ) parameterized by θ that take the sequence x as input and predict the distribution of the next token x T +1 ∈ V . For both models that we study in this paper, the next token probability is modeled as the softmax applied to the logits l θ ( ·| x ) ∈ R M of each token in the vocabulary, i.e., p θ ( y | x ) = exp( l θ ( y | x )) ∑ v ∈V exp( l θ ( v | x )) , ∀ y ∈ V . Also, each token v ∈ V has a corresponding (fixed or learnable) embedding vector u v ∈ R d .

3 Bilinear Models

We start analyzing the reversal curse under bilinear models, which can be viewed as simplified one-layer transformers with input length one and decoder layer only. Also, in this section, we assume the embeddings of each token are fixed, so we directly use the embedding vector to represent a token.

Datasets. Assume the vocabulary has size m where each token v 1 , v 2 , . . . , v m i.i.d. ∼ N d (0 d , 1 d I d ) . Let V = { v 1 , . . . , v m } and let X = { x 1 , . . . , x n } and Y = { y 1 , . . . , y n } be disjoint random subsets of V . Assume all training and test sequences have a length of two. For any 2 ≤ i ≤ n , the training dataset contains both sequence ( x i , y i ) and ( y i , x i ) . In addition, the training set contains ( x 1 , y 1 ) while the test set only contain one example ( y 1 , x 1 ) . During training, the model learns both ( x i , y i ) and ( y i , x i ) for i ≥ 2 to conclude that ( x i , y i ) is equivalent to ( y i , x i ) . For example, X is a set of names and Y is a set of books. The sequence ( x i , y i ) means ' x i is the author of y i ', and the sentence ( y i , x i ) means ' y i is written by x i '. We test whether the model is able to infer an unseen sequence ( y 1 , x 1 ) given the training data which includes the other direction ( x 1 , y 1 ) .

Bilinear model. We consider a bilinear model parameterized by Θ ∈ R d × d of which the input contains single token x ∈ V . The logits of the next token y ∈ V is defined as l Θ ( y | x ) = x ⊤ Θ y which is bilinear in x and y , and thus the next token probability is p Θ ( y | x ) = exp( l Θ ( y | x )) ∑ v ∈V exp( l Θ ( v | x )) . The training loss for the bilinear model is the cross-entropy loss L (Θ) = 1 2 n -1 ( ∑ n i =1 -log p Θ ( y i | x i ) + ∑ n i =2 -log p Θ ( x i | y i )) and the test loss (reversal loss) is L rev (Θ) = -log p Θ ( x 1 | y 1 ) . We study the training dynamics of gradient flow d Θ t dt = -∇L (Θ t ) with the initialization Θ 0 that can be either randomly sampled from N ( 0 ⊗ 2 , σ 2 I ⊗ 2 ) or set as a pretrained parameter satisfying 1 2 m < p Θ 0 ( y i | x i ) , p Θ 0 ( x i | y i ) < 2 m for all i ∈ [ n ] . The following theorem shows a separation during training dynamics.

Theorem 1 (Separation of training dynamics (informal statement of Theorem 5)) . Fix any δ, ϵ ∈ (0 , 1) . For small σ and d ≥ poly( n, m, 1 /ϵ, log(1 /δ )) , with probability at least 1 -δ , we have

Theorem 1 shows that the reversal loss is lower bounded by the training loss. Note that for large d and small ϵ close to 0, ( L (Θ t ) / L (Θ 0 )) ϵ is close to 1 and thus L rev (Θ t ) ≳ L rev (Θ 0 ) which implies that p Θ ( x 1 | y 1 ) remains small during training. We summarize the above argument in Theorem 2.

Theorem 2 (Lower bound of reversal loss (informal statement of Theorem 6) . Fix arbitrary c > 0 and C ≤ log( m/ 2) . Suppose σ is small and d ≥ poly( n, m, log(1 /δ ) , log c, 1 / log C ) . With probability at least 1 -δ , it holds that L rev (Θ τ ) ≥ C , where τ denotes the first time such that L (Θ t ) ≤ c .

The proofs of Theorems 1 and 2 are deferred to Appendix B. Theorem 2 implies that for large d 2 , while the training loss can be trained to be arbitrarily small, the reversal loss remains large. In other words, the model fails to infer an unseen sequence ( y 1 , x 1 ) given the training data which includes the other direction ( x 1 , y 1 ) . Furthermore, even if the model is fine-tuned on new data from a pre-trained parameter Θ 0 that initially grasps the concept of reversal and satisfies 1 2 m < p Θ 0 ( y i | x i ) , p Θ 0 ( x i | y i ) < 2 m for new data, it fails to extend this understanding to new, unseen data.

A core reason that the reversal curse happens on the above bilinear model is that the parameter matrix Θ t is asymmetric. Consequently, the logits l Θ t ( y | x ) = x ⊤ Θ t y and l Θ t ( x | y ) = y ⊤ Θ t x generally differ. Consider a special case where each v i is a one-hot vector. Then l Θ ( y | x ) = x ⊤ Θ y = Θ ij and l Θ ( x | y ) = y ⊤ Θ x = Θ ji for x = e i , y = e j . Training on ( x, y ) can increase Θ ij but not Θ ji , which means the model does not automatically learn the reversal direction ( y, x ) from ( x, y ) . In Section 4, we will show that for one-layer transformers, the reversal curse is mainly caused by the same reason, i.e., the asymmetry of the model weights.

4 One-Layer Transformers

In Section 3, we analyzed the reversal curse under a bilinear model. In this section, we analyze the reversal curse for one-layer transformers in a similar setting to [14] via training dynamics. We also extend our analysis to chain-of-thought in Section 4.2.

Basic notations. Let V = [ M ] be the vocabulary. For any token x ∈ [ M ] , we also use the corresponding one-hot vector x = e x ∈ R M to represent it. Let U = [ u 1 , u 2 , . . . , u M ] ⊤ ∈ R M × d be the embedding matrix, where u x ∈ R d is the embedding of token x ∈ [ M ] . Note that U ⊤ x = u x . Consider the i -th training sample in the dataset, x [ i ] = ( x 1 [ i ] , . . . , x T [ i ] [ i ] , x T [ i ]+1 [ i ]) , a sequence of tokens of length T [ i ] + 1 . Here, x T [ i ]+1 [ i ] is the next token (or equivalently the label) to be predicted, x T [ i ] [ i ] is the query token, and ( x 1 [ i ] , . . . , x T [ i ] -1 [ i ]) are contextual tokens. For token x t [ i ] , we also use its one-hot vector x t [ i ] = e x t [ i ] ∈ R M to represent it. Define the contextual token matrix X [ i ] = [ x 1 [ i ] , . . . , x T [ i ] -1 [ i ]] ⊤ ∈ R ( T [ i ] -1) × M . We omit all i in notations when the context is clear.

One-layer transformer. For a training sample x = ( x 1 , . . . , x T , x T +1 ) , its contextual token matrix X = [ x 1 , . . . , x T -1 ] ⊤ and thus XU = [ u x 1 , . . . , u x T -1 ] ⊤ contains the contextual token embeddings. We study one-layer transformers in the same setting as [14]. In particular, for an input token sequence ( x 1 , . . . , x T ) , after the one-layer self-attention, we can obtain ˜ u T = U ⊤ LN ( X ⊤ b T ) √

(after softmax) that query token x T attend to each contextual token 3 , LN ( x ) = x / ∥ x ∥ 2 is the ℓ 2 -normalization operator, and W Q , W K ∈ R d × d k are trainable query and key matrices respectively. The logit of x ∈ [ M ] is then calculated by a decoder layer, i.e., l θ ( x | x 1 , . . . , x T ) = u ⊤ x W V ˜ u T , where θ encodes all parameters in the transformer, and W V ∈ R d × d can be viewed as a reparameterization of value and output matrices. Finally, the next token prediction probability is obtained by

We use the cross-entropy loss function to train the model over the whole training set D train , i.e., max U,W K ,W Q ,W V J ≜ E D train [log p θ ( x T +1 | x 1 , . . . , x T )] .

Reparameterization. Similar to [14], we define Z = UW Q W ⊤ K U ⊤ / √ d ∈ R M × M and Y = UW ⊤ V U ⊤ ∈ R M × M and are interested in their dynamics after reperameterization. Then, the attention score (after softmax) and next token probability become

and the objective can be written as

Let η Y , η Z be the learning rate of matrices Y and Z respectively. Then the gradient of Y and Z can be characterized by the following lemma:

Lemma 1 (Gradient of Y and Z for 1-layer transformer, Lemma 1 of [14]) . The gradient of Y and Z w.r.t. (2) of batch size 1 and learning rate η Y and η Z can be written as

where P ⊥ v ≜ I -vv ⊤ / ∥ v ∥ 2 2 projects any vector to orthogonal complement of v , α = [ α 1 , α 2 , . . . , α M ] ⊤ ∈ R M with α = exp ( Y ⊤ LN ( X ⊤ b T ) ) / 1 ⊤ exp ( Y ⊤ LN ( X ⊤ b T ) ) .

Proof. ˙ Y , ˙ Z can be obtained by direct calculation. One can refer to the proof of Lemma 1 of [14].

4.1 Main results for the reversal curse

In this section, we analyze the reversal curse where data points are three-token sentences ' A → B ' or ' B ← A '. For each sentence, A and B are two distinct tokens that represent two entities, and ' → ' and ' ← ' are two special tokens representing a pair of relationships inverse to each other.

Datasets. Let N train > 0 , N (1) test > 0 and N (2) test > 0 and denote N total = N train + N (1) test + N (2) test . Let A i , B i ∈ V , ∀ i ∈ [ N total ] be 2 N total distinct tokens representing distinct entities. Let → , ←∈ V be two additional different tokens that represent two inverse relationships. Specifically, we have A i → B i and B i ← A i for all i ∈ [ N total ] . For notation convenience, we define the following three index sets

The training set D train consists of all A i → B i and B i ← A i for i ∈ I train. In addition, D train contains A i → B i for i ∈ I (1) test and B i ← A i for i ∈ I (2) test . For convenience, we let N = |D train | to be the size of the training set. The test set D test consists of B i ← A i for i ∈ I (1) test and A i → B i for i ∈ I (2) test . Under our construction of the dataset, the LLM will learn the relationship between A i and B i for i ∈ I train in both directions to deduce that → is reverse to ← , and learn the relationship between A i and B i for i ∈ I (1) test ∪ I (2) test in one direction and will be tested for the other.

We use p θ ( A i | B i ← ) and p θ ( B i | A i → ) to more compactly represent p θ ( x 3 = A i | x 1 = B i , x 2 = ← ) and p θ ( x 3 = B i | x 1 = A i , x 2 = → ) , respectively. Our goal is to prove through the training dynamics of one-layer transformers that the test probability remains negligible during training. In particular, we are interested in p θ ( A i | B i ← ) , ∀ i ∈ I (1) test and p θ ( B i | A i → ) , ∀ i ∈ I (2) test .

For convenience, we assume zero-initialization Y (0) = 0 and Z (0) = 0 . This is the same as [14] and is reasonable since empirically, Y and Z are usually initialized as inner products of d -dimensional vectors with i.i.d Gaussian entries, and thus are almost zero (Lemma 8 in Appendix B). The following proposition shows the initial train/test probabilities are uniform over the vocabulary V .

Proposition 4.1 (Initial probability under zero initializaion) . Assume the transformer is under zero-initialization θ (0) = ( Y (0) , Z (0)) with Y (0) = 0 and Z (0) = 0 . For any i ∈ [ N total ] , we have

The proof is deferred to Appendix C.1.1. Proposition 4.1 shows that initially, the probability of predicting any B (or A , respectively) given any A → (or B ← , respectively) as input is uniform over the whole vocabulary. When Y (0) and Z (0) are not exactly 0 but close to 0 , the initial prediction will still be close to the uniform distribution, which is similar to Lemma 6. Next we analyze the dynamics of p θ ( t ) ( B i | A i → ) and p θ ( t ) ( A i | B i ← ) .

Proposition 4.2 (Next token probability) . For input sequence ( x 1 , x 2 ) , the next token probability under parameters θ ( t ) is p θ ( t ) ( x | x 1 , x 2 ) = exp ( Y ( t ) x 1 ,x ) / ∑ x ′ ∈ [ M ] exp( Y ( t ) x 1 ,x ′ ) , where Y ( t ) i,j is the entry of the matrix Y ( t ) at row i and column j .

The proof is also deferred to Appendix C.1.1. According to Proposition 4.2, the next token probability when the entity in the input is x 1 is determined by the x 1 -th row of the matrix Y ( t ) . Another nice property indicated by Proposition 4.2 is that we don't need to keep track of the dynamics of Z ( t ) , which could greatly simplify the analysis. The following lemma shows the dynamics of Y ( t ) .

Lemma 2 (Dynamics of Y ( t ) ) . Assume we run SGD with batch size 1 4 , and assume M ≫ 100 and 1 M 0 . 99 ≪ η Y < 1 . Let t ≳ N ln M η Y and let Y ( t ) i denote the i -th row of Y ( t ) and Y ( t ) ij denote the ( i, j ) -th entry of Y ( t ) . Then for training sequence ( x 1 , x 2 , x 3 ) ∈ D train at time t , we have

̸

and for any test sequence ( x 1 , x 2 , x 3 ) ∈ D test, we have Y ( t ) x 1 ,x = 0 , ∀ x ∈ [ M ]

The proof of Lemma 2 is presented in Appendix C.1.2. Lemma 2 implies the asymmetry of the model weights Y ( t ) : for two tokens x 1 , x 3 , when x 1 appears as a contextual token and x 3 serves as the next token in the same training sequence, the model weights Y ( t ) x 1 ,x 3 gets increased during training while Y ( t ) x 3 ,x 1 will not get increased. Combining Proposition 4.2, we can obtain our main theorem for the reversal curse.

Theorem 3 (Reversal curse) . Assume we run SGD with batch size 1, and assume M ≫ 100 and 1 M 0 . 99 ≪ η Y < 1 . Let t ≳ N ln M η Y denote the time step which also satisfies ln t ≳ ln( NM/η Y ) . For training sequence ( x 1 , x 2 , x 3 ) ∈ D train at time t , we have

for some constant c > 0 , and for any test sequence ( x 1 , x 2 , x 3 ) ∈ D test that is not included in the training set D train, we have p θ ( t ) ( x 3 | x 1 , x 2 ) ≤ 1 /M .

Theorem 3 shows that although the direction presented in the training set can be learned nearly perfectly, the model's next token prediction of the reverse direction is almost a random guess. The proof is deferred to Appendix C.1.3. We also empirically validate the above results for multi-layer transformers in Section 5.

4.2 Chain-of-thought

In this section, we extend our analysis in Section 4.1 to study other logical relationships. In particular, we study chain-of-thought (COT) [4] and show its importance via training dynamics. COT encourages LLMs to output a series of intermediate reasoning steps to increase their performance. Consider the simplest example, where the model learns two facts that A → B and B → C , and we want to test whether the model is able to directly conclude that A ❀ C . COT indicates that if an LLM is only trained on A → B and B → C , it would be easier for the model to deduce A ❀ C during the inference time if the model can first output the intermediate steps A → B and B → C , instead of directly predicting the next token C given the input ' A ❀ '. The failure of directly deducing A ❀ C is also empirically observed by [9].

Theoretically, [15] shows the importance of COT for some complex reasoning tasks through the lens of the expressivity of transformers. In this section, we show the importance of COT through a different angle, i.e., training dynamics. We show that for the above simplest two-step reasoning, without COT, the model is not able to directly predict C given the input ' A ❀ ' even if it learns A → B and B → C .

Theorem 4 (Importance of chain-of-thought, informal statement of Theorem 7) . Under certain assumptions as stated in Theorem 7, for any A i , B i , C i s.t. A i → B i and B i → C i are in the training set but A i ❀ C i is not, we have

We defer the details of the dataset construction and proof to Appendix C.2. Theorem 4 shows that although the LLM learns A i → B i and B i → C i nearly perfectly, it cannot directly deduce A i ❀ C i . Analogous to the asymmetry of causal transformer weights as we discussed in Section 4.1, our analysis of COT reveals another property, i.e., intransitivity: training the weights associated with A to B and B to C does not necessarily increase the weights associated with A to C .

We also emphasize that the model fails to directly deduce A i ❀ C i when the two intermediate steps A i → B i and B i → C i are trained separately . If the two steps are concatenated into a single training sequence, it is possible that the model learns A i ❀ C i directly [19].

4.3 Roles of the attention score matrix

During the analysis of Sections 4.1 and 4.2, we show that the reversal curse and the importance of COT are largely due to the asymmetry and intransitivity of causal transformer weights (in our case, the weight matrix Y ( t ) ). However, it seems that the dynamics of the attention score matrix Z ( t ) do not impact the model performance. Below, we briefly discuss the role of the attention score matrix Z ( t ) .

In (1), the attention score is used to calculate the weights b tT , where a contextual token x t with a larger attention score attended by the query token x T has a larger weight. Note that we use the same formulation as the previous work [14] where the query token will not attend to itself. Therefore, for a three-token training sequence, the weights b 12 is always one since there is only one contextual token x 1 , no matter whether the value of the attention score is high or low.

However, consider a slightly different setting, where the relationship is represented by two tokens. In that case, x 1 = A i , x 2 = R 1 , x 3 = R 2 , x 4 = B i , and there are two contextual tokens A i and R 1 . The role of the attention score is then to select the important token, i.e., A i , by putting more weights on it. Theorem 2 of [14] showed that under certain assumptions, the query token R 2 will attend more to 'distinct tokens' A i and less to the 'common token' R 1 . Therefore, the query token R 2 will eventually put all weights to A i , and the remaining analysis remains the same as in Sections 4.1 and 4.2. See Appendix C.3 for a more rigorous analysis.

5 Experiments

In this section, we conduct experiments to further validate our theoretical results in Section 4 on multi-layer transformers. We show experimental results of the reversal curse in this section and COT in Appendix D. Note that in Sections 3 and 4, we theoretically proved the reversal curse for both the bilinear model and one-layer transformer under certain assumptions. Now, we empirically show that the reversal curse still happens even for multi-layer transformers. In Appendix E.2.3, we also provide empirical results that the reversal curse does not happen in ICL settings.

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

70

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

Forward direction logits (Train)

Reverse direction logits (Train)

Seen direction logits (Val)

Unseen direction logits (Val)

60

40

20

0

20

40

60

80

Dataset construction. Below, we describe how we generate our synthetic dataset for experiments on the reversal curse. We choose the vocabulary V = { 0 , 1 , . . . , N } for a specified N > 0 . We randomly sample two disjoint sets of entities A , B ⊂ V with |A| = |B| = |V| / 4 , and reserve two additional tokens for relationships → and ← , respectively. 5 Next, we specify a bijection from A to B uniformly at random. For each A i ∈ A and its corresponding B i ∈ B , we can obtain a pair of sequence ( A i → B i , B i ← A i ). We split the set of all pairs into training pairs and validation pairs. For each training pair, both sequences will be included in the training set, while for the validation pair, we randomly select one sequence for the training set and the other for the validation set. Therefore, the model will learn both directions for the training pairs and only one direction for each validation pair while being tested in the unseen direction.

Model architectures. We train multi-layer transformers based on GPT-2 architecture [63]. Figure 1 shows the results where the model has 24 layers, 12 attention heads per layer, uses absolute positional encoding, and we choose the vocabulary size of 800. The training set size is 340, and the validation set size is 60 (resulting from 140 training pairs and 60 validation pairs). We also conducted experiments with various model configurations and vocabulary sizes in Appendix E.2. Besides, all hyperparameters and different model configurations are presented in Appendix E.1.

Results. Figure 1 shows that during the training, the next token probability for training data increases a lot while the next token probability for validation data remains unchanged or gets even smaller. This is consistent with our theoretical results of Theorem 3.

According to our theoretical analysis, the reversal curse happens due to the asymmetry of model (reparameterized) weights (i.e., logits of a token given another token as input), and we also empirically validate the asymmetry for multi-layer transformers. Figure 2 shows the model weights from a token x 1 to x 3 is trained large for a training sequence ( x 1 , x 2 , x 3 ) as represented by the diagonals of the

first three matrices, while the weights from a token x 1 to x 3 remains nearly zero for a validation sequence ( x 1 , x 2 , x 3 ) as represented by the diagonals of the last matrix, which is consistent with Lemma 2. This implies that if a pair of tokens ( A , B ) only appear in the training set in one direction, then the model weights associated with the other direction will hardly get trained.

6 Conclusions

In this paper, we study the reversal curse theoretically via training dynamics of (1) a bilinear model, which is a simplification of the one-layer transformer; (2) one-layer transformers under certain technical assumptions similar to [14]. Our theoretical results suggest that a core reason the reversal curse happens in auto-regressive LLMs is the asymmetry of the model weights, and we apply our technique to prove the necessity of COT for one-layer transformers, which is mainly due to the intransitivity of model weights. The asymmetry and intransitivity of model weights caused by unconstrained optimization of CE loss indicate that an auto-regressive LLM might mainly focus on learning text sequences during training separately instead of automatically deducing indirect conclusions under the current popular training paradigms. This highlights the importance of ICL, data augmentation, or planning for current auto-regressive LLMs to solve complex reasoning tasks.

As for future directions, it would be interesting and important to study: (1) What is a unified way to characterize and study the reversal curse, COT, and other similar logical reasoning tasks? (2) Our paper mainly focuses on three-token sequences, where each entity or relationship is represented by a single token. While we empirically explored the setting where each entity might consist of multiple tokens and distinct entities might share a few tokens, it would be interesting to analyze the multiple-token setting theoretically. (3) We theoretically analyzed the bilinear model and one-layer transformer, and it would be an important future direction to extend the analysis to multi-layer transformers.

References

Laria Reynolds and Kyle McDonell. Prompt programming for large language models: Beyond the few-shot paradigm. In Extended Abstracts of the 2021 CHI Conference on Human Factors in Computing Systems , pages 1-7, 2021.

Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems , 35:22199-22213, 2022.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems , 35:24824-24837, 2022.

Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171 , 2022.

Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114 , 2021.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021.

Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems , 35:15476-15488, 2022.

Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation. arXiv preprint arXiv:2309.14402 , 2023.

Lukas Berglund, Meg Tong, Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, and Owain Evans. The reversal curse: Llms trained on" a is b" fail to learn" b is a". arXiv preprint arXiv:2309.12288 , 2023.

Qingyan Guo, Rui Wang, Junliang Guo, Xu Tan, Jiang Bian, and Yujiu Yang. Mitigating reversal curse via semantic-aware permutation training. arXiv preprint arXiv:2403.00758 , 2024.

Olga Golovneva, Zeyuan Allen-Zhu, Jason Weston, and Sainbayar Sukhbaatar. Reverse training to nurse the reversal curse. arXiv preprint arXiv:2403.13799 , 2024.

Ang Lv, Kaiyi Zhang, Shufang Xie, Quan Tu, Yuhan Chen, Ji-Rong Wen, and Rui Yan. Are we falling in a middle-intelligence trap? an analysis and mitigation of the reversal curse. arXiv preprint arXiv:2311.07468 , 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, 2023.

Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems , 36, 2024.

Wenlong Huang, Pieter Abbeel, Deepak Pathak, and Igor Mordatch. Language models as zero-shot planners: Extracting actionable knowledge for embodied agents. In International Conference on Machine Learning , pages 9118-9147. PMLR, 2022.

Jaehun Jung, Lianhui Qin, Sean Welleck, Faeze Brahman, Chandra Bhagavatula, Ronan Le Bras, and Yejin Choi. Maieutic prompting: Logically consistent reasoning with recursive explanations. arXiv preprint arXiv:2205.11822 , 2022.

Simon Jerome Han, Keith J Ransom, Andrew Perfors, and Charles Kemp. Inductive reasoning in humans and large language models. Cognitive Systems Research , 83:101155, 2024.

Xinyi Wang, Alfonso Amayuelas, Kexun Zhang, Liangming Pan, Wenhu Chen, and William Yang Wang. Understanding the reasoning ability of language models from the perspective of reasoning paths aggregation. arXiv preprint arXiv:2402.03268 , 2024.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 , 2021.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895 , 2022.

Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 , 2024.

Jannik Brinkmann, Abhay Sheshadri, Victor Levoso, Paul Swoboda, and Christian Bartelt. A mechanistic analysis of a transformer trained on a symbolic multi-step reasoning task. arXiv preprint arXiv:2402.11917 , 2024.

Karthik Valmeekam, Alberto Olmo, Sarath Sreedharan, and Subbarao Kambhampati. Large language models still can't plan (a benchmark for llms on planning and reasoning about change). arXiv preprint arXiv:2206.10498 , 2022.

Yupeng Chang, Xu Wang, Jindong Wang, Yuan Wu, Linyi Yang, Kaijie Zhu, Hao Chen, Xiaoyuan Yi, Cunxiang Wang, Yidong Wang, et al. A survey on evaluation of large language models. ACM Transactions on Intelligent Systems and Technology , 2023.

Dylan Zhang, Curt Tigges, Zory Zhang, Stella Biderman, Maxim Raginsky, and Talia Ringer. Transformer-based models are not yet perfect at learning to emulate structural recursion. arXiv preprint arXiv:2401.12947 , 2024.

Chengwen Qi, Bowen Li, Binyuan Hui, Bailin Wang, Jinyang Li, Jinwang Wu, and Yuanjun Laili. An investigation of llms' inefficacy in understanding converse relations. arXiv preprint arXiv:2310.05163 , 2023.

Ruilin Luo, Tianle Gu, Haoling Li, Junzhe Li, Zicheng Lin, Jiayi Li, and Yujiu Yang. Chain of history: Learning and forecasting with llms for temporal knowledge graph completion. arXiv preprint arXiv:2401.06072 , 2024.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009.11264 , 2020.

Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09286 , 2020.

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819 , 2018.

Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing-complete. Journal of Machine Learning Research , 22(75):1-35, 2021.

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread , 1:1, 2021.

Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint arXiv:2106.03764 , 2021.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08117 , 2023.

Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115 , 2021.

Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint arXiv:2207.04901 , 2022.

Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems , 35:21750-21764, 2022.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637 , 2023.

Shuai Li, Zhao Song, Yu Xia, Tong Yu, and Tianyi Zhou. The closeness of in-context learning and weight shifting for softmax regression. arXiv preprint arXiv:2304.13276 , 2023.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749 , 2022.

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems , 35:12071-12083, 2022.

Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309.11420 , 2023.

Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint arXiv:2310.08566 , 2023.

Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations , 2020.

Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466 , 2022.

Enric Boix-Adsera, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Transformers learn through gradual rank increase. arXiv preprint arXiv:2306.07042 , 2023.

Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. arXiv preprint arXiv:2306.00802 , 2023.

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836, 2022.

Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576 , 2023.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. Advances in Neural Information Processing Systems , 36, 2024.

Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 , 2023.

Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535 , 2023.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, October 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/ 2020.emnlp-demos.6 .

Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selection. Annals of Statistics , pages 1302-1338, 2000.

Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing , 568:127063, 2024.

A Additional Notations

Let δ ij = 1 for i = j and δ ij = 0 for i = j . For a squared matrix A ∈ R d × d , its trace is Tr( A ) = ∑ d i =1 A ii . For two matrices A,B ∈ R m × n of the same shape, their inner product is defined as ⟨ A,B ⟩ = Tr( AB ⊤ ) . For any matrix A ∈ R m × n , its (Frobenius) norm is defined as ∥ A ∥ = √ ⟨ A,A ⟩ . For any vector x = ( x 1 , . . . , x d ) ⊤ ∈ R d or a matrix A ∈ R m × n , we define the zero-norm as ∥ x ∥ 0 = ∑ d i =1 ✶ { x i = 0 } or ∥ A ∥ 0 = ∑ m i =1 ∑ n j =1 ✶ { A ij = 0 } where ✶ {·} is the indicator function.

̸

B Missing Proofs of Section 3

Theorem 5 (Separation of training dynamics, formal statement of Theorem 1) . Fix arbitrary δ, ϵ ∈ (0 , 1) . Let v 1 , . . . , v m be independently sampled from N d ( 0 d , 1 d I d ) . Let x 1 , . . . , x n and y 1 , . . . , y n be sampled uniformly at random from { v 1 , . . . , v m } without replacement. Define

Consider the gradient flow Θ t : t ≥ 0

̸

where Θ 0 ∼ N ( 0 ⊗ 2 , σ 2 · I ⊗ 2 ) or Θ 0 satisfies 1 2 m < p Θ 0 ( y i | x i ) , p Θ 0 ( x i | y i ) < 2 m for all i ∈ [ n ] . Suppose σ ≤ 1 100 ln(64 m 2 /δ ) and

With probability at least 1 -δ , we have

Proof. Let v = √ 400 n 2 m 2 log(64 m 2 n 2 /δ ) d . By Lemma 3 and Lemma 4, with probability at least 1 -δ ,

By definition of d , we have 8 v (2 n -1) log 2 (2 m ) ≤ ϵ . Notice that L (Θ t ) L (Θ 0 ) < 1 , thus

Theorem 6 (Lower bound of reversal loss, formal statement of Theorem 2) . Fix arbitrary c > 0 and C ≤ log( m/ 2) . Under the setting of Theorem 5, suppose σ ≤ 1 100 ln(64 m 2 /δ ) and

With probability at least 1 -δ ,

.

where τ denotes the first time such that L (Θ t ) ≤ c .

̸

Proof. By continuity, L (Θ τ ) = c . By Theorem 5, when

with probability at least 1 -δ ,

B.1 Training dynamics

Lemma 3 (Dynamics of the forward loss) . Let x 1 , . . . , x n , y 1 , . . . , y n , L (Θ) , and Θ t ( t ≥ 0) be defined as in Theorem 5. When σ ≤ 1 100 ln(16 n 2 /δ ) and d ≥ 1600 n 3 m 2 log(8 m 2 n 2 /δ ) , with probability at least 1 -δ , we have

Proof. For convenience, we assume x 1 , . . . , x n = v 1 , . . . , v n and y 1 , . . . , y n = v n +1 , . . . v 2 n WLOG.

Furthermore,

Under this event, applying Lemma 6, we can obtain

To ensure that the right hand side is C , we set ϵ = log C log( m/ 2) log c log(2 m ) . One may check that the definition of d satisfies Eq. (4). It follows that

Similarly, we have

and

Applying Lemma 7, with probability at least 1 -δ/ 2 , for any t ≥ 0 we have

and

Furthermore, Lemma 6 implies that with probability at least 1 -δ/

The following arguments are based on the event that the above inequalities hold.

We first show that

Now we show

By Eq. (5),

̸

̸

where the second inequality is due to ∑ j = i + n | α ( t ) i,j | = α ( t ) i,i + n = 1 -p Θ t ( y i | x i ) , ∑ j = i | β ( t ) i + n,j | = β ( t ) i + n,i = 1 -p Θ t ( x i | y i ) for all i ∈ [ n ] ; the third inequality applies Cauchy-Schwarz inequality;

̸

We prove this by contradiction. Let

By Eq. (7), it is obvious that τ > 0 . Assume without loss of generality that p Θ τ ( y 1 | x 1 ) ≤ 1 2 m . It follows that there exists δ > 0 such that p Θ t ( y 1 | x 1 ) is a decreasing function in ( τ -δ, τ ) and p Θ t ( y 1 | x 1 ) = min { p Θ t ( x i | y i ) , p Θ t ( y i | x i ) : i ∈ [ n ] } for any t ∈ ( τ -δ, τ ) . It follows that for t ∈ ( τ -δ, τ ) ,

̸

where the first inequality is from Eq. (6); the second inequality is due to ∑ j = i + n | α ( t ) i,j | = α ( t ) i,i + n = 1 -p Θ t ( y i | x i ) , ∑ j = i | β ( t ) i + n,j | = β ( t ) i + n,i = 1 -p Θ t ( x i | y i ) for all i ∈ [ n ] ; the third inequality is because p Θ t ( y 1 | x 1 ) = min { p Θ t ( x i | y i ) , p Θ t ( y i | x i ) : i ∈ [ n ] } . However, p Θ t ( y 1 | x 1 ) is a decreasing function in ( τ -δ, τ ) , a contradiction. Therefore, we conclude that

the last inequality uses the fact that p Θ t ( x i | y i ) , p Θ t ( y i | x i ) > 1 2 m for any t ≥ 0 , i ∈ [ n ] and the inequality 1 -x ≥ log x 2 log(1 / (2 m )) for x ∈ ( 1 2 m , 1) .

By Lemma 10, we conclude that ∀ t ≥ 0 ,

which completes the proof.

Lemma 4 (Dynamics of the reversal loss) . Let x 1 , . . . , x n , y 1 , . . . , y n , L rev (Θ) , and Θ t ( t ≥ 0) be defined as in Theorem 5. When σ ≤ 1 100 ln(16 n 2 /δ ) and d ≥ 400 n 2 m 2 log(8 m 2 n 2 /δ ) /ϵ 2 , with probability at least 1 -δ ,

Proof. Similar to Lemma 3, we assume x 1 , . . . , x n = v 1 , . . . , v n and y 1 , . . . , y n = v n +1 , . . . v 2 n WLOG. Let α ( t ) i,j = -p Θ t ( v j | v i ) + δ i,j -n , β ( t ) i,j = -p Θ t ( v j | v i ) + δ i -n,j . By Lemma 5,

Applying Lemma 7, with probability at least 1 -δ/ 2 , for any t ≥ 0 we have

Furthermore, Lemma 6 implies that with probability at least 1 -δ/ 2 ,

The following arguments are based on the event that the above inequalities hold. By Eq. (8),

Notice that

̸

where the first inequality uses ∑ j = i + n | α ( t ) i,j | = α ( t ) i,i + n = 1 -p Θ t ( y i | x i ) , ∑ j = i | β ( t ) i + n,j | = β ( t ) i + n,i = 1 -p Θ t ( x i | y i ) for all i ∈ [ n ] ; the third inequality uses the fact that 1 -x ≤ -log x ; the last inequality applies Lemma 3. Similarly,

̸

Combining, we have

By Lemma 10, we conclude that ∀ t ≥ 0 ,

This completes the proof.

Lemma 5 (Gradient of the loss function) . Define

Then we have

Proof. We have

The statements follow immediately.

B.2 Initialization

Lemma 6 (Initial distributions are all close to uniform) . Fix any δ ∈ (0 , 1) . Let x 1 , x 2 , . . . , x n i.i.d. ∼ N d ( 0 d , 1 d I d ) . Let Θ ∈ R d × d , where each Θ ij i.i.d. ∼ N (0 , σ 2 ) independent of x 1 , . . . , x n . For any i, j ∈ [ n ] , define

Then when σ 2 ≤ 1 100 ln(4 n 2 /δ ) and d ≥ 400 log(1 / (2 δn 2 )) /ϵ 2 , with probability at least 1 -δ ,

Proof. Let v = 0 . 2 . By Lemma 8, with probability at least 1 -δ , we have

Conditioned on the above high-probability event, we can further obtain that for any j ∈ [ n ]

and

It follows that

which completes the proof.

B.3 Subspace embedding

Lemma 7 ( ℓ 1 -subspace embedding of Gaussian second-order tensors) . Let z 1 , . . . , z n be independently sampled from N d ( 0 d , 1 d I d ) . Let I 1 , I 2 ⊆ [ n ] × [ n ] be two index sets. Let I 0 = I 1 ∩ I 2 . If d ≥ 64 log(2 n 2 /δ ) /ϵ 2 , then with probability at least 1 -δ ,

holds for any α i,j , β i,j . Furthermore, if d ≥ 64 k 2 log(2 n 2 /δ ) /ϵ 2 , then with probability at least 1 -δ ,

Proof. Using Lemma 8 and Cauchy-Schwarz inequality, we have

̸

The statements follow directly from plugging in suitable values of d .

Lemma 8 (Almost orthonormal) . Let x 1 , x 2 , . . . , x n i.i.d. ∼ N d ( 0 d , 1 d I d ) . For any ϵ, δ ∈ (0 , 1) , when d ≥ 16 log(2 n 2 /δ ) /ϵ 2 , it holds that with probability at least 1 -δ ,

̸

Proof. Fix i = j ∈ [ n ] . Notice

Using Lemma 9, we have that with probability at least 1 -δ/n 2 ,

Furthermore, fix i ∈ [ n ] , with probability at least 1 -δ/n 2 ,

|⟨

x

i

, x

i

⟩ -

1

|≤

ϵ.

The statement then follows from union bound over i, j ∈ [ n ] .

Lemma 9 (Almost normal for a fixed vector) . For a d -dimensional random vector x ∼ N d ( 0 d , 1 d I d ) and any v ∈ (0 , 1 / 2) ,

In particular, when d ≥ 16 log(1 / (2 δ )) /v 2 , we have

Proof. Let x = ( x 1 , . . . , x d ) . By Lemma 11 (letting x = v 2 d/ 16 ),

The second inequality follows from simple arithmetics.

̸

B.4 Useful results

Lemma 10 (ODE bound) . Let c 1 , c 2 , c 3 > 0 . Suppose the function f 1 , f 2 : R + → R satisfies f 1 (0) > 0 , f 2 (0) > 0 and

Then

Proof. The conditions imply that

Rearranging the above inequalities, one can obtain the desired results.

Lemma 11 ( χ 2 -concentration bound, Lemma 1 of [64]) . Let g 1 , . . . , g t be i.i.d. N (0 , 1) random variables. Then for any x ≥ 0 ,

C Missing Proofs of Section 4

In this section, we show missing proofs in Section 4.

C.1 Proofs of Section 4.1

C.1.1 Proofs of Proposition 4.1 and Proposition 4.2

We first show the proofs of Proposition 4.1 and Proposition 4.2, respectively.

Proof of Proposition 4.1. Actually, for any three tokens x 1 , x 2 , x 3 , it holds that

since Y (0) = 0

It follows that and

Proof of Proposition 4.2. Note that the input sequence length T = 2 . By (1),

where b 2 = [ b 12 ] and b 12 = 1 . Also, X = [ x 1 ] ⊤ is a one-hot row vector. Therefore, LN ( X ⊤ b 2 ) = LN ( x 1 b 12 ) = LN ( x 1 ) = x 1 , and thus

where Y ( t ) i,j is the entry of the matrix Y ( t ) at row i and column j .

C.1.2 Proof of Lemma 2

Proof of Lemma 2. We first calculate the gradient of Y when the current batch is a sequence ( x 1 , x 2 , x 3 ) . Note that the input sequence length T = 2 and by the proof of Proposition 4.2, we have LN ( X ⊤ b T ) = LN ( X ⊤ b 2 ) = x 1 . By Lemma 1, we have

where α = [ α 1 , α 2 , . . . , α M ] ⊤ ∈ R M with α = exp ( Y ⊤ x 1 ) / 1 ⊤ exp ( Y ⊤ x 1 ) . Note that x 1 ( x 3 -α ) ⊤ is a matrix with only x 1 -th row non-zero since x 1 is one-hot. Therefore, the update of each row of Y are independent and only x 1 -th row of Y gets updated at the current time step.

Now we consider any fixed x 1 ∈ V . Let t x 1 ,i be the time step that the x 1 -th row of Y gets updated (i.e., the first token of the training data is x 1 for the current batch) for the i -th time and let t x 1 , 0 = 0 for notation convenience, then

For convenience, we denote y ( i ) = Y ( t x 1 ,i ) ⊤ x 1 , and thus

where y (0) = 0 and α ( i -1) = exp( y ( i -1)) / 1 ⊤ exp( y ( i -1)) . Note that for a fixed x 1 , x 3 is also fixed by our construction of the dataset. By Lemma 5 of [14], we can obtain that

where ξ x 3 = M M -1 ( x 3 -1 M 1 ) and h ∗ ( i ) can be derived recursively as

with h ∗ (0) = 0 . Combining Lemma 7 and 9 in [14], we have

Note that the update of each row of Y are independent, the training set has size N , and the batch size is 1, we have

where the training data at time step t is ( x 1 , x 2 , x 3 ) . Combining (11), we can obtain that

and

̸

On the other hand, for any sequence ( x 1 , x 2 , x 3 ) in the test set, since the x 1 -th row of Y has never been updated, we have

C.1.3 Proof of Theorem 3

Proof. We first consider training sequence ( x 1 , x 2 , x 3 ) at time t . By Proposition 4.2, we have

for some constant c > 0 . Therefore,

where the last inequality holds since

Finally, for any sequence ( x 1 , x 2 , x 3 ) ∈ D test, since Y ( t ) x 1 ,x = 0 , ∀ x ∈ [ M ] according to Lemma 2, we have

which completes the proof.

C.2 Additional details and proof of Section 4.2

Datasets. Let N train > 0 , N test > 0 be two positive integers and let N total = N train + N test . Let A i , B i , C i ∈ V , ∀ i ∈ [ N total ] be 3 N total distinct tokens. Let → , ❀ ∈ V = [ M ] be two additional different tokens that represent 'direct implication' and 'indirect implication' respectively. Specifically, we have A i → B i , B i → C i and A i ❀ C i for all i ∈ [ N total ] . For notation convenience, we define the following two index sets

The training set D train consists of all A i → B i , B i → C i and A i ❀ C i for i ∈ I train. In addition, D train contains A i → B i and B i → C i for i ∈ I test. For convenience, we let N = |D train | to be the size of the training set. The test set D test consists of A i ❀ C i for i ∈ I test. Under our construction of the dataset, the LLM will learn the relationship between A i , B i and C i for i ∈ I train in both direct and indirect implication, and learn the relationship between A i , B i and C i for i ∈ I test only in direct implication and will be tested for indirect implication.

Similar to the reversal curse in Section 4.1, we aim to prove through the training dynamics of onelayer transformers that the test probability remains negligible during training. In particular, we are interested in

We also use p θ ( B i | A i → ) , p θ ( C i | B i → ) and p θ ( C i | A i ❀ ) to more compactly represent p θ ( x 3 = B i | x 1 = A i , x 2 = → ) , p θ ( x 3 = C i | x 1 = B i , x 2 = → ) and p θ ( x 3 = C i | x 1 = A i , x 2 = ❀ ) , respectively.

The following theorem shows the importance of the chain-of-thought method:

̸

Theorem 7 (Importance of chain-of-thought, formal statement of Theorem 4) . Assume we run SGD with batch size 1, and assume M ≫ 100 and 1 M 0 . 99 ≪ η Y < 1 . Let t ≳ N ln M η Y denote the time step which also satisfies ln t ≳ ln( NM/η Y ) . For any test index i ∈ I test, we have

for some constant c > 0 and and by Lemma 2, we have

̸

for some constant c > 0 . Therefore, using the same proof as Theorem 3, we have

Additionally, according to the proof of Lemma 2, Y ( t ) A i ,x has the same value across all x = B i , which implies

Similarly, applying Lemma 2 to Y ( t ) B i , we can obtain that

C.3 Analysis of the reversal curse for the four-token sequences

In this section, we analyze the reversal curse where data points are four-token sentences ' AR 1 R 2 B ' or ' BR 1 R 2 A '. For each sentence, A and B are two distinct tokens that represent two entities, and R 1 , R 2 are two special tokens jointly representing a relationship that is inverse to itself (e.g., R 1 R 2 represents 'is the friend of', then AR 1 R 2 B means ' A is the friend of B ' and BR 1 R 2 A means ' B is the friend of A ' ).

Datasets. Let N train > 0 and N test > 0 and denote N total = N train + N test . Let A i , B i ∈ V , ∀ i ∈ [ N total ] be K ≜ 2 N total distinct tokens representing distinct entities. WLOG, we assume A i , B i ∈ [ K ] , ∀ i ∈ [ N total ] . Let R 1 , R 2 ∈ V be two additional different tokens that jointly represent a relationship that is inverse to itself. Specifically, we have A i R 1 R 2 B i and B i R 1 R 2 A i for all i ∈ [ N total ] . For notation convenience, we define the following two index sets

The training set D train consists of all A i R 1 R 2 B i and B i R 1 R 2 A i for i ∈ I train. In addition, D train contains A i R 1 R 2 B i for i ∈ I test. For convenience, we let N = |D train | to be the size of the training set. The test set D test consists of B i R 1 R 2 A i for i ∈ I test. Under our construction of the dataset, the LLM will learn the relationship between A i and B i for i ∈ I train in both directions to deduce that the relationship ' R 1 R 2 ' is reverse to itself, and learn the relationship between A i and B i for i ∈ I test in one direction and will be tested for the other. WLOG, we assume for each training sequence ( x 1 , R 1 , R 2 , x 4 ) ∈ D train , x 4 ∈ [ N ] .

Proof. Recall that by Proposition 4.2, we have

̸

Now we assume that the learning rate η Y ≫ η Z and therefore can treat X ⊤ b T as fixed when analyzing the dynamics of Y . For any sequence ( x 1 , x 2 = R 1 , x 3 = R 2 , x 4 = n ) in training or test dataset with T = 3 , we define f n = LN ( X ⊤ b T ) . Then the gradient of Y in (3) becomes

Note that for three-token sequences, f n is a one-hot vector, and thus ˙ Y has only one non-zero row, and each row of Y can be analyzed independently. For the four-token sequences, we use the same reparameterization strategy as in [14], where we denote W = [ w 1 , w 2 , . . . , w K ] ⊤ ≜ F ⊤ Y ∈ R K × M with F = [ f 1 , . . . , f K ] ∈ R M × K .

Note that the parameter W can be viewed as a combination of Y and Z where Z is fixed. The following lemma shows the dynamics of W , assuming we are performing gradient updates on W instead of Y .

Lemma 12 (Dynamics of W ) . Assume we perform gradient updates directly on W instead of Y with learning rate η Y and batch size 1, and assume M ≫ 100 and 1 M 0 . 99 ≪ η Y < 1 . Let t ≳ N ln M η Y and let W ( t ) i denote the i -th row of W ( t ) and W ( t ) ij denote the ( i, j ) -th entry of W ( t ) . Then for training sequence ( x 1 , R 1 , R 2 , x 4 ) ∈ D train at time t , we have

̸

and for any test sequence ( x 1 , R 1 , R 2 , x 4 ) ∈ D test, we have W ( t ) x 4 ,x = 0 , ∀ x ∈ [ M ] .

Proof. According to Lemma 3 of [14], for training sequence ( x 1 , R 1 , R 2 , x 4 ) ∈ D train at time t , only the x 4 -th row of W will be updated and the gradient

where α ( t ) = exp( w x 4 ( t )) / ( 1 ⊤ exp( w x 4 ( t ))) . Therefore, the dynamics of W is nearly identical to the dynamics of Y in Lemma 2, and we can use the proof of Lemma 2 to conclude the results.

With the dynamics of W , we can obtain the following result:

Proposition C.1 (Reversal curse for the four-token sequences) . Assume we perform gradient updates directly on W instead of Y with learning rate η Y and batch size 1, and assume M ≫ 100 and 1 M 0 . 99 ≪ η Y < 1 . Let t ≳ N ln M η Y denote the time step which also satisfies ln t ≳ ln( NM/η Y ) . For training sequence ( x 1 , R 1 , R 2 , x 4 ) ∈ D train at time t , we have

for some constant c > 0 , and for any test sequence ( x 1 , R 1 , R 2 , x 4 ) ∈ D test that is not included in the training set D train, we have

Proof. For any sequence ( x 1 , x 2 = R 1 , x 3 = R 2 , x 4 ) where T = 3 ,

The above next token probability formulation is almost identical to Proposition 4.2 after replacing Y with W . Combining the dynamics of W as shown in Lemma 12, we can use the proof of Theorem 3 to conclude the result.

Finally, we discuss the role of Z in the four-token-sequence settings. Note that Proposition C.1 assumes gradient update on W instead of Y . While we are not able to perform gradient updates on W directly, it is equivalent to modifying the gradient of Y to be

according to Lemma 3 of [14], where the next token x T +1 = n , E ′ = ( I + E ) -1 -I , E = F ⊤ F -I , and F = [ f 1 , . . . , f N ] ∈ R M × N which only contains the next token that appears in the training set. Compared to the original gradient of Y

one can obtain that the modification of the gradient of Y is small if λ 1 ( E ) is small.

Note that E ii = ∥ f i ∥ 2 2 -1 = 0 and E ij = f ⊤ i f j . Also, for any sequence ( x 1 = n ′ , x 2 = R 1 , x 3 = R 2 , x 4 = n ) in training set,

Note that R 1 is a common token that appears in each training sentence, and n ′ is a distinct token that only appears in one training sentence as a contextual token. By Theorem 2 of [14], under certain technical assumptions, ˙ Z R 2 ,n ′ > 0 and ˙ Z R 1 ,n < 0 and one can expect Z R 2 ,n ′ -Z R 1 ,n to be sufficiently large after sufficient time of training. Therefore,

with ˜ b 23 close to 0. Consider a simple case where for each n ∈ [ N ] , f n = √ 1 -c 2 e n ′ + c e R 1 for c sufficiently small. Then

Therefore, one can calculate that λ 1 ( E ) = c 2 ( N -1) . When c ≪ 1 √ N , we have λ 1 ( E ) ≪ 1 , and thus the gradient update of W and gradient update of Y are almost the same.

D Experiments for chain-of-thought

In this section, we conduct experiments for COT on multi-layer transformers to validate theoretical results in Section 4.2.

Dataset construction. Similar to Section 5, we randomly sample three disjoint sets of entities A , B , C ⊂ V , and reverse two additional tokens for → and ❀ , respectively. Next, we specify a bijection from A to B , and a bijection from B to C randomly. For each A i ∈ A and its corresponding B i ∈ B and C i ∈ C , we can obtain a triple of sequences ( A i → B i , B i → C i , A i ❀ C i ) , and split the set of all triples into training triples and validation triples. All three sequences of a training triple will be added to the training set, while for a validation triple, we add A i → B i and B i → C i to the training set and add A i ❀ C i to the validation set. Therefore, the model will learn both direct and indirect implications for the training triples and only learn the direct implications for each validation triple while being tested on the indirect implication.

Results. Figure 3 shows the experiment results for COT using the same model architecture and configurations as in Figure 1 (the training set size is 540, and the validation set size is 60 resulting from 140 training triples and 60 validation triples), which is consistent with Theorem 7. One can refer to Appendix E.3 for additional experiments with various model configurations and vocabulary sizes. We also empirically validate the intransitivity of model weights (i.e., logits) for multi-layer transformers in Figure 4, which shows that for a validation triple ( A i , B i , C i ) of which only the direct implication ' A i → B i ' and ' B i → C i ' appears in the training set, although the weights from A i to B i and from B i to C i are trained large as indicated by the diagonals of the first two bottom matrices, the weights from A i to C i gets hardly trained as indicated by the diagonals of the last matrix. We also emphasize that another reason that COT is necessary is that all tokens A i , B i , and C i are different tokens with randomly initialized embedding and thus irrelevant. When these tokens are relevant and show specific patterns, the validation loss can also get better. See more details in Appendix E.4.

where

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

E Additional Experimental Results

In this section, we show additional experimental results for Section 5.

E.1 Details of model architectures and hyperparameters

For both the reversal curse and COT experiments, we used the GPT2 model architecture [63] 6 and trained the model with the AdamW optimizer for 3000 epochs of batch size 64. See Table 2 for a full list of hyperparameters. We also conducted experiments under various model configurations and vocabulary sizes to show that the results in Section 5 and appendix D are consistent under different settings. See Table 3 for a complete list of different configurations, where the default choices are boldened. For each curve in all figures, the results are averaged over 10 trials, and the error bar is calculated using standard deviation. We run each trial on an Nvidia A100 GPU and it typically takes 0.5-1.5 hours for each trial.

E.2 Additional experimental results for the reversal curse

In this section, we show additional experimental results for the reversal curse under different configurations, including different vocabulary sizes (Figure 5), different number of layers (Figure 6),

Logits of B given A (Train)

Logits of C given B (Train)

Logits of C given A (Train)

Logits of B given A (Val)

Logits of C given B (Val)

Logits of C given A (Val)

75

50

25

0

25

50

75

100

different positional encoding (Figure 7), different entity lengths (Figure 8) and whether token and positional embeddings are trainable or fixed (Figure 9). Our experimental results consistently show that the reversal curse happens under different settings.

We provide additional experimental results that (1) the reversal curse still happens even if the embedding dimension is much smaller than the vocabulary size in Appendix E.2.1; (2) the embeddings of different tokens are nearly orthogonal; (3) the reversal curse does not happen under the in-context learning settings.

E.2.1 The reversal curse under small embedding dimensions

Although for theoretical analysis in Section 3, we assumed the embedding dimension is polynomial in the vocabulary size, in practice, the embedding dimension only needs to be the order of logarithm of the vocabulary size. Figure 10 shows that for a much smaller embedding size, the reversal curse still happens.

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) Vocabulary size 20

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) Vocabulary size 50

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(c) Vocabulary size 200

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(d) Vocabulary size 2000

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) 12 layers

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) 48 layers

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) No positional encoding

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

70

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) Relative positional encoding

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

Negative Log Probability

Negative Log Probability over Epochs

Train

Train (word)

Val

Val (word)

(a) Entity length 2

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Train (word)

Val

Val (word)

(b) Entity length 3

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

E.2.2 Near orthogonal embeddings

Note that in Section 3, our analysis relies on the fact that embeddings of different tokens are nearly orthogonal, and in Section 4, the embeddings are effectively one-hot. In Figure 11, We show that in practice, even if the embedding dimension is much smaller than the vocabulary size, the near orthogonality condition still holds.

E.2.3 The reversal curse does not happen in ICL settings

We also emphasize that the reversal curse does not happen in ICL settings, which means if ' A → B ' is provided as part of the prompt, then the model is able to answer ' B ← A '. Figure 12 shows preliminary results of ICL. All the sentences in the dataset have the format of ' A i RB j ⇔ B j R -1 A i ', which is a seven-token sentence where R and R -1 is a pair of relationships inverse to each other, and ' ⇔ ' is another reserved token representing equivalence. There are ten different B j and n different A i where n = 100 for the left figure in Figure 12 and n = 200 for the right figure. For each A i , we construct ten sentences using different B j , and we randomly chose three of them to be included in the validation set and seven other sentences in the training set. The result of Figure 12 shows that the reversal curse does not happen during the ICL setting.

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

100

Negative Log Probability

Train

Val

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

100

Negative Log Probability

Train

Val

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

100

Negative Log Probability

Train

Val

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

1.00

E.3 Additional experimental results for chain-of-thought

In this section, we show additional experimental results for COT under different configurations, including different vocabulary sizes (Figure 13), different number of layers (Figure 14), different positional encoding (Figure 15), different entity lengths (Figure 16) and whether token and positional embeddings are trainable or fixed (Figure 17). Note that our experimental results consistently show the necessity of COT under different settings.

E.4 Chain-of-thought with relevant tokens

In Appendix D, we briefly mentioned that the irrelevance of different entity tokens is one of the reasons that COT is necessary. Now, we show that if the entity tokens are correlated and show specific patterns, it is possible for a model to deduce indirect implications automatically.

Instead of using single tokens A i , B i , C i to represent each entity, now we use two tokens Ai , Bi , Ci to represent entities, where A , B and C are three common tokens shared by each triple, and token i are distinct for each triple. Figure 18 shows that for the above version of COT where tokens in the same chain are correlated, the model is able to 'deduce' Ai ❀ Ci after training on Ai → Bi , Bi → Ci and other training samples.

0

200

400

600

800

1000

Number of Epochs

2

0

2

4

6

8

10

12

14

Negative Log Probability

Vocab size: 400

Train

Val

0

200

400

600

800

1000

Number of Epochs

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Negative Log Probability

Vocab size: 800

Train

Val

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) Vocabulary size 20

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) Vocabulary size 50

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

60

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(c) Vocabulary size 200

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(d) Vocabulary size 2000

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) 12 layers

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) 48 layers

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) No positional encoding

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

20

40

60

80

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) Relative positional encoding

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

Negative Log Probability

Negative Log Probability over Epochs

Train

Train (word)

Val

Val (word)

(a) Entity length 2

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

Negative Log Probability

Negative Log Probability over Epochs

Train

Train (word)

Val

Val (word)

(b) Entity length 3

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

40

50

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

0

500

1000

1500

2000

2500

3000

Number of Epochs

0

10

20

30

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(a) Training set size 160

0

500

1000

1500

2000

2500

3000

Number of Epochs

5

0

5

10

15

20

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(b) Training set size 250

0

500

1000

1500

2000

2500

3000

Number of Epochs

5

0

5

10

15

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(c) Training set size 400

0

500

1000

1500

2000

2500

3000

Number of Epochs

5

0

5

10

15

20

Negative Log Probability

Negative Log Probability over Epochs

Train

Val

(d) Training set size 1600

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: See abstract and Section 1.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See Sections 1 and 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: See Sections 3 and 4 and appendices B and C. Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: See Section 5 and appendices D and E.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our code is available at https://github.com/marlo-z/reversal_ curse_analysis/ .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See Section 5 and appendices D and E.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conforms, in every respect, with the NeurIPS code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: See Section 1 and Section 6.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Our code is available at https://github.com/marlo-z/reversal_ curse_analysis/ .

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a theoretical analysis of the recently discovered reversal curse phenomenon of LLMs [1], whereby a model trained or fine-tuned on the sentence \"A is B\" often fails to generalize to the reverse direction \"B is A\". The authors study the gradient dynamics of two auto-regressive models, a bilinear model and a one-layer transformer, when optimizing for cross-entropy loss of next-token prediction.\n\nThe models are trained on sentences containing both implications $A\\to B$, $B\\leftarrow A$ for train tokens $A,B$ but only one implication $A'\\to B'$ for test tokens $A',B'$, and tested on the reverse implication $B'\\leftarrow A'$. The paper reveals a strict separation between training and test error, showing that train error can become arbitrarily small while test error is still close to random guessing. The analysis exploits the asymmetry of the reparametrized model weights, where the weights from $A$ to $B$ do not directly affect the weights from $B$ to $A$. An extension to analyzing chain-of-thought and experiments on multi-layer transformers are also given.\n\n[1] Berglund et al. The Reversal Curse: LLMs Trained on \"A is B\" Fail to Learn \"B is A\". ICLR 2024.\n\nSTRENGTHS:\nThe paper provides a timely study on a fundamental limitation of the reasoning capabilities of LLMs. The analysis is able to incorporate the next-token prediction mechanism with CE loss with a detailed breakdown of the gradient dynamics and sheds light on various aspects of the reversal curse empirically demonstrated in [1]. Moreover, the modification in Section 4.2 for CoT is very interesting and provides a new approach to the theory of CoT, which so far has been mostly limited to complexity-based expressivity analyses.\n\nWEAKNESSES:\n* **Meta-learning:** I am concerned that the paper's approach precludes discussion of the meta-learning capability of Transformers. The surprising aspect of the reversal curse is not just that the LLM fails to learn the reverse implication for a single test instance $A'\\to B'$, but fails to *generalize* from a large training set containing both $A_i\\to B_i$ and $B_i\\leftarrow A_i$. While the problem setup follows this setting, the proof depends on the near-orthonormality of the embeddings, which in turn requires the dimensionality to be much higher than the number of relevant tokens. This assumption feels very strong: in effect, it forces the model to disassociate different tokens and reduces learning relationships to encoding specific parameters, e.g. $\\Theta_{i,j}$ similar to memorizing a lookup table. The authors then rely on the asymmetry or intransitivity of these 'lookup' parameters to draw their conclusions. However, I think this reduction may prevent a meaningful analysis of generalization behavior in the first place, and does not explain why the reversal curse happens in practical LLMs (where the training dataset is nearly infinite and the model has had ample time to generalize from a continuum of examples).\n* **Asymmetry:** Aside from the meta-learning aspect and just looking at the asymmetric encoding perspective, the idea that LLMs store factual associations as directed representations has been empirically proposed multiple times before, see e.g. the references in Sections 3-4 of [1]. The presented analysis, while theoretically very solid and well-detailed, feels to me more like a verification of this 'obvious explanation' in a toy setup rather than providing new insights into the phenomenon, and hence may be less impactful. The authors should include more in-depth discussion on previous empirical studies (the Reversal Curse paragraph in Section 1.1) to better contextualize their contributions rather than just pointing out the lack of theoretical works.\n* **In-context learning:** With the above points in mind, it would be interesting to extend the analysis to e.g. theoretically compare the settings of pre-training (or fine-tuning) and in-context learning, given the different training approaches in [1,2] and the observation that ICL manages to overcome the reversal curse [1].\n\n[2] Grosse et al. Studying Large Language Model Generalization with Influence Functions. 2023.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nSome directions for future work are suggested in the last section. See also weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores the difficulty auto-regressive LLMs face with simple logical reasoning tasks, such as reversal curse. Through theoretical analysis and experiments, the study examines the training dynamics of gradient descent in bilinear models and one-layer transformers, suggesting that the asymmetry in model weights during training—where learning the relationship from A to B does not ensure learning from B to A—is a significant factor behind the reversal curse. This issue arises from the optimization process under the commonly used cross-entropy loss, which does not enforce symmetry in learning these relationships. The paper also discusses the implications for other logical reasoning tasks, like chain-of-thought reasoning, and validates its theoretical findings with experiments on multi-layer transformers.\n\nSTRENGTHS:\n1. The exploration of the \"reversal curse\" is a timely and relevant topic within the field. This paper stands out as a pioneering theoretical analysis on this subject, which is commendable. \n\n2. From a technical standpoint, this paper is well-founded. The depth of the analysis is impressive and may hold value not only within the context of this study but also as a reference point for future research. The methodology and the thoroughness of the analysis contribute significantly to the paper's technical soundness.\n\n3. The investigation into the problem, specifically highlighting that the issue stems from training dynamics rather than expressiveness, is particularly insightful. This perspective is refreshing and challenges conventional approaches, potentially inspiring a new stream of research in this direction.\n\nWEAKNESSES:\n1. Concerning the problem formulation, I found myself questioning its identifiability (Question 1).\n\n2. Regarding the assumptions made and the proofs provided, I have a few questions that arose during my review (Question 2, 3, 4).\n\nQUESTIONS:\n1. In the context of this paper's system model, two rules merit consideration:\n- RULE 1: If A leads to B (A->B), then it is also true that B leads to A (B->A).\n- RULE 2: If A leads to B (A->B) and A belongs to a specific set $\\mathcal{A}$, then B leads to A (B->A).\n\nThe LLMs learn RULE 2. Given that RULE 1 and RULE 2 are indistinguishable under the training dataset, learning RULE 2 instead of RULE 1 is acceptable and not a reverse curse. For example, considering -> as the relationship \"mother\" and <- as the relationship \"son\", A->B implies B->A under the condition that B is male.\n\n2. This paper assumes that tokens follow a Gaussian distribution and that only one token is used as the input for attention. I understand that this is used for landscape and generalization analysis. However, these conditions diverge significantly from practical scenarios. Could the authors discuss the applicability of the results in settings where the distribution of tokens is more realistic, and sequences contain more than one token?\n\n3. The explanation for the reversal curse provided between lines 173-179, which employs one-hot vectors for x and y, seems more aligned with current LLM inputs.\n- For Theorem 1 and 2, could the proofs be simplified while still achieving similar results if tokens are selected from one-hot vectors?\n- Is it possible to apply this intuition to a one-layer Transformer model as well?\n\n4. The paper operates under the assumption that $T$ is limited to 2, and the model considers merely a single input token. This premise leads to a scenario where the attention mechanism is simplified, rendering the problem convex. Does this approach risk oversimplifying the complexities inherent in the model's dynamics? Moreover, is it feasible to extend the analysis to accommodate scenarios involving a greater number of input tokens, such as 2 tokens, to explore how the model behaves under a non-convex setting?\n\n5. Regarding the experimental setup employing a 24-layer 12-head attention Transformer, this configuration seems considerably complex for testing the theoretical findings.\n- Could you provide the rationale behind choosing such a large Transformer model?\n- I suggest including experiments with a 1-layer Transformer to verify if the observed training dynamics align with the theoretical predictions.\n- I am particularly interested in the case where the embedding dimension is (much) smaller than the vocabulary size, specifically for a single-layer Transformer. In this scenario, the Transformer may not have sufficient memory to store all the necessary information. How would this affect the empirical results?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work analyzes the optimization trajectory of (a) a bilinear model trained on sequences of random Gaussian vectors, and (b) a simplified one-layer transformer on sequences of three tokens. It studies how well the trained model generalizes to sequences in reverse order. This study is motivated by an empirical phenomenon observed in large language models, called the \"reversal curse.\" Furthermore, the paper considers some variations of the main simplified settings, notably a setup that mimics the so-called chain of thought technique. Finally, experiments are presented that demonstrate the validity of the theoretical results.\n\nSTRENGTHS:\nThis work studies a well-documented empirical phenomenon of autoregressive language models, namely their inability to operate as databases (since the order of a query matters) - a phenomenon called \"the reversal curse.\" A better theoretical understanding of the subject could have a substantial impact. The current work appears to be the first to study it theoretically from an optimization perspective. The paper has certain merits: in particular, it offers an interesting theoretical analysis of training dynamics of simplified models on synthetic data. The proven theorems/propositions seem correct to me (I did not manage to verify all the steps - see questions on how the presentation of some results could be improved), and the techniques, whilst appearing to be standard, are interesting.\n\nWEAKNESSES:\nUnfortunately, I identified the following issues with the paper in its current form:\n\n* Unclear motivation of the theoretical framework: As far as I understand, the so-called \"reversal curse\" is interesting for cases where we expect a model to have deduced the transitivity of a relation before evaluating. For instance, this could have happened from pretraining on some large corpus and then fine-tuning/testing on Question-Answer data, like the experiments considered in [1]. However, the setup in this work only considers fine-tuning on well-structured, synthetic datasets. As a result, a constant question that I had in mind while reading the theoretical part of this paper was \"why should I expect the contrary of what is proven?\". Is there any expectation from a model to learn a reverse relation solely from training on only one direction? The answer seems obviously to be \"no\", in the same way that we do not expect the probabilities of other unseen sentences to increase. I would appreciate it if the authors could elaborate on that in their rebuttal and perhaps add a discussion of it in the paper. Ideally, an analysis of pretrained models would have been appropriate for theoretically studying this phenomenon.\n\n* Imprecise writing: The paper is not very well written in places, and it is hard to follow with terms that are not being defined before usage. For example, lines 35-51 try to motivate the paper by arguing that a specific bias in the architecture weights could potentially resolve the \"reversal curse,\" but they quickly dismiss this proposal on the (correct!) basis that this would bias the model in learning nonsense inputs. Thus, I do not see why this argument is worth mentioning. In the same passage, the authors use the term \"symmetric loss\" without defining it before. It becomes clear to someone what the authors mean after reading the paper and why someone (naively) might consider this as an option, but a priori, it is not clear and does not help with the understanding of the paper. See also questions for specific questions/suggestions on other parts of the paper (and in particular in the technical part of the proofs).\n\n1. The Reversal Curse: LLMs trained on \"A is B\" fail to learn \"B is A\". Lukas Berglund, Meg Tong, Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, Owain Evans.\n\nQUESTIONS:\nQuestions/suggestions:\n\n* Define terms used in the text: examples include the term \"symmetric loss\" in line 47, the term \"reparameterization\" in line 59.\n* Consistent notation: In Section 2, line 130, you define input sequences of tokens $x = (x_1, \\ldots, x_T)$, but in the next page you redefine (or specify) input sequences as $(x_i, y_i)$ or $(y_i, x_i)$ (line 144). While it is not a major issue, being consistent would help readers here.\n* Can you please clarify the sentence in line 297-298 \"Note that ... will not attend to itself\". It is not clear to me how the query token does not attend to itself in the adopted framework.\n* A criticism of Theorems 1, 2 could be that they rely on the large dimension of the matrix $\\Theta$ to basically make it operate as an associative memory (if I understand correctly). Could you please comment on that and perhaps discuss this in a revised version? How does the dimensionality of the model affect the conclusions? I understand that it needs to be large enough, but should we expect the same   conclusions to hold for models of smaller dimensionality?\n* The proofs of the Theorems and Propositions appear in a convoluted order in the Appendix. For instance, the proof of Theorem 1 (Theorem 5 in the Appendix) starts almost immediately by invoking Lemma 3 and Lemma 4 (with no commentary), which in turn immediately invokes Lemma 5. Thus, it is difficult to follow and verify the logic of the conclusions. Please consider presenting the proofs with more details and/or change the order of presentation.\n* line 625: It is not clear why you set $u = 0.2$.\n* line 626: Can you please add some explanation on the first inequality? You probably use the fact that the matrix $\\Theta$ is Gaussian, but there exists no justification for that. As far as I understand this the only element of the proof that really makes use of the properties of $\\Theta$ so there is value in being explicit about your proof arguments.\n* line 641, 'Using Lemma 9': Can you please elaborate on that? The mentioned Lemma shows a result for the norm of a Gaussian vector, not for the inner product of two different vectors.\n* Lemma 10 in page 22: A little bit more care should have been taken in the proof of this Lemma. In particular, below line 653, an inequality for the inverse function is being stated. However, there is no guarantee that the function cannot attain the 0 value at some point $t$ (in fact, $f(0) > 0$ and $f'(0) < 0$).\n* Theorem 4: I do not understand the conclusion of this theorem. This theorem asserts that a certain conditional probability is no better than random chance, but it does not demonstrate a **neccessity** of chain-of-thought. A statement about the necessity of chain-of-thought should prove that no matter which other inference technique we use, if it is not chain-of-thought then it will not succeed. This is very different from the proved theorem. Can you please address this in your rebuttal?\n* Suggestion on experiments: An interesting observation of [1] is that the \"reversal curse\" is remedied if the reverse sentence appears in context. Do you believe that a similar observation could be made in your synthetic setup? If the answer is affirmative, then this would actually alleviate some of the current concerns with the validity of the framework in explaining the \"reversal curse\".\n* Figure 15 is interesting, the authors could perhaps consider add some discussion of it in the main paper.\n\nA few typos/grammatical errors I caught:\n\n* line 87: \"phenomenonin\" -> \"phenomenon in\"\n* line 88: \"Several paper studies\".\n* line 98: \"There are rich literatures\"\n* line 117: \"We use $poly(x_1, \\ldots, x_n)$ to denote the polynomial of $x_1, \\ldots, x_n$.\" -> what polynomial? you mean **a** polynomial function.\n\n\n1. The Reversal Curse: LLMs trained on \"A is B\" fail to learn \"B is A\". Lukas Berglund, Meg Tong, Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, Owain Evans.\n\nLIMITATIONS:\nPlease consider being more elaborate on the NeurIPS paper checklist and your justifications of your answers (the paper currently mostly points to Sections of the paper, without any further justification).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper works toward a theoretical explanation of why the \"Reversal Curse\" occurs in language models. This \"curse\" refers to the empirical phenomenon that even when trained on certain relationships between entities (\"Sue is an aunt of Alice\"), the models fail at reversing those relationships in a semantically equivalent way (\"Alice is a niece of Sue\").\nA priori, we might expect that the training dynamics of large language models on large amounts of data would encourage a model that places a high probability on the forward relationship to do the same on the reverse one. But empirical results show this is not the case, and this paper attempts to break down the training dynamics to give a theoretical explanation why.\n\n- Section 3 shows that in a synthetic setting with a bilinear model, the \"reversal loss\" (predicting B<--A when the training data only contains A-->B for that example) is lower-bounded by a function of the training loss. For certain parameter regimes, this lower bound is 1, indicating that the reversal loss does not decrease due to the training dynamics even when the forward relationship A-->B is present in the training data.\n\n- Section 4 extends this analysis to one-layer transformers, a model also studied in prior work in understanding transformer training dynamics. It also extends the analysis of the reversal curse to show why Chain-of-Thought reasoning is necessary.\n\nSTRENGTHS:\n- Analyzes the training dynamics of increasingly more complex transformer-like models to show why the reversal curse emerges.\n\n- The techniques used in the paper can also be used to analyze why language models sometimes can't distill chain-of-thought reasoning: even though they can directly learn A-->B and B-->C, they can't always directly predict A-->C.\n\n- On a technical level, this paper relaxes some assumptions from other work on training dynamics of the bilinear model in Section 3.\n\nWEAKNESSES:\n- If I understand correctly, the ultimate conclusion is that there's no real hope beyond reversal training, i.e., having both A-->B and B-->A in the training data. It would be stronger if the analysis revealed some kind of (partial?) solution to the reversal curse.\n\nQUESTIONS:\n- L42 \"Since it is intractable to manually hard-code the constraints to the model parameter, one can alternatively expect the model to learn the higher-level regularity by training samples under unconstrained optimization\"\n\n    On my first read, I didn't quite understand what this meant. Is the claim that we might learn B <-- A from A --> B samples even without constraints? Or are the authors trying to say that we could just learn from samples of both A --> B and B <-- A? The former point is unintuitive (how could this be true? Because of some implicit structure in model parameter space / inductive bias?) and the latter point is arguably clear. It's clear from further reading that the authors mean the former, but it was confusing on a first reading. Mentioning \"due to training dynamics\" would be helpful here.\n\n- What is $\\sigma$ in Theorem 2? Defined in Line 156? It's not clear because L157 also says that I could take the parameters to be some pretrained value satisfying certain conditions. Does Theorem 2 only apply to the first case?\n\n- What does Theorem 2 add beyond the interpretation of Theorem 1 in L161--L162?\n\n- In Section 4, why do we need to assume the query token doesn't attend to itself?\n\n- Is there any reason not to just present the reparameterized form starting in L204? This would make the exposition simpler, though I suppose the connection to the transformer might be less clear.\n\n- In Section 4, what happens if the embeddings are clustered in such a way that the set of entities {Bi} and entities {Ai} are disjoint and (say) linearly separable? Intuitively if we already know what kind of entity a new B is, that might help with reversal? As far as I can tell, this is not the case in Section 3, but does that change in S4?\n\n- L87 \"phenomenonin\"\n\nLIMITATIONS:\nThe conclusion includes a fair discussion of the main limitations of the work (it only considers two- or three-token sequences, and the deepest model is considers is a one-layer transformer).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}