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

We thank the reviewer for their valuable comments. Below are our responses.

Meta-learning:

First, we emphasize that the near-orthonormality of the embeddings does not require the dimensionality to be much higher than the number of relevant tokens. In fact, it only requires the dimension $d$ and vocabulary size $m$ satisfies $d \geq \Omega(\log m)$, which means the dimension $d$ can be even much smaller than the number of tokens. This is also empirically verified in our additional experiments (Figure 1 in the global response pdf) The assumption of large dimensionality is only made to show the provable reversal curse phenomenon in bilinear settings (Theorem 2): in large dimensionality, the outer products of the embedding form a $\ell_1$ subspace embedding, which is critical to the theoretical result.

Also, our analysis reveals that given a large training set containing both Ai->Bi and Bi<-Ai where Ai and Bi can be chosen arbitrarily, the model simply learns (or memorizes) these facts/sentences instead of understanding the relationship between ‘->’ and ‘<-’. Therefore, even if the training set contains nearly infinite samples, the model might end up learning an (artificial) mapping between A and B but still fails to generalize to a new pair (Aj, Bj) since Aj and Bj can be arbitrarily chosen.

Asymmetry:

We argue that our result is not simply a verification of the obvious explanation but actually provides important new insights. Note that the previous empirical observation that ‘LLMs store factual associations as directed representations’ is a qualitative and vague argument. However, our analysis quantifies the phenomenon in a rigorous manner: the directed relationship between token $i$ and token $j$ is associated with the logits $Y_{ij}$. Our insights not only explained the reversal curse in a token-level granularity but can also be applied to explain other phenomena such as chain-of-thought. Also, as Reviewer yXHq mentioned, ‘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.’ We also thank the reviewer for the suggestion of including more discussion on previous empirical studies and we will better contextualize our contributions in the revision.

In-context learning:

We totally agree that comparing pre-training/fine-tuning and ICL would be an interesting future direction. As suggested also by the Reviewer bwDn, we added additional experiments on ICL (see the ‘In-context learning’ paragraph in the global response).

We thank the reviewer for their valuable comments. Below are our responses.

Answer to Question 1:

Thank the reviewer for pointing out this interesting question. Note that 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, it is not because the model really learns the rule but simply because the model fails to infer B<-A no matter whether A belongs to the specific set. Therefore, the reversal curse happens under both rules. It would be an interesting future direction to anlyze the difference between model behaviors under these two rules.

Answer to Question 2:

Actually, we only assumed the Gaussian distribution of token embeddings for bilinear settings, and the actual condition we used in our proof is that those embeddings are nearly orthogonal to each other, which is empirically validated in our experiments (see Figure 1 in the global response pdf). Therefore, the result is applicable to practical LLM token embedding distributions. For sequences containing more than one contextual token, see responses to Question 4.

Answer to Question 3:

Yes. If tokens are one-hot vectors, the proofs can be simplified since we don’t need Lemma 7 which essentially shows that ‘near-orthonormal’ embeddings behave nearly identically to the one-hot embeddings. Also, we don’t need the embedding dimension $d$ to be as large as $poly(m)$. The intuition between lines 173-179 can be applied to (and actually is the same as) a one-layer transformer, where the model parameters $\Theta_{ij}$ in bilinear setting corresponds to the logits $Y_{ij}$ for transformers.

Answer to Question 4:

We study the case where $T=2$ in our main results since we aim to more clearly convey the core reason of the reversal curse, i.e., the asymmetry of model weights. Our analysis can definitely be extended to scenarios involving a greater number of contextual tokens. Actually, we discussed the case $T=3$ in Section 4.3 and provided analysis for $T=3$ in Appendix C.3.

Answer to Question 5:

