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

Authors: hanlin  zhu
baihe  huang
shaolun  zhang
michael  jordan
u c berkeley
jiantao  jiao
yuandong  tian
meta  ai
stuart  russell
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:54:59

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/).

Main paper: 
 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 usef (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 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))E D [f ] = 1 n n i=1 f (x i ). See additional notations in Appendix A. Entities Forward Backward Direct Indirect Others A, B, C, A i , B i , C i → ← → ; R 1 , R 2v∈V exp(l θ (v|x)) , ∀y ∈ V. Also, each token v ∈ V has a corresponding (fixed or learnable) embedding vector u v ∈ R d .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. 
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.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 lossL(Θ) = 1 2n-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 satisfying1 2m < 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 haveL rev (Θ t )/L rev (Θ 0 ) ≥ (L(Θ t )/L(Θ 0 )) ϵ , ∀t ≥ 0.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 d2 , 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 2m < 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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 vectorx = 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 vectorx t [i] = e xt[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 x1 , . . . , 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 ũT = U ⊤ LN(X ⊤ b T )where b tT = exp(u ⊤ x T W Q W ⊤ K ux t / √ d) T -1 t ′ =1 exp(u ⊤ x T W Q W ⊤ K ux t ′ / √ d) , b T = [b 1T , . . . , b T -1,T ]⊤ contains attention scores (after softmax) that query token x T attend to each contextual token 3 , LN(x) = x/∥x∥ 2 is the ℓ 2normalization 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 ũ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 byp θ (x|x 1 , . . . , x T ) = exp(l θ (x|x 1 , . . . , x T )) x ′ ∈[M ] exp(l θ (x|x 1 , . . . , x T )) = exp(u ⊤ x W V ũT ) x ′ ∈[M ] exp(u ⊤ x ′ W V ũT ).We use the cross-entropy loss function to train the model over the whole training set D train , i.e., maxU,W K ,W Q ,W V J ≜ E Dtrain [log p θ (x T +1 |x 1 , . . . , x T )].Reparameterization. Similar to [14] , we defineZ = U W Q W ⊤ K U ⊤ / √ d ∈ R M ×M and Y = U W ⊤ V U ⊤ ∈ R M ×Mand are interested in their dynamics after reperameterization. Then, the attention score (after softmax) and next token probability becomeEQUATIONand the objective can be written as maxEQUATIONLet η 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 asEQUATIONwhereP ⊥ 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. Ẏ , Ż can be obtained by direct calculation. One can refer to the proof of Lemma 1 of [14] .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.EQUATIONtest > 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 2N 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 setsEQUATIONEQUATIONtest ]\I train , IEQUATIONtest ). 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 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 havep θ(0) (B i |A i →) = p θ(0) (A i |B i ←) = 1/M.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) x1,x ) / x ′ ∈[M ] exp (Y (t) x1,x ′ ), where Y (t) i,jis 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 and1 M 0.99 ≪ η Y < 1. Let t ≳ N ln M η Yand 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 haveY (t) x1,x3 ≳ ln (M η Y t/N ) , and Y (t) x1,x ≲ -ln (M η Y t/N ) /M, ∀x ̸ = x 3 ,and for any test sequence(x 1 , x 2 , x 3 ) ∈ D test , we have Y (t) x1,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) x1,x3 gets increased during training while Y (t) x3,x1 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 and1 M 0.99 ≪ η Y < 1. Let t ≳ N ln M η Ydenote the time step which also satisfies ln t ≳ ln(N M/η Y ). For training sequence (x 1 , x 2 , x 3 ) ∈ D train at time t, we havep θ(t) (x 3 |x 1 , x 2 ) ≥ 1 -(M -1)(M η Y t/N ) -c → 1, as t → ∞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.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 anyA 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 p θ(t) (B i |A i →) → 1, p θ(t) (C i |B i →) → 1, p θ(t) (C i |A i ;) ≤ 1/M, as t → ∞.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] .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. 
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. 3 ). The curves represent the (average) negative log probability of the model predicting the next token to be B i when the input is "A i →", or to be A i when the input is "B i ←". While the sentences in the training set can be learned nearly perfectly (as shown by the training curve where the next token probability converges to one), the model is not able to predict the correct next token in the validation set better than a uniformly random guess. Both curves are averaged over 10 random seeds. 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. 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. 
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.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. 
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] Justification: See Appendix E.1.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. 
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]Justification: See Appendix E.1.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). 
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? 
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). 
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] Justification: See Appendix E.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. 
: 
 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 B1 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 LLM Reasoning. 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. 
 
 Appendix: Let δ ij = 1 for i = j and δ ij = 0 for i ̸ = j. For a squared matrix A ∈ R d×d , its trace isTr(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 1{x i ̸ = 0} or ∥A∥ 0 = m i=1 n j=1 1{A ij ̸ = 0}L(Θ) = 1 2n -1 n i=1 -log p Θ (y i |x i ) + n i=2 -log p Θ (x i |y i ) L rev (Θ) = -log p Θ (x 1 |y 1 ).t : t ≥ 0 dΘ t dt = -∇L(Θ t ).whereΘ 0 ∼ N (0 ⊗2 , σ 2 • I ⊗2 ) or Θ 0 satisfies 1 2m < p Θ0 (y i |x i ), p Θ0 (x i |y i ) < 2 m for all i ∈ [n]. Suppose σ ≤ 1 100 ln(64m 2 /δ) and d ≥ 10 6 n 4 m 2 log 4 (2m) log(64m 2 n 2 /δ) ϵ 2 .With probability at least 1 -δ, we haveL rev (Θ t ) L rev (Θ 0 ) ≥ L(Θ t ) L(Θ 0 ) ϵ , ∀t ≥ 0. Proof. Let v = 400n 2 m 2 log(64m 2 n 2 /δ) d. By Lemma 3 and Lemma 4, with probability at least 1 -δ,L rev (Θ t ) ≥ L rev (Θ 0 ) • 1 + L(Θ 0 ) • t 8(2n -1) log 2 (2m) -8v(2n-1) log 2 (2m) . ≥ L rev (Θ 0 ) • L(Θ t ) L(Θ 0 ) 8v(2n-1) log 2 (2m)By definition of d, we have 8v(2n -1) log 2 (2m) ≤ ϵ. Notice that L(Θt) L(Θ0) < 1, thus .EQUATIONWith probability at least 1 -δ,L rev (Θ τ ) ≥ C.where τ denotes the first time such that L(Θ t ) ≤ c.Proof. By continuity, L(Θ τ ) = c. By Theorem 5, whenEQUATIONwith probability at least 1 -δ,L rev (Θ τ ) ≥ L rev (Θ 0 ) • L(Θ τ ) L(Θ 0 ) ϵ ≥ L rev (Θ 0 ) • c L(Θ 0 ) ϵ .Under this event, applying Lemma 6, we can obtainL rev (Θ τ ) ≥ log(m/2) • c log(2m) ϵ .To ensure that the right hand side is C, we set ϵ= log C log(m/2) log c log(2m). One may check that the definition of d satisfies Eq. ( 4). It follows thatL rev (Θ τ ) ≥ C.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(16n 2 /δ) and d ≥ 1600n 3 m 2 log(8m 2 n 2 /δ), with probability at least 1 -δ, we haveL(Θ t ) ≤ 1 t 8(2n-1) log 2 (2m) + 1 L(Θ0), ∀t ≥ 0.Furthermore,inf {p Θt (x i |y i ), p Θt (y i |x i ) : t ≥ 0, i ∈ [n]} > 1 2m .Proof. For convenience, we assume x 1 , . . . , x n = v 1 , . . . , v n and y 1 , . . . ,y n = v n+1 , . . . v 2n WLOG. Let ϵ = 400n 2 m 2 log(8m 2 n 2 /δ) d . Then ϵ ≤ 1 2 √ n . Define l i (Θ) = -log p Θ (y i |x i ) and l rev i (Θ) = -log p Θ (x i |y i ). 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, dL(Θ t ) dt = ∇L(Θ t ), dΘ t dt = -⟨∇L(Θ t ), ∇L(Θ t )⟩ = - 1 2n -1 n i=1 x i (y i -E pΘ t (•|xi) [y]) ⊤ + n i=2 y i (x i -E pΘ t (•|yi) [x]) ⊤ 2 = - 1 2n -1   n i=1 m j=1 α (t) i,j v i v ⊤ j + 2n i=n+2 m j=1 β (t) i,j v i v ⊤ j   2 .Similarly, we haveEQUATIONdl rev i (Θ t ) dt = ∇l rev i (Θ t ), dΘ t dt = -⟨∇l rev i (Θ t ), ∇L(Θ t )⟩ = -y i (x i -E pΘ t (•|yi) [x]) ⊤ , 1 2n -1 n i=1 x i (y i -E pΘ t (•|xi) [y]) ⊤ + n i=2 y i (x i -E pΘ t (•|yi) [x]) ⊤ = - m j=1 β (t) i+n,j v i+n v ⊤ j , 1 2n -1   n i=1 m j=1 α (t) i,j v i v ⊤ j + 2n i=n+2 m j=1 β (t) i,j v i v ⊤ j   .Applying Lemma 7, with probability at least 1 -δ/2, for any t ≥ 0 we haveEQUATIONandEQUATION)dl rev i (Θ t ) dt + 1 2n -1 m j=1 (β (t) i+n,j ) 2 ≤ ϵ • 1 2n -1   m j=1 (β (t) i+n,j ) 2   1/2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   1/2 .Furthermore, Lemma 6 implies that with probability at least 1 -δ/2EQUATIONThe following arguments are based on the event that the above inequalities hold.We first show thatinf {p Θt (x i |y i ), p Θt (y i |x i ) : t ≥ 0, i ∈ [n]} > 1 2m .We prove this by contradiction. Letτ = inf t ≥ 0 : ∃i ∈ [n], s.t. min{p Θt (y i |x i ), p Θt (x i |y i )} ≤ 1 2m .By Eq. ( 7), it is obvious that τ > 0. Assume without loss of generality that p Θτ (y 1 |x 1 ) ≤ 1 2m . It follows that there exists δ > 0 such that p Θt (y 1 |x 1 ) is a decreasing function in (τ -δ, τ ) and p Θt (y1 |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 ∈ (τ -δ, τ ), dl 1 (Θ t ) dt ≤ 1 2n -1   - m j=1 (α (t) 1,j ) 2 + ϵ •   m j=1 (α (t) 1,j ) 2   1/2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   1/2    ≤ 1 2n -1   - m j=1 (α (t) 1,j ) 2 + ϵ •   m j=1 (α (t) 1,j ) 2   1/2 2 n i=1 (α (t) i,i+n ) 2 + 2 2n i=n+2 (β (t) i,i-n ) 2 1/2    ≤ 1 2n -1   m j=1 (α (t) 1,j ) 2   1/2 •   -   m j=1 (α (t) 1,j ) 2   1/2 + ϵ • 4n(α (t) 1,n+1 ) 2 1/2    ≤ 0where 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 thatinf {p Θt (x i |y i ), p Θt (y i |x i ) : t ≥ 0, i ∈ [n]} > 1 2m.L(Θ t ) ≤ 1 t 8(2n-1) log 2 (2m) + 1 L(Θ0), ∀t ≥ 0.By Eq. ( 5),dL(Θ t ) dt ≤ - 1 -ϵ (2n -1) 2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   ≤ - 1 -ϵ (2n -1) 2 n i=1 (1 -p Θt (y i |x i )) 2 + n i=2 (1 -p Θt (x i |y i )) 2 ≤ - 1 -ϵ (2n -1) 3 n i=1 (1 -p Θt (y i |x i )) + n i=2 (1 -p Θt (x i |y i )) 2 ≤ - 1 -ϵ 8(2n -1) log 2 (2m) L(Θ t ) 2 ≤ - 1 8(2n -1) log 2 (2m) L(Θ t ) 2 ,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; the last inequality uses the fact that p Θt (x i |y i ), p Θt (y i |x i ) > 1 2m for any t ≥ 0, i ∈ [n] and the inequality 1 -x ≥ log x 2 log(1/(2m)) for x ∈ ( 1 2m , 1). By Lemma 10, we conclude that ∀t ≥ 0,L(Θ t ) ≤ 1 t 8(2n-1) log 2 (2m) + 1 L(Θ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(16n 2 /δ) and d ≥ 400n 2 m 2 log(8m 2 n 2 /δ)/ϵ 2 , with probability at least 1 -δ,L rev (Θ t ) ≥ L rev (Θ 0 ) • 1 + L(Θ 0 ) • t 8(2n -1) log 2 (2m) -8ϵ(2n-1) log 2 (2m).Proof. Similar to Lemma 3, we assume x 1 , . . . , x n = v 1 , . . . , v n and y 1 , . . . ,y n = v n+1 , . . . v 2n 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, dL rev (Θ t ) dt = ∇L rev (Θ t ), dΘ t dt = -y 1 (x 1 -E pΘ t (•|y1) [x]) ⊤ , 1 2n -1 n i=1 x i (y i -E pΘ t (•|xi) [y]) ⊤ + n i=2 y i (x i -E pΘ t (•|yi) [x]) ⊤ = - m j=1 β (t) n+1,j v n+1 v ⊤ j , 1 2n -1   n i=1 m j=1 α (t) i,j v i v ⊤ j + 2n i=n+2 m j=1 β (t) i,j v i v ⊤ j   .Applying Lemma 7, with probability at least 1 -δ/2, for any t ≥ 0 we haveEQUATIONanddL(Θ t ) dt + 1 (2n -1) 2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   ≤ϵ • 1 (2n -1) 2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   ,as well asdl i (Θ t ) dt + 1 2n -1 m j=1 (α (t) i,j ) 2 ≤ ϵ • 1 2n -1   m j=1 (α (t) i,j ) 2   1/2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   1/2 , dl rev i (Θ t ) dt + 1 2n -1 m j=1 (β (t) i+n,j ) 2 ≤ ϵ • 1 2n -1   m j=1 (β (t) i+n,j ) 2   1/2   n i=1 m j=1 (α (t) i,j ) 2 + 2n i=n+2 m j=1 (β (t) i,j ) 2   1/2 .Furthermore, Lemma 6 implies that with probability at least 1 -δ/2,EQUATIONThe following arguments are based on the event that the above inequalities hold.By Eq. ( 8),EQUATIONwhere 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.A t ≤ 2(1 -p Θt (x 1 |y 1 )) 2 ≤ 2L rev (Θ t ) 2 .Combining, we havedL rev (Θ t ) dt ≥ -ϵ • 1 2n -1 • 2L rev (Θ t ) • (2n -1) • 1 t 8(2n-1) log 2 (2m) + 1 L(Θ0) ≥ -8ϵ(2n -1) log 2 (2m) • L rev (Θ t ) • 1 t + 8(2n-1) log 2 (2m) L(Θ0).By Lemma 10, we conclude that ∀t ≥ 0,L rev (Θ t ) ≥ L rev (Θ 0 ) • 1 + L(Θ 0 ) • t 8(2n -1) log 2 (2m) -8ϵ(2n-1) log 2 (2m), This completes the proof.L(Θ) = 1 2n -1 n i=1 -log p Θ (y i |x i ) + n i=2 -log p Θ (x i |y i ) , L rev (Θ) = -log p Θ (x 1 |y 1 ).∇L(Θ) = - 1 2n -1 n i=1 x i (y i -E pΘ(•|xi) [y]) ⊤ + n i=2 y i (x i -E pΘ(•|yi) [x]) ⊤ , ∇L rev (Θ) = -y 1 (x 1 -E pΘ(•|y1) [x]) ⊤ . Proof. We have ∇(-log p Θ (y|x)) = - ∇p Θ (y|x) p Θ (y|x) = - 1 p Θ (y|x) • xy ⊤ exp(x ⊤ Θy) y∈V exp(x ⊤ Θy) -exp(x ⊤ Θy) y∈V xy ⊤ exp(x ⊤ Θy) y∈V exp(x ⊤ Θy) 2 = - 1 p Θ (y|x)   p Θ (y|x)xy ⊤ -p Θ (y|x) y∈V p Θ (y|x)xy ⊤   = -x   y - y∈V p Θ (y|x)y   ⊤ = -x y -E y∼pΘ(•|x) [y] ⊤ .The statements follow immediately.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 p Θ (x j |x i ) = exp(l Θ (x j |x i )) n k=1 exp(l Θ (x k |x i )), where l Θ (xj |x i ) = x ⊤ i Θx j .Then when σ 2 ≤ 1 100 ln(4n 2 /δ) and d ≥ 400 log(1/(2δn 2 ))/ϵ 2 , with probability at least 1 -δ,|p Θ (x j |x i ) -1/n|≤ 1 2n , ∀i, j ∈ [n].Proof. Let v = 0.2. By Lemma 8, with probability at least 1 -δ, we have|⟨x i , x j ⟩ -δ ij |≤ v, ∀i, j ∈ [n].Conditioned on the above high-probability event, we can further obtain that for any j ∈ [n]EQUATIONp Θ (x j |x i ) = exp(l Θ (x j |x i )) n k=1 exp(l Θ (x k |x i )) ≥ exp(-v) n k=1 exp(v) = exp(-2v) n , It follows that 1 2n < p Θ (x j |x i ) < 3 2n =⇒ |p Θ (x j |x i ) -1/n|< 1 2n .which completes the proof.Lemma 7 (ℓ 1 -subspace embedding of Gaussian second-order tensors). Let z 1 , . . . , z n be independently sampled fromN 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(2n 2 /δ)/ϵ 2 , then with probability at least 1 -δ, (i,j)∈I1 α i,j z i z ⊤ j , (i,j)∈I2 β i,j z i z ⊤ j - (i,j)∈I0 α i,j β i,j ≤ ϵ •   (i,j)∈I1 |α i,j |     (i,j)∈I2 |β i,j |  holds for any α i,j , β i,j . Furthermore, if d ≥ 64k 2 log(2n 2 /δ)/ϵ 2 , then with probability at least 1 -δ,(i,j)∈I1 α i,j z i z ⊤ j , (i,j)∈I2 β i,j z i z ⊤ j - (i,j)∈I0 α i,j β i,j ≤ ϵ •   (i,j)∈I1 α 2 i,j   1/2   (i,j)∈I2 β 2 i,j   1/2holds for any α i,j , β i,j such that ∥α∥ 0 ≤ k, ∥β∥ 0 ≤ k.Proof. Using Lemma 8 and Cauchy-Schwarz inequality, we have(i,j)∈I1 α i,j z i z ⊤ j , (i,j)∈I2 β i,j z i z ⊤ j - (i,j)∈I0 α i,j β i,j ≤ (i,j)∈I0 α i,j β i,j (∥z i ∥ 2 ∥z j ∥ 2 -1) + (i,j)∈I1,(k,l)∈I2,(i,j)̸ =(k,l) α i,j β k,l z ⊤ i z k z ⊤ j z l ≤ 64 log(2n 2 /δ) d •   (i,j)∈I |α i,j β i,j |+ (i,j)∈I1,(k,l)∈I2,(i,j)̸ =(k,l) |α i,j β k,l |   = 64 log(2n 2 /δ) d •   (i,j)∈I1 |α i,j |   •   (i,j)∈I2 |β i,j |   ≤ 64k 2 log(2n 2 /δ) d •   (i,j)∈I2 α 2 i,j   1/2   (i,j)∈I2 β 2 i,j   1/2 .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(2n 2 /δ)/ϵ 2 , it holds that with probability at least 1 -δ,|⟨x i , x j ⟩ -δ ij |≤ ϵ, ∀i, j ∈ [n]. Proof. Fix i ̸ = j ∈ [n]. Notice |⟨x i , x j ⟩|= |⟨x i + x j , x i + x j ⟩ -⟨x i -x j , x i -x j ⟩|/4 ≤ (|⟨x i + x j , x i + x j ⟩ -1|+|⟨x i -x j , x i -x j ⟩ -1|)/4.Using Lemma 9, we have that with probability at least 1 -δ/n 2 , |⟨x i , x j ⟩|≤ ϵ.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),P (|⟨x, x⟩ -1|≥ v) ≤ 2e -v 2 d/16 .In particular, when d ≥ 16 log(1/(2δ))/v 2 , we have 16 . The second inequality follows from simple arithmetics.P (|⟨x, x⟩ -1|≥ v) ≤ δ Proof. Let x = (x 1 , . . . , x d ). By Lemma 11 (letting x = v 2 d/16), P (|⟨x, x⟩ -1|≥ v) ≤ P d i=1 ( √ d • x i ) 2 -d ≥ vd ≤ 2e -v 2 d/. 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 df 1 (t) dt ≤ -c 1 • f 1 (t) 2 , df 2 (t) dt ≥ -c 2 • f 2 (t) • 1 t + c 3 .Thenf 1 (t) ≤ 1 c 1 t + 1 f1(0) , f 2 (t) ≥ f 2 (0) • 1 + t c 3 -c2.Proof. The conditions imply thatdf -1 1 (t) dt = - 1 f 2 1 (t) • df 1 (t) dt ≥ c 1 , d log f 2 (t) dt = 1 f 2 (t) • df 2 (t) dt ≥ -c 2 /(t + c 3 ).It follows thatf -1 1 (t) ≥ c 1 t + f -1 1 (0), log f 2 (t) ≥ -c 2 log(1 + t/c 3 ) + log f 2 (0).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, Pr t i=1g 2 i ≥ t + 2 √ tx + 2x ≤ exp(-x),andPr t i=1 g 2 i ≤ t -2 √ tx ≤ exp(-x).In this section, we show missing proofs in Section 4. Proof of Proposition 4.1. Actually, for any three tokens x 1 , x 2 , x 3 , it holds thatp θ(0) (x 3 |x 1 , x 2 ) = exp x ⊤ 3 Y (0) ⊤ LN(X ⊤ b 2 ) x ′ ∈[M ] exp x ′ ⊤ Y (0) ⊤ LN(X ⊤ b 2 ) = exp (0) x ′ ∈[M ] exp (0) = 1/M, since Y (0) = 0.Proof of Proposition 4.2. Note that the input sequence length T = 2. By (1),p θ(t) (x|x 1 , x 2 ) = exp x ⊤ Y (t) ⊤ LN(X ⊤ b 2 ) x ′ ∈[M ] exp x ′ ⊤ Y (t) ⊤ LN(X ⊤ b 2 ) 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 p θ(t) (x|x 1 , x 2 ) = exp x ⊤ Y (t) ⊤ x 1 x ′ ∈[M ] exp x ′ ⊤ Y (t) ⊤ x 1 = exp (Y (t) x1,x ) x ′ ∈[M ] exp (Y (t) x1,x ′ )where Y (t) i,j is the entry of the matrix Y (t) at row i and column j.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 Ẏ = η Y LN(X ⊤ b T )(x T +1 -α) ⊤ = η Y x 1 (x 3 -α) ⊤ where α = [α 1 , α 2 , . . . , α M ] ⊤ ∈ R M with α = exp Y ⊤ x1 /1 ⊤ exp Y ⊤ x1 .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 x1,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 x1,0 = 0 for notation convenience, thenY (t x1,i ) x1 = Y (t x1,i-1 ) x1 + η Y (x 3 -α) ⊤ .For convenience, we denote y(i) = Y (t x1,i ) ⊤ x1 , and thus y(i) = y(i -1) + η Y (x 3 -α(i -1)) (10) 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 thaty(i) = (M -1)h * (i)ξ x3 ,whereξ x3 = M M -1 (x 3 -1 M 1) and h * (i) can be derived recursively as h * (i) = h * (i -1) + η Y (M -1) + exp(M h * (i -1))with h * (0) = 0. Combining Lemma 7 and 9 in [14] , we haveEQUATION)Note that the update of each row of Y are independent, the training set has size N , and the batch size is 1, we haveY (t) ⊤ x1 = (M -1)h * (⌈t/N ⌉) ξ x3where the training data at time step t is (x 1 , x 2 , x 3 ). Combining (11) , we can obtain thatY (t) x1,x3 ≳ (M -1) • M M -1 (1 - 1 M ) • 1 M ln (M η Y ⌈t/N ⌉) ≥ ln M η Y t N and Y (t) x1,x ≲ (M -1) • M M -1 (- 1 M ) • 1 M ln (M η Y ⌈t/N ⌉) ≤ - 1 M ln M η Y t N , ∀x ̸ = x 3 .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 haveY (t) x1,x = Y (0) x1,x = 0, ∀x ∈ [M ].Proof. We first consider training sequence (x 1 , x 2 , x 3 ) at time t. By Proposition 4.2, we havep θ(t) (x 3 |x 1 , x 2 ) = exp (Y (t) x1,x3 ) x ′ ∈[M ] exp (Y (t) x1,x ′) and by Lemma 2, we haveEQUATIONY (t) x1,x ≤ - c M ln M η Y t N , ∀x ̸ = x 3for some constant c > 0. Therefore,p θ(t) (x 3 |x 1 , x 2 ) ≥ exp c ln M η Y t N exp c ln M η Y t N + (M -1) ln -c M ln M η Y t N ≥ M η Y t N c M η Y t N c + (M -1) =1 - M -1 M η Y t N c + (M -1) ≥1 - M -1 2 M η Y t N cwhere the last inequality holds sinceln t ≳ ln(N M/η Y ) =⇒ t ≥ N (M -1) 1/c M η Y =⇒ M η Y t N c ≥ M -1.Finally, for any sequence(x 1 , x 2 , x 3 ) ∈ D test , since Y (t) x1,x = 0, ∀x ∈ [M ]according to Lemma 2, we havep θ(t) (x 3 |x 1 , x 2 ) = exp (Y (t) x1,x3 ) x ′ ∈[M ] exp (Y (t) x1,x ′ ) = exp(0) M • exp(0) = 1/M,which completes the proof.Datasets. Let N train > 0, N test > 0 be two positive integers and letN total = N train + N test . Let A i , B i , C i ∈ V, ∀i ∈ [N total ] be 3N total distinct tokens. Let →, ; ∈ V = [M ]be two additional different tokens that represent "direct implication" and "indirect implication" respectively. Specifically, we haveA 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 setsI train = {1, 2, . . . , N train }, I test = {N train + 1, . . . , N total }. 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 inp θ (x 3 = C i |x 1 = A i , x 2 =;), i ∈ I test .We also usep θ (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 and1 M 0.99 ≪ η Y < 1. Let t ≳ N ln M η Ydenote the time step which also satisfies ln t ≳ ln(N M/η Y ). For any test index i ∈ I test , we havep θ(t) (B i |A i →) ≥ 1 - M -1 2 M η Y t N c , p θ(t) (C i |B i →) ≥ 1 - M -1 2 M η Y t N cfor some constant c > 0 andp θ(t) (C i |A i ;) ≤ 1 M .Proof. Recall that by Proposition 4.2, we havep θ(t) (x 3 |x 1 , x 2 ) = exp (Y (t) x1,x3 ) x ′ ∈[M ] exp (Y (t) x1,x ′) and by Lemma 2, we haveEQUATIONY (t) Ai,x ≤ - c M ln M η Y t N , ∀x ̸ = B ifor some constant c > 0. Therefore, using the same proof as Theorem 3, we havep θ(t) (B i |A i →) ≥ 1 - M -1 2 M η Y t N c .Additionally, according to the proof of Lemma 2, Y (t) Ai,x has the same value across all x ̸ = B i , which impliesp θ(t) (C i |A i ;) ≤ 1 M .Similarly, applying Lemma 2 to Y (t) Bi , we can obtain thatp θ(t) (C i |B i →) ≥ 1 - M -1 2 M η Y t N c .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 denoteN total = N train + N test . Let A i , B i ∈ V, ∀i ∈ [N total ] be K ≜ 2N 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 haveA 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 setsI train = [N train ], I test = [N total ]\I train .The training set D train consists of allA 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 ].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 Ẏ = η Y f n (x T +1 -α) ⊤ = η Y f n (e n -α) ⊤ .Note that for three-token sequences, f n is a one-hot vector, and thus Ẏ 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 denoteW = [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 and1 M 0.99 ≪ η Y < 1. Let t ≳ N ln M η Yand 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 haveEQUATIONW (t) x4,x ≲ -ln (M η Y t/N ) /M, ∀x ̸ = x 4 ,and for any test sequence(x 1 , R 1 , R 2 , x 4 ) ∈ D test , we have W (t) x4,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ẇx4 (t) = η Y (x 4 -α x4 (t))where α(t) = exp(w x4 (t))/(1 ⊤ exp(w x4 (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 and1 M 0.99 ≪ η Y < 1. Let t ≳ N ln M η Ydenote the time step which also satisfies ln t ≳ ln(N M/η Y ). For training sequence (x 1 , R 1 , R 2 , x 4 ) ∈ D train at time t, we havep θ(t) (x 4 |x 1 , R 1 , R 2 ) ≥ 1 - M -1 (M η Y t/N ) c → 1, as t → ∞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 havep θ(t) (x 4 |x 1 , R 1 , R 2 ) ≤ 1/M. Proof. For any sequence (x 1 , x 2 = R 1 , x 3 = R 2 , x 4 ) where T = 3, p θ(t) (x|x 1 , R 1 , R 2 ) = exp x ⊤ Y (t) ⊤ LN(X ⊤ b T ) x ′ ∈[M ] exp x ′ ⊤ Y (t) ⊤ LN(X ⊤ b T ) = exp x ⊤ Y (t) ⊤ f x4 x ′ ∈[M ] exp x ′ ⊤ Y (t) ⊤ f x4 = exp x ⊤ w x4 (t) x ′ ∈[M ] exp x ′ ⊤ w x4 (t) = exp (W (t) x4,x ) x ′ ∈[M ] exp (W (t) x4,x ′ ).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Ẏ = η Y (f n -F E ′ e n )(e n -α n ) ⊤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Ẏ = η Y f n (e n -α n ) ⊤ ,one can obtain that the modification of the gradient of Y is small if λ 1 (E) is small.Note thatE 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, f n = LN(X ⊤ b T ) = b 13 e n ′ + b 23 e R1 b 2 13 + b 2 23, where b 13 = exp(Z R2,n ′ ), b 13 = exp(Z R2,R1 ). 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, ŻR2,n ′ > 0 and ŻR1,n < 0 and one can expect Z R2,n ′ -Z R1,n to be sufficiently large after sufficient time of training. Therefore,f n = b13 e n ′ + b23 e R1with b23 close to 0. Consider a simple case where for each n ∈ [N ], f n = √ 1 -c 2 e n ′ + ce R1 for c sufficiently small. Then E ij = f ⊤ i f j = ( 1 -c 2 e i ′ + ce R1 ) ⊤ ( 1 -c 2 e j ′ + ce R1 ) = c 2 . Therefore, one can calculate that λ 1 (E) = c 2 (N -1). When c ≪ 1 √ N , we have λ 1 (E) ≪ 1,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. 3 ). The curves represent the (average) negative log probability of the model predicting the next token to be: (1) B i given the input "A i →", (2) C i given the input "B i →", or (3) C i given the input "A i ;". Similar to the reversal curse experiment, while the sentences in the training set can be learned nearly perfectly, the model is not able to predict the correct next token in the validation set better than a uniformly random guess. Both curves are averaged over 10 random seeds.In this section, we show additional experimental results for Section 5.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.Values Learning Rate 0.01 Weight Decay λ 0.9 (β 1 , β 2 ) (0.9, 0.999) Batch Size 64 Number of Epochs 3000 Table 2 : Full list of hyperparameters for AdamW optimizer and training.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 ), 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.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. 3 . 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.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).The weights from A to B can be viewed as the logits of token B when the input is A.Empirically, d only needs to be of the order of logarithm of the vocabulary size. See Appendices E.2.1 and E.2.2 for additional results.Note that we assume the query token will not attend to itself as in[14].The lemma holds even if the batch size is larger than 1 and the analysis is essentially the same.By default, each entity consists of one token. See multi-token experiments in Appendices E.2 and E.3Apache License 2.0 


Decision: This paper presents a theoretical analysis of the "Reversal Curse" observed in language models. Specifically, this study examines the training dynamics of gradient descent in bilinear models and one-layer transformers, suggesting that the asymmetry in model weights during training is a significant factor behind the reversal curse. This issue is related to the commonly used cross-entropy loss, which does not enforce symmetry in learning process. The analysis is further extended to other logical reasoning tasks, like chain-of-thought reasoning, and the paper validates its theoretical findings with experiments on multi-layer transformers. The presented analysis provides novel insights into the Reversal Curse phenomenon and the paper is well written.

Reviews: review: paper_summary: paper_summary: This 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").
A 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.

- 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.

- 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.
strength_weakness: strength_weakness: - 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.
questions: questions: - 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"

    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.

- 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?

- What does Theorem 2 add beyond the interpretation of Theorem 1 in L161--L162?

- In Section 4, why do we need to assume the query token doesn't attend to itself?

- 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.

- 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?

- L87 "phenomenonin"
limitations: limitations: The 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).
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: Unfortunately, I identified the following issues with the paper in its current form:

