Paper: When can transformers reason with abstract symbols? (ICLR-2024)

Reviewer comment (Writing):
The abstract and introduction refer to 'transformer large language models'. However, if I understand correctly, the experiments in Figures 1 and 2 are performed on transformers trained from scratch, not on LLMs (though an LLM is tested later on). If this is correct, I think it would be clearer to simply refer to transformers (rather than LLMs) in the abstract and introduction.

Author response:
Thank you for the suggestion. We will edit the intro and in particular our presentation of Figures 1 and 2 to make it clearer that those experiments are on language transformers and not pretrained language models.

Reviewer comment (Writing):
It would be helpful if an intuitive characterization of the proposed modifications ($aI$ and $bI$) could be provided earlier in the paper, along with some sense of why this modification helps.

Author response:
Thank you for the suggestion. We will try to incorporate this into Section 1.1, and will expand on the intuition of the modifications in Section 4.2 and Section 5.

Reviewer comment (Experiments):
Do the wikitext experiments involve OOD generalization in any sense (in the same way that the reasoning tasks do)? If not, I wonder whether this partly explains the smaller relative improvement from the modified models vs. GPT-2.

Author response:
Good point. The Wikitext task is next-token prediction on Wikipedia, so it does not explicitly involve OOD generalization. There may be some implicit reasoning tasks in this problem, which could explain the improvement.

Reviewer comment (Experiments):
I was somewhat confused by the distinction between GPT-2 and 'GPT-2 pretrained' -- isn't GPT-2 by default pretrained (since this term refers to the trained model)? Just to clarify, do the results in Figure 4 reflect pretrained GPT-2, which is then fine-tuned on the wikitext datasets, or both of the models in this figure are trained from scratch only on wikitext?

Author response:
These are models trained from scratch only on wikitext. We will clarify in the text.

Reviewer comment (Presentation):
The legend is cut off in figure 3a. In general the figure text is very small and difficult to read.

Author response:
Our apologies. This got accidentally introduced when we regenerated the figures right before submitting.

Reviewer comment (Theory):
Do the authors expect that larger models (e.g. GPT-3 and GPT-4) may be able to perform the abstract reasoning tasks from this paper without the proposed modifications, and if so how does that relate to the theoretical results? Do the authors expect that this capability would similarly depend on the emergence of heads that implement a diagonal attention strategy? Does the analysis in this paper have anything to say about the ability to learn these kinds of abstract reasoning tasks through *in-context* learning (as opposed to direct training of the model parameters)?

Author response:
* This is an excellent question. Unfortunately, our theory does not bear upon in-context learning. It seems that to understand that it will be necessary to understand feature-learning for these template tasks – our paper provides a baseline to build off of when conducting those future analyses.

* We expect that the feature learning in the form of strong diagonal weights in the attention heads plays an important role in practice for both copying and for checking whether two symbols are equal. We do not know whether this learned weight structure is necessary for in-context learning to succeed, although our intuition is that it is. This is indeed a fantastic direction for future work.

Reviewer comment (Novelty):
How might the proposal in this work relate to the idea of 'induction heads' [1]? Intuitively these seem to be related, in that both exploit the inductive biases of transformers to perform abstract tasks involving copying.

Author response:
There is an interesting relation between our result and that work. Our result on copying, and our experiment in Figure 3a, indicate that the larger the embedding dimension, the more difficult it becomes for the model to learn to copy symbols it has not seen copied at training time. Thus, we should expect infinitely-large models to be unable to learn induction heads that generalize outside of the training data. So it would seem that transformer modifications are necessary for infinitely-large models if you want data-efficiency in learning to copy.

Reviewer comment (Theory):
The theory for the transformers seems to rely on training only the final layer of the transformer. Does that mean that there are other architectures that can also do symbol binding?

Author response:
* Our analysis identifies the condition number of the N matrix as a key driver of the sample complexity for learning the template task. This condition number is always infinite for MLPs, and so they do not perform symbol binding. (In fact, we prove that MLPs do not perform symbol binding even when all parameters are trained.)