The reason we chose a large transformer model for experiments is that theoretically we proved the reversal curse for one-layer transformers and revealed a core reason for the reversal curse, i.e., model weights asymmetry, so we aim to empirically verify that the result still holds for multi-layer transformers. We added experiments with 1-layer transformers and also the case where the embedding dimension is much smaller than the vocabulary size, and the result still holds (see Figures 1 and 2 in the global response pdf). Actually, the embedding dimension $d$ only needs to be as large as the log of vocabulary size to satisfy the near-orthogonal property. We assumed $d = poly(m)$ in the proof of bilinear setting solely for a rigorous theoretical guarantee, and in practice, $d \gtrsim \log m$ suffices.

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.

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 their valuable comments. Below are our responses.

Unclear motivation of the theoretical framework: ...

First, we want to emphasize that according to our analysis in this work, the reversal curse is hard to mitigate even on a well-structured , synthetic dataset both theoretically and empirically, which implies that the reversal curse is even harder to mitigate in a more complex real-world dataset. Therefore, our setup is appropriate for revealing the most essential reason of the reversal curse.

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.

For the analysis of pertained models, in fact, our result for bilinear models also works if the model parameter $\Theta_0$ is initialized to a pre-trained value satisfying the condition $\frac{1}{2m} < p_{\Theta_0}(y_i|x_i), p_{\Theta_0}(x_i|y_i) < \frac{2}{m}$ (as shown in Theorem 5). This implies that the reversal curse would also happen in the pre-trained models. However, we also want to point out that we haven't considered the learning of embedding in our theoretical study, which may lead to further generalization in the pre-trained model (and might be a (partial) solution to the reversal curse).

Imprecise writing: ...

Thank you for your suggestions and we will incorporate them in the revision.

Define terms used in the text: examples include the term "symmetric loss" in line 47, the term "reparameterization" in line 59.

One example of the symmetric loss is $L = \sum_{t=1}^T \log p(x_t|x_{1…t-1}) + \log p(x_t|x_{T…t+1})$ and more broadly, it can be defined as $L(x_{1…T}) = L(x_{T...1})$ for any sequence. Reparameterization is a general term, and in this paper, it refers to $Z = UW_QW_K^TU^T/\sqrt{d}$ and $Y = UW_V^TU^T$. We will make these terms more clear and easier to understand in the revision.

Can you please clarify the sentence in line 297-298 "Note that ... will not attend to itself" ...

In Line 196, the attention score vectors $b_T$ only contains $b_{1,T}, …, b_{T-1,T}$ which are the attentions of contextual tokens $x_1, … x_{T-1}$ attended by the query token $x_T$. Note that this framework does not take into account $b_{T,T}$, which is the score that $x_{T}$ attends to itself.

A criticism of Theorems 1, 2 could be that they rely on the large dimension of the matrix $\Theta$...

Please check the 'Small embedding dimension and number of layers' section in the global response.

The proofs of the Theorems and Propositions appear in a convoluted order in the Appendix ...

Thank you for your valuable suggestions. In revision, we will improve the structure and readability of the proof.

line 625: It is not clear why you set 𝑢=0.2.

$v$ is set to make sure that $p_\Theta(x_i|x_j)$ is of order $1/n$. The proof works for any small constant.

line 626: Can you please add some explanation on the first inequality? ...

Fixing $x_i,x_j$, the logit $l_{\Theta}(x_j|x_i)$ is a Gaussian random variable and the bound follows from standard concentration inequality for Gaussians. In revision, we will make this argument clearer.

line 641, 'Using Lemma 9': Can you please elaborate on that? ...

In line 640, we use the rectangular trick to reduce vector inner products to vector norms. Since the vectors $x_i+x_j$ and $x_i- x_j$ are also Gaussians, we apply Lemma 9 to bound their norms and thus by the right-hand side of line 640, we can bound the inner products.

Lemma 10 in page 22: A little bit more care should have been taken in the proof of this Lemma ...

Thank you for your valuable suggestion. The functions cannot attain the 0 value because letting $\tau = \inf$ { $t: f(t) \leq 0$ }, a contradiction would happen in the interval $(0,t)$ by the same argument of the proof. We will add this point in the revision.