* 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.

* 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).

1. 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.
questions: questions: Questions/suggestions:

* Define terms used in the text: examples include the term "symmetric loss" in line 47, the term "reparameterization" in line 59.
* 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.
* 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.
* 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?
* 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.
* line 625: It is not clear why you set $u = 0.2$.
* 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.
* 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.
* 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$).
* 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?
* 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".
* Figure 15 is interesting, the authors could perhaps consider add some discussion of it in the main paper.

A few typos/grammatical errors I caught:

* line 87: "phenomenonin" -> "phenomenon in"
* line 88: "Several paper studies".
* line 98: "There are rich literatures"
* 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.


1. 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.
limitations: limitations: Please 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).
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. Concerning the problem formulation, I found myself questioning its identifiability (Question 1).

2. Regarding the assumptions made and the proofs provided, I have a few questions that arose during my review (Question 2, 3, 4).
questions: questions: 1. In the context of this paper's system model, two rules merit consideration:
- RULE 1: If A leads to B (A->B), then it is also true that B leads to A (B->A).
- 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).

The 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.

2. 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?

3. 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.
- For Theorem 1 and 2, could the proofs be simplified while still achieving similar results if tokens are selected from one-hot vectors?
- Is it possible to apply this intuition to a one-layer Transformer model as well?