* For non-transformer architectures, the condition number of the N matrix might be finite, in which case those architectures would also perform symbol binding. However, we only prove this for transformers in this paper. We will improve the clarity of the proof sketch in Section 4.2 of Lemmas 4.6 and 4.7 in the revision to better explain why the condition number is finite. The main reason is due to the self-attention layer, which, roughly speaking, allows the architecture to access the matrix XX^T. This matrix contains information about which symbols in the string are repeated and where. We will also improve the explanation in Section 4.3 for how this insight motivates the $W_KW_Q^T + aI$ modification. Here it may be illustrative to give concrete bounds on the condition number before/after the modification for the same-different task.

* Re: training only the last layer, please also see our answer to the question 2 below.

Reviewer comment (Experiments):
It is unclear which aspect of transformers are important for symbol binding: is it the attention, the MLP, the invariance properties or something else?

Author response:
* Our analysis identifies the condition number of the N matrix as a key driver of the sample complexity for learning the template task. This condition number is always infinite for MLPs, and so they do not perform symbol binding. (In fact, we prove that MLPs do not perform symbol binding even when all parameters are trained.)

* For non-transformer architectures, the condition number of the N matrix might be finite, in which case those architectures would also perform symbol binding. However, we only prove this for transformers in this paper. We will improve the clarity of the proof sketch in Section 4.2 of Lemmas 4.6 and 4.7 in the revision to better explain why the condition number is finite. The main reason is due to the self-attention layer, which, roughly speaking, allows the architecture to access the matrix XX^T. This matrix contains information about which symbols in the string are repeated and where. We will also improve the explanation in Section 4.3 for how this insight motivates the $W_KW_Q^T + aI$ modification. Here it may be illustrative to give concrete bounds on the condition number before/after the modification for the same-different task.

* Re: training only the last layer, please also see our answer to the question 2 below.

Reviewer comment (Theory):
A limitation of most theory that uses the PAC framework often has vacuous bounds. While the paper derives asymptotic trends, it is unclear if they will be predictive of what happens in practice as a result.

Author response:
* Yes, our sample complexity bounds are indeed non-quantitative in terms of the number and structure of the templates. The bounds could in principle be numerically computed by estimating the network’s kernel matrix.

* Nevertheless, we do get a qualitative result on the capabilities of transformers on symbol binding tasks. And we believe that our theory is relevant to practice since it suggests modifications to the transformer that improve data-efficiency on the reasoning tasks in practice by an order of magnitude.

Reviewer comment (Theory):
The proof for the failure of MLPs relies on taking an expectation over $\theta_t$ obtained through gradient descent. Does this also hold for a single run of an MLP training? Is it possible for MLPs to converge some of the time (and not with probability at least $1 - \delta$)?

Author response:
The proof of failure of MLPs relies on a coupling between the weights of the MLP trained with the true data distribution D, and the weights of the MLP trained with a fake, permuted data distribution D’. This coupling holds almost surely.
Using this coupling, we show that the loss of the MLP after training is lower-bounded in expectation over the randomness of the initialization, the random shuffling of the samples, and the noise in the samples. So it is indeed possible that the MLP may in exceptional cases have good test loss due to the randomness of the algorithm (this will just not be true in expectation). It would be interesting to extend the result to bound the probability of these events.

Reviewer comment (Theory):
The proof seems to rely on the kernel ridge regression estimator. Does this mean that representation learning is not critical for symbol binding?

Author response:
* It was quite unexpected to us that the transformer kernel analysis is sufficient to characterize when the relational reasoning tasks can be learned. We began this work with the mindset that we would have to study feature-learning dynamics. Nevertheless, with the kernel analysis, we were able to *characterize* when relational reasoning can be learned by large transformers (it can be learned anytime there is no copying of placeholder symbols; it cannot be learned when there is).

* A consequence of our analysis is that any separation between NTK and feature learning would only be quantitative in terms of the number of samples. I.e., to show such a separation we would need to prove bounds on the number of samples with explicit quantitative dependence on the structure and number of the templates. We do not presently know how to prove these, but we feel that it is very interesting for future work. This is a very intriguing question and we will add discussion in the final section.

* Also, in the revision, we will add experiments comparing feature learning to NTK training – these indicate that feature learning is intriguingly more data-efficient on the relational reasoning tasks (although NTK also learns given enough data).

Reviewer comment (Theory):
If the data diversity is small (data follows a Zipf distribution), then do these results hold? In practice we expect datasets to follow such distributions and I am unsure if I am misunderstanding the theorem incorrectly.

