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

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.

Strengths

The paper provides a timely study on a fundamental limitation of the reasoning capabilities of LLMs. The analysis is able to incorporate the next-token prediction mechanism with CE loss with a detailed breakdown of the gradient dynamics and sheds light on various aspects of the reversal curse empirically demonstrated in [1]. Moreover, the modification in Section 4.2 for CoT is very interesting and provides a new approach to the theory of CoT, which so far has been mostly limited to complexity-based expressivity analyses.

Weaknesses

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

See weaknesses.

Rating

5

Confidence

4

Soundness

4

Presentation

3

Contribution

2

Limitations

Some directions for future work are suggested in the last section. See also weaknesses.

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.

Strengths

1. The exploration of the "reversal curse" is a timely and relevant topic within the field. This paper stands out as a pioneering theoretical analysis on this subject, which is commendable. 2. From a technical standpoint, this paper is well-founded. The depth of the analysis is impressive and may hold value not only within the context of this study but also as a reference point for future research. The methodology and the thoroughness of the analysis contribute significantly to the paper's technical soundness. 3. The investigation into the problem, specifically highlighting that the issue stems from training dynamics rather than expressiveness, is particularly insightful. This perspective is refreshing and challenges conventional approaches, potentially inspiring a new stream of research in this direction.

Weaknesses

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

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?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes

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.

Strengths

This work studies a well-documented empirical phenomenon of autoregressive language models, namely their inability to operate as databases (since the order of a query matters) - a phenomenon called "the reversal curse." A better theoretical understanding of the subject could have a substantial impact. The current work appears to be the first to study it theoretically from an optimization perspective. The paper has certain merits: in particular, it offers an interesting theoretical analysis of training dynamics of simplified models on synthetic data. The proven theorems/propositions seem correct to me (I did not manage to verify all the steps - see questions on how the presentation of some results could be improved), and the techniques, whilst appearing to be standard, are interesting.

Weaknesses

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

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

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

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.

Strengths

- Analyzes the training dynamics of increasingly more complex transformer-like models to show why the reversal curse emerges. - The techniques used in the paper can also be used to analyze why language models sometimes can't distill chain-of-thought reasoning: even though they can directly learn A-->B and B-->C, they can't always directly predict A-->C. - On a technical level, this paper relaxes some assumptions from other work on training dynamics of the bilinear model in Section 3.

Weaknesses

- 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

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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

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