4. 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?

5. Regarding the experimental setup employing a 24-layer 12-head attention Transformer, this configuration seems considerably complex for testing the theoretical findings.
- Could you provide the rationale behind choosing such a large Transformer model?
- I suggest including experiments with a 1-layer Transformer to verify if the observed training dynamics align with the theoretical predictions.
- 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?
limitations: limitations: Yes
paper_summary: paper_summary: The 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.

The 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.

[1] Berglund et al. The Reversal Curse: LLMs Trained on "A is B" Fail to Learn "B is A". ICLR 2024.
strength_weakness: strength_weakness: * **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).
* **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.
* **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].

[2] Grosse et al. Studying Large Language Model Generalization with Influence Functions. 2023.
questions: questions: See weaknesses.
limitations: limitations: Some directions for future work are suggested in the last section. See also weaknesses.

Comments: comment: We are happy that we have addressed your question and thank you again for your helpful feedback. We will polish our work according to your feedback and our discussion.
Our intuition/guess of [1] is that the model did not really learn RULE 1 or RULE 2. For a training sentence A->B, the model learns a large weight from A to B, and by reverse training, the model also learns a large weight from B to A, and therefore, the model is able to infer B<-A. Our paper shows that in the three-token-sentence setting of 3.1 in [1] (which assumed RULE 1 as the underlying ground-truth rule), the above intuition is true and we provided theoretical proof for one-layer transformers. However, it requires more theoretical analysis/experiments to further validate the above intuition/guess in a more general setting. 