Author response:
* This is a great question. If the data follows a Zipf distribution with word k occurring with probability $p_k \propto 1/k$, then it will occur a total number of times $p_k \\approx ck / \\log K$, where $K$ is the total number of words in the language. As the total number of words increases, then this probability vanishes and so the condition on our data holds with diversity parameter $\rho \propto log(K)$.

* However, $log(K)$ is a slowly-growing and so we could hope to do better. Perhaps we should not think of common words as those words that substitute placeholder wildcard tokens. Instead, if a word is sufficiently common then maybe it should be modeled as a regular token in a template. For example, by a large margin the most common words in English are words like “the”, “to”, “be”, “and”, and “of”. These words have specific meanings and we should not generally think of them as abstract, interchangeable symbols, in the same way that we think of names like “Alice”, “Bob”, etc… as interchangeable. Interestingly, large language models appear to sometimes fail when these more common words are used as part of abstract reasoning: see [1]. So maybe our analysis is tight in some sense?
* There definitely seems more to learn by exploring in this direction. It seems that to get a tighter understanding we may need to derive new sample complexity bounds that depend quantitatively on the number (and potentially structure) of the templates.

[1] https://arxiv.org/abs/2305.15507, “The Larger They Are, the Harder They Fail: Language Models do not Recognize Identifier Swaps in Python” by Miceli-Barone et al.

Reviewer comment (Theory):
Are there technical reasons for using the kernel regression with gradient flow in section 4, but the cross-entropy loss with next-token prediction for section 5? Is it possible to conduct an analysis similar to section 4 if we consider next token prediction?

Author response:
The technical reason is that we do not have a closed-form solution to the kernel regression with cross-entropy loss. It would be very interesting to find ways to prove the result in this cross-entropy setting without such a closed-form solution. We will mention this in the paper.

Reviewer comment (Evaluation):
I think the authors are over-claiming their results at several points. The simplified architecture that is studied in this work, i.e., a self-attention + MLP model, is not justified to be labeled a Transformer. Residual connections play a huge role in both optimization stability and, e.g., some of the seemingly emergent abilities of Transformers. In this sense, I would argue the paper is actually focused on understanding inductive biases and limitations of a feedforward, MLP model with self-attention. This is totally fine, but the over-claiming can yield, in my opinion, an incorrect takeaway that Transformers broadly cannot perform the variable binding tasks. It may also be worth mentioning that the results focus on a inference with a single pass through the model, i.e., tools like in-context learning or chain-of-though can impact the results. The latter is just a nitpick to better contextualize the paper's findings.

Author response:
* We are sorry that we give the impression of overclaiming. It is not at all our intention to overclaim. Our reasoning is as follows:
    *  As far as theoretical results, the self-attention+MLP setting in our work is not unusual and even goes beyond many theoretical works. All known theoretical analyses of transformer training dynamics use a combination of simplifications such as: (a) removing the softmax to get a linear attention mechanism, (b) replacing the $W_KW_Q^T$ matrices with just one matrix $W_{KQ}$ and similarly for $W_VW_O^T$, (c) considering only one layer, (d) assuming diagonal weight matrices, (e) removing layernorms and residual connections, and (f) even removing the MLP layers.
    * Furthermore, we believe that the same ideas in our analysis could be used to prove our theoretical results in the multilayer setting with residual connections – but this would require an even more involved analysis than what we currently have here, and we did not write this because we feel that the more convoluted technical arguments would distract from the main ideas of the proof which only need one self-attention layer + one MLP layer.

    * In our experiments, we do train standard multilayer transformers, including practical models such as GPT-2 where we show a benefit of our theory-inspired modifications.

* With respect to “the takeaway that Transformers broadly cannot perform the variable binding tasks”, this was not our intention. Our intention was only to highlight that a large amount of data is needed for simple variable binding tasks, and that theory-inspired modifications can lower that amount of data substantially. Also, we uncovered an intriguing inverse scaling law for copying which shows that making transformers larger can actually break down some of their reasoning abilities on unseen symbols.

Reviewer comment (Experiments):
Lack of experimental details: While the authors often state that their theoretical claims focus on a simplified architecture, it is unclear to me if the experiments involve this simplified architecture or an actual Transformer, i.e., with residual connections, is used. Generally, I think the experimental details can be improved.