Theorem 4: I do not understand the conclusion of this theorem ...

Thanks for pointing this out. We agree that it would be more appropriate to change the ‘necessity of COT’ to the ‘importance of COT’ since Theorem 4 does not exclude all other inference techniques than COT. The main point we want to emphasize here is that the model fails to directly output an indirect conclusion.

Suggestion on experiments: An interesting observation of [1] is that the "reversal curse" is remedied if the reverse sentence appears in context ...

Thanks for your suggestions on experiments. A similar observation can be made in our settings, and please check the ‘In-context learning’ paragraph in the global response.

Figure 15 is interesting, the authors could perhaps consider add some discussion of it in the main paper.

Thanks for your suggestions. We will use the additional content page for the camera-ready version to add more discussions.

A few typos/grammatical errors I caught: ...

Thank you for catching these typos, and we will fix them in the revision.

References:

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

We thank the reviewer for their valuable comments. Below are our responses.

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.

According to the currently widely-used training paradigm (CE loss for next token prediction using SGD/AdamW, etc., without hard-coding constraints on model parameters), the reversal curse might be unavoidable. Therefore, we might need to change the training paradigm. An ultimate solution could be a method that really enforces the model to understand the higher-level regularity that two relationships are inverse to each other, instead of just learning how to predict the next token. A potential partial solution might be learning in a latent space where the model parameter in this latent space shows symmetry to some degree inspired by our model weights asymmetry analysis. All these are interesting and important future directions.

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.

The constraints in L42 refer to hard constraints on model parameters during training/optimization (e.g., a concrete constraint on model parameters $\Theta$ can be $\Theta$ is symmetric, or $\| \Theta \|_2 = 1$). Since some easy-to-implement constraints might hurt the model performance and the correct constraints could be rather complicated and hard to enforce, in practice, we often use unconstrained optimization, and thus the model learns the relationship between ‘->’ and ‘<-’ through lots of examples ‘A->B’ and ‘B<-A’. Thanks for pointing this out, and we will make the statement more clear in the revision.

What is 𝜎 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?

$\sigma$ is the standard deviation of the model parameter $\Theta\_0$. Theorem 5 shows that the reversal curse would happen (with high probability) under either of these two conditions. Therefore, we say that $\Theta\_0$ can also be initialized as a pre-trained value: as long as it satisfies the condition $\frac{1}{2m} < p\_{\Theta\_0}(y\_i|x\_i), p\_{\Theta\_0}(x\_i|y\_i) < \frac{2}{m}$.

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

Theorem 2 is a direct corollary of Theorem 1, stating more explicitly that with high probability the forward loss is low while the reverse loss is high.

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

This assumption is mainly for the convenience of theoretical analysis so that the query tokens and contextual tokens can be disjoint, and this assumption is also adopted by previous work both theoretically [1] and empirically [2]. Our experimental results show that when the query token can attend itself, the conclusion still holds. Moreover, even if we assume the query token can attend to itself, previous results [1] indicate that the query token will gradually put more attention weights to distinct tokens (i.e., the entity token Ai or Bi in our settings) and lose attention to common tokens (i.e., the query token itself in our setting) and thus is effectively the same as our assumption.

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.

Yes, the reason that we start with the actual parameterization and then the reparameterized form is to draw a more clear connection to the transformer.

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?

In section 4, even if $\{A\_i\}$ and $\{B\_i\}$ are disjoint and linearly separable, the reversal curse still happens. The most straightforward example is when all $A\_i$ and $B\_i$ have one-hot embedding, the two sets are disjoint and linearly separable. Our analysis in section 4 then shows the reversal curse still happens. In practice, the embeddings of different tokens are not exactly one-hot but still nearly orthogonal, which is indicated by Figure 1 in the global response pdf. Therefore, even if the embeddings of two sets are linearly separable, knowing what kind of entity a new B is does not help with reversal.

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

[2] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ICLR, 2020.