[1] Reverse Training to Nurse the Reversal Curse
The discussion between you and Reviewer bwDn has reminded me of a question that I had forgotten before. Regarding question 1, could you provide more insight into the relationship between your results and the experiments conducted in [1]? Given that the weights in [1] are asymmetric, I am curious whether the model learns RULE 1, or if it simply increases the size of the specific set $\mathcal{A}$ in RULE 2. Please note that your response to this question will not influence my decision, as I acknowledge that my comment is being posted quite late and I had forgotten this question earlier.

[1] Reverse Training to Nurse the Reversal Curse
Thank you for your very quick reply. I hope that my feedback and our discussion will help you improve the presentation of your work.
We thank the reviewer for reading our response, recognizing our clarification and ICL experiments, and giving further comments.

Regarding the motivation of our work, we agree that there are two possible cases/rules as also mentioned by Reviewer yXHq:

1. RULE 1: If A->B, then it is always true that B<-A.

2. RULE 2: If A->B, and A/B belongs to specific sets, then B<-A (which means it is not always true that B<-A).

An example of RULE 1 is ‘A is the parent of B’ is equivalent to ‘B is the child of A’. The example pointed out by the reviewer that 'a cat is black', but 'black is not a cat' is under RULE 2 which requires that B is not an adjective. However, no matter which rule we are choosing, the model is not able to infer B<-A trained on A->B. Although sometimes this is a correct behavior under rule 2 since under rule 2 the reverse is not always true, it is not because the model really learns the rule but simply because the model fails to infer B<-A no matter whether A/B belongs to the specific sets (due to model weights asymmetry). Therefore, the reversal curse happens under both rules, and our paper theoretically explained why the reversal curse happens when we do not expect it to happen (i.e., it happens under RULE 1). It would be an interesting future direction to anlyze the difference between model behaviors under these two rules.