Author response:
* Thank you for your suggestion. We will make it clearer that our experiments are on standard multilayer transformers with residual connections and layer-norm.

Reviewer comment (Novelty):
Relation with prior theoretical works on inductive biases of self-attention: A specific work that addresses a similar question as this paper is Edelman et al. [1]. Therein, the authors show a self-attention model can express sparse boolean circuits and performs implicit variable creation during forward inference. I think the work deserves thorough discussion in the current paper's related work. A few more empirical works also exist on this problem, e.g., Davies et al. [2], that will be worth discussing.

Author response:
* Thank you for the references. These seem interesting and we will certainly incorporate discussion in the revision.

Reviewer comment (Novelty):
Discussion with Systematic generalization: I would argue the tasks that the authors studied are extremely related to what NLP community has studied under the label of systematic generalization before. I think such papers warrant discussion as well (e.g., see [3]).

Author response:
* Thank you. We were unaware of this reference. Interestingly, our results show that reasoning is possible even without neuron/head specialization, which appears quite relevant to this reference. We will incorporate a discussion in the revision.

Reviewer comment (Theory):
The proposed changes based on the theoretical models involve addition of a scaled identity matrix to parts of the attention matrices. This looks, at a surface, like addition of a residual connection in fact. Can the authors clarify this further? I found the discussion of this point to be quite unclear.

Author response:
The $W_VW_O^T + bI$ modification is adding an "attention-modulated" skip connection. It is a skip connection except that it is multiplied by the attention matrix.

The $W_KW_Q^T + aI$ modification does not seem to have this interpretation. Instead, it encourages the attention head to use the matrix XX^T of the data, which encodes which symbols are repeated and where.

Reviewer comment (Theory):
Though the paper provides some generic results about architectures for MLPs, the main result for transformers relies on operating the model in the kernel regime. It is unclear at the moment if any of the results would change at finite width.

Author response:
Re: (1) finite width

* We will add experiments in the appendix to the revision where we vary the width of the network.

Reviewer comment (Experiments):
It is unclear at the moment if any of the results would change in the feature learning regime.

Author response:
Re: (2) feature learning vs. kernel analysis

* As far as we are aware, there are **no previous results that prove the ability of transformers to learn to perform symbolic reasoning tasks** when trained. Our paper is the first to formalize a framework of such reasoning tasks and to prove it is learned by transformers when trained.


* It was quite **unexpected to us that the transformer kernel analysis is sufficient** to characterize when the relational reasoning tasks can be learned. We began this work with the mindset that we would have to study feature-learning dynamics. Nevertheless, with the kernel analysis, we were able to *characterize* when the relational reasoning can be learned by large transformers (it can be learned anytime there is no copying of placeholder symbols; it cannot be learned when there is).


* **The analysis was considerably involved**: (i) This is not a typical kernel method analysis where it is proved that the data can be memorized – instead, we must show that the network *generalizes in the correct way out of distribution*, for any task in our framework. (ii) Also, insofar as we know, this is also the first paper that manages to analyze how the transformer kernel learns a nontrivial task, due to the technical difficulties involved in studying this kernel. The technical tools we develop here may thus be useful for future work.


* Note also that **all known theoretical analyses of transformer dynamics use a combination of simplifications** such as: (a) removing the softmax to get a linear attention mechanism, (b) replacing the $W_KW_Q^T$ matrices with just one matrix $W_{KQ}$ and similarly for $W_VW_O^T$, (c) considering only one layer, (d) assuming diagonal weight matrices, (e) removing layernorms and residual connections.



* We believe that our theory is relevant since it suggests modifications to the transformer that **improve data-efficiency on the reasoning tasks in practice by an order of magnitude**.


* **Our work opens up many directions** that we feel are outside of the scope of the current paper. One direction is to study feature learning. We completely agree that it is an excellent question to ask: can we characterize if/when feature learning is more powerful than NTK? But, to ask this question, an NTK analysis as in this paper is needed first. In the revision, we will add experiments comparing feature learning to NTK training – these indicate that feature learning is intriguingly more data-efficient on the relational reasoning tasks (although NTK also learns given enough data).

* On a more philosophical note, our results raise the question: to what extent is transformer reasoning due to NTK, and to what extent is it due to feature learning such as head specialization? This is a highly relevant question in the context of interpreting the internal workings of transformers. The more NTK-like the solution, the more polysemantic heads that the network will have, and so the harder it will be to interpret a network by extracting sparse subnetworks (which is currently the main method of the mechanistic interpretability literature).