Please feel free to let us know if you have any remaining questions. We will also make sure to include our clarifications in the rebuttal (especially the COT part) in the revised paper.
Thank you for your reply, and I apologize for the delay in my response. I appreciate the clarifications, and I hope they are included in the revised paper as well (especially the COT part).

> Also, while we definitely agree that we cannot expect a model to learn a reverse relation solely from training in only one direction, we emphasize that our training data contains both directions. Ideally, we expect the model to learn the reverse relationship between ‘$\to$’ and ‘$\gets$’ from training samples ‘$A_i\to B_i$’ and ‘$B_i\gets A_i$’, and then infer ‘$B_j\gets A_j$’ given ‘$A_j\to B_j$’ for new entity tokens $A_j$ and $B_j$, which is the same as what we would expect for human.

I disagree with the conclusion. The reverse relationship is not always true. Humans understand that when 'A is B', 'B is A' is not always true (e.g., 'a cat is black', but 'black is not a cat').

The motivation of the paper still remains weak/unclear to me, and this is why I am maintaining my score. I appreciate the ICL experiments, and I am sure that they strengthen the paper.
Thank you for reading our response. We are happy to make further discussion if you have any remaining questions.
Thank you for the detailed reply. I will maintain my score of the submission.
We sincerely thank the reviewer for recognizing our work and raising their score. Also, we thank the reviewer again for their great effort in reviewing our paper, reading our response, and providing helpful and insightful feedback.
I have adjusted the score accordingly. I hope this paper will serve as a starting point for the training dynamic analysis in higher-level intelligence tasks, such as reasoning and planning, within LLMs.
We thank the reviewer for reading our response and the further questions. In Section C.3, we assumed that the learning rate of the decoder layer $Y$ is much faster than the learning rate of the attention layer $Z$, and used another reparameterization $W = F^TY$. For a general setting with contextual token sequence length $T \geq 2$, one should
1. Analyze the dynamics of $W$.
2. Use $W$ to derive the dynamics of $Y$.
3. Use dynamics of $Y$ to further analyze the dynamics of $Z$.

The analysis of $Z$ is non-convex and definitely doable, as shown by previous work [1]. However, for the specific question that we are investigating in this paper (i.e., the reversal curse), it suffices to only analyze the dynamics of $Y$ (when $T=1$) or analyze the dynamics of $W$ (when $T \geq 2$). Note that this does not mean the dynamics of $Z$ do not matter. Actually, $W$ is determined by $Y$ and $F$, and $F$ is calculated using $Z$, so the dynamics of $W$ implicitly incorporate dynamics of both $Y$ and $Z$. One can definitely follow the above steps 1-3 to analyze the reversal curse or other settings, where the problem is non-convex, but for this paper, it would be more convenient to directly analyze $Y$ or $W$. Therefore, our method does not risk oversimplifying the complexities of the problem or the model dynamics.

**References:**

[1] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, 2023
I appreciate that the rebuttal has addressed most of my concerns. However, I am still unclear on how the analysis presented in Appendix Section C.3 extends beyond a simple convex setting, as mentioned in Question 4. Could you please clarify if this would result in an oversimplification? I am inclined to increase my score once this matter is resolved.