Reviewer comment (Theory):
The proposed theory bounds the generalization error in terms of the data diversity metric, which does not have an obvious scaling with samples n, making it hard to reason about the number of samples needed to attain good generalization.

Author response:
A large number of samples n is not enough, because all n samples could be copies of each other, in which case the algorithm will not be able to generalize beyond the training data distribution. Therefore, it is necessary to have a control on the data diversity. We opt in our case to have a parameter rho which measures diversity, but otherwise we are agnostic to the data distribution.

Reviewer comment (Theory):
Do the sample complexity results / rank of N matrix change if you allow the network to learn its internal weights rather than just the readout weights?

Author response:
Thanks for the question – we will put discussion of this in the paper: Our result would not change; what we have proved about NNGP is a strictly stronger statement than the corresponding statement if all layers are trained together.

Reviewer comment (Experiments):
Do the authors have any empirical evidence showing the emergence of this block diagonal structure in the NNGP kernel at initialization (I imagine this would be easy to compute)? The block-diagonal approximation seems reasonable, but some empirical support and visualization could be useful. A similar comparison of the predictor from kernel regression and the average of the labels over a template could also be useful.

Author response:
As we promised above, we will add experiments for feature learning vs. NTK. These show that the NTK does succeed albeit with more samples than in feature learning. We will also add explicit approximation of the matrix N for the same-different task before/after the addition of the $W_KW_Q^T$ identity perturbations, illustrating how the condition number is vastly improved by our modifications.

Reviewer comment (Theory):
Do the authors have a sense of how wide the transformer must be compared to samples n, templates r, etc. for the theory to be accurate (N non-singular etc.)? Wouldn’t the kernels become singular at some finite width? What if every layer is trained at large width (NTK vs NNGP)? Do real networks (which learn features) have to be as wide as kernel analysis would suggest?

Author response:
This is an excellent question for future work, but is outside of the scope of our current paper. The width and number of samples would depend on the particular number of templates and potentially on the structure of these templates. In this setting with quantitative bounds, we could hope for a separation between training just readout vs. also training the internal weights. However, we do not presently know how this would work. We will add discussion of this question in the final section.

Reviewer comment (Evaluation):
The perplexity improvements on GPT + Wikitext seem rather small. Does this suggest that transformers trained in practice evolve to have well conditioned N matrices, removing the need to add an explicit identity matrix? Is this related to the pronounced diagonals in Wv Wo the authors report in the next section? Or is the idea that real data behaves differently than the proposed template tasks?

Author response:
* We are adding a total of 288 parameters to a 117M parameter architecture, which is a minuscule fraction of the total number of parameters. Per parameter, the drop in perplexity is significant. The transformer architecture is state-of-the-art and ubiquitous, so we believe that any improvement can be impactful.

* For comparison, for this rebuttal we ran GPT-2 medium (345M parameters) on Wikitext103 for 5 epochs on a V100. This takes 2.14 days and gets a perplexity of 17.904. Stopping GPT-2 small at those same 2.14 days gives a perplexity of 17.08, and stopping GPT-2 small plus the modifications for 2.14 days gives a perplexity of 16.693. Despite the fact that GPT-2 medium has 228M more parameters it actually underperforms because with the same computational budget it has not yet converged (we are continuing to run this comparison).

* Finally, you are correct that Wikitext is not a reasoning task, so seeing any significant improvements at all is a welcome surprise. We will try to benchmark our modifications on more pure reasoning tasks in the revision if we have time, but we do recall that this is not the main point of the paper, which is to theoretically characterize when transformer architectures can learn to perform symbol binding reasoning.

Reviewer comment (Theory):
In Theorem 1.1 it may be helpful to clarify that the gradient flow is only performed on the last layer.

Author response:
Thanks for the suggestion. It is in the formal statement, and in the revision we will add it to the informal statement.

Reviewer comment (Presentation):
Figure 3b caption should read W_V W_O^T (transpose dropped in second writing).

Author response:
Thank you, we will correct these.

Reviewer comment (Writing):
Page 28 d/dx in the Wronskian should be d/dt.

Author response:
Thank you, we will correct these.