Birth of a Transformer: A Memory Viewpoint

Summary

The paper studies in detail two-layer transformers and extend the setting of pure associative recall based on data in-context with mixing these tasks with tasks coming from a global birgram model. They then study these two-layer Transformers by freezing layers and probing.

Strengths

The paper studies an interesting problem i.e. in-context learning within Transformers, that is believed to be an important characteristic of language models in general. Based on recent work, they aim to mechanistically interpret / reverse engineer the two layer Transformer which in my opinion is an interesting direction to better understand and get intuition on how Transformers work. I like the mixed bigram data setup that the authors propose and study and I share the opinion of the authors that the tension between storing knowledge and quick adaptation / learning from in-context is very interesting and not explored. Progress in this direction is important.

Weaknesses

It feels that the paper is very rushed. I think therefore that the presentation can be vastly improved as well as results and analyses refined. Furthermore, I think that the proposed method of freezing the vast majority of layers in the network leads to very biased results and therefore results which might not hold when training on all weights. A couple of design decisions seem also quite arbitrary and it would be I think the job of the paper to justify and investigate them. See Questions. Therefore, I think that the paper is very promising but can and should be improved. Minor things / Language: Key, Query, Value Matrix must not be square and the positional encodings are usually not randomly initialized. Please at least comment on these design decisions. Also you do not use layer norm, I would also expect at least a comment on this. Please use LaTex in your Figures, quite hard to read the legends etc (maybe different background color as well?). The sum in Figure 2 over the positional encodings should be over t? Please describe more precisely Figure 3&4 in the caption. These are not concise i.e. it was not possible for me to understand what you are plotting what the Figures show. Also please increase the font size. At least one interesting citation is missing: https://arxiv.org/abs/2212.07677 The paper studies in details 1) single and two layer Transformer and provide evidence about copying in the first layer and 2) in-context learning in the second layer by gradient descent. This is, I believe, equivalent to the Hebbian-rule / associative recall with orthogonal inputs.

Questions

In general I think the papers focus on the tension between memory and in-context learning is great. Nevertheless, I think that a couple of very interesting experiments, ablations and analyses about design decisions are missing. 1) Can you provide analyses in both extremes i.e. when there are only "triggers" i.e. when all sequences can be learned by quick associative learning and the other way around. This should lead to quite different circuits that can be contrasted against each other. 2) Why do you include a single (linear) feedforward layer in the architecture? The problem should be solvable without (as studied in the induction head paper). If you want to include this layer, then I think the right thing to do would be to include it also after the second layer. 3) Why are you so heavily relying on the freezing of some layers? I actually ran your experiments (with only triggers) and I believe that my matrices do actually show very different behavior. I think that the freezing of the layers, forces you to focus on the superposition too much. If you wouldn't freeze your vocabulary and the input matrix, the model could (and I see this in my experiments) actually produce embeddings which leave token memory free to use for layer computations i.e. essentially concatenate x_i = w_E(z_i) + w_P(i) \approx [z_i, p_i] since the positional encoding is basically alternating 0 and 1 after some dimension. By that, the circuits and mechanism that you find could be very different and therefore not stable wrt. to design decisions. 4) Building on 3): I am quite convinced that all design decisions will change the circuits and mechanism of how the two layers solve your problem. Therefore I think it is very interesting and actually needed to investigate these and make suggestion what these design decisions control. E.g. Add or concat frozen/learnable random/sinusoidal PEs, control the dimension of the vocab and the token dimension, include feedforward layer, etc. 5) The theoretical insights on Learning Dynamics seem a bit vague and less supported by empirical study (didnt look in the appendix). As I understand, in-context learning kicks in later in training, so why study on the random init? At least comment on this shortcoming, if I understood this correctly. Minor things: 5) Can you please better your probing experiments, how exactly do you do that? Which vectors are probed? 6) What do you mean with discrete tokens line 81? 7) Equation 7 - its confusing for me why exactly these weights recover the induction head mechanism. Please elobatore, I can somehow see it but please make this more rigorous. 8) Please explain a bit better your data generation and the statistics. How large is the possibility of conflicts i.e. a -> based on global knowledge and -> based on in-context. 9) Figure 4 left. Why cant the in-context learning loss go to zero? Is it because of the conflicts? Please make the Figures a bit less crowded, hard to see all the lines. Would be great to have scaling plots wrt to K imo.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

I dont think the limitations of design decisions are well discussed or investigated. See Questions.

Summary

This paper provides a detailed analysis of how in-context learning behavior emerges in a simplified version of the transformer architecture on a toy task. This work can provide important insight into how in-context learning emerges in LLMs. The toy task is this: given a sequence of tokens of the form $\ldots, a, b, \ldots, a$, predict $b$. The token $a$ is a "trigger" token. If token $b$ (the "output" token) comes after the first occurrence of $a$, then the model must predict $b$ after the second occurrence of $a$. The model is expected to use in-context learning to predict the correct $b$. The rest of the sequence follows a bigram language model distribution. The authors zoom in on two aspects of the model: the "induction head" mechanism that implements in-context learning, and the "associative memory" mechanism in the key/query weight matrices that allows the model to learn the global bigram distribution. They give mathematical justification for the emergence of both mechanisms in their simplified transformer model trained on infinite data. They provide experimental results that provide evidence that these mechanisms are learned in a particular order, and that learnability depends on attributes of the training data distribution, such as the number of trigger tokens per sequence, or whether the trigger token types are randomized.

Strengths

This paper provides a very useful case study of how in-context learning behavior emerges in transformers. In-context learning is a hot topic, and it is important to understand architecturally how this behavior emerges, as well as the conditions that are conducive to it. This paper is a good step towards better understanding this phenomenon. I think their bigram task is an appropriate choice of case study. The authors analyze two mechanisms and provide mathematical and empirical justification for both. Originality: Good. I think the choice of task is very useful for the purposes of studying in-context learning. Quality: The mathematical analysis and experimental design appear to be sound. Clarity: The paper seems to be well-contextualized with respect to previous work. Significance: High. After reading this paper, I feel I have a better grasp on how in-context learning works in the transformer architecture, and this has important implications for any NLP applications that use transformers.

Weaknesses

**Edit:** I have read the rebuttal, and it addressed my biggest concerns. I have two major criticisms of this paper, which is why I did not immediately assign it a higher score: 1. The transformer architecture used in this paper is drastically simplified from real transformers used in LLMs. There is no layer normalization or dropout, and the transformer only has 2 layers. They use linear layers instead of feedforward layers. The key and value matrices are merged into one matrix. Many of the parameters are frozen during experiments and during gradient analysis. I think the paper requires a much more detailed, readable discussion justifying why the authors expect their analysis on this simplified transformer model to carry over to real transformers. I think there is some discussion scattered throughout the paper, but if so, I think it needs to be organized more clearly. I think the results are useful regardless of these simplifying assumptions, but the authors need to talk about this more. 2. Clarity. This was a very difficult paper to read, and this negatively impacted my ability to interpret the results. As I explain in more detail in the Questions section, a recurring issue I had while reading this paper was that information is frequently presented in *reverse* order; I often needed to read several lines, paragraphs, or sections ahead in order to clarify something I was confused about. I think the whole paper would benefit from a round of editing that alleviates these issues. See the Questions section for more details.

Questions

Most important questions: 1. 77, 90, 93, 97: Does the decision to ignore layer normalization make the findings of this paper less applicable to real transformers and LLMs? Same question for using only a single head. What about dropout? Same question about using linear transformations instead of feed-forward networks. 1. Fig 1: I had a very hard time understanding this figure, and it's essential for understanding the paper. I don't know what Attn1 and Attn2 are supposed to signify. In Layer 1 column 2, $w_1(a)$ and $w_E(b)$ aren't the only values in superposition there, right? It would also have $p_t$. It's not possible to completely isolate $w_E(b)$ from $p_t$, because they are in superposition in Layer 0, right? I think it would really help to give a symbolic representation of the query, key, and value functions at each layer. If I understand correctly, the inputs at Layer 1 are essentially $(x_t, t)$, where $x_t$ is the input token at timestep $t$. Then $\mathrm{query}(x_t, t) = t-1$ (assuming you can compute this with a linear layer), $\mathrm{key}(x_t, t) = t$, and $\mathrm{value}(x_t, t) = x_t$. This adds $x_{t-1}$ to the inputs of Layer 2. Then in Layer 2, $\mathrm{query}(x_t, t, x_{t-1}) = x_t$, $\mathrm{key}(x_t, t, x_{t-1}) = x_{t-1}$, and $\mathrm{value}(x_t, t, x_{t-1}) = x_t$. Then if $x_T = x_{t-1}$, this lets you predict $x_t$. Is this right? How do you ensure that this only applies when $x_T = a$, and not other token types in the vocabulary? 1. Lemma 1: Could you include a proof sketch in the main text? Could you include a longer discussion of the limitations of your simplifying assumptions in Lemma 1 and how they apply to real transformers? The main problem with this section is that I cannot verify that Lemma 1 is true, and that it applies to real transformers in the way you claim that it does. Won't it totally change the dynamics if the other weights are not fixed? I think it's probably a useful result either way, but you should talk about this more. 1. 197: Could you justify this decision a bit more? How do we know this doesn't drastically change the situation from real transformers? 1. 308: Can you include a proof sketch for Lemma 2? Other questions: 1. Can you clarify which key-value associations you expect the linear layers to learn? Is it the global bigram distribution? 1. 51: If all the examples are generated by a single bigram language model, isn't all of that "global" knowledge? That is, how will you know if the model isn't just memorizing the training data? Are you counting on the fact that the transformer isn't big enough to memorize so many examples? Do you test on out-of-distribution examples that verify that this is the case? 1. 98: What about tied embeddings? 1. 199: I think this assumption is pretty reasonable. But I think it's harder to understand how the induction head would work without a separate query matrix. Could you spend a little time explaining how this would affect the construction? 1. In general it would be very helpful to explain in detail how to design an induction head under these constraints. 1. 208: Again, it would be useful to discuss how W_F would implement this. 1. 219: How? 1. 220: Where do $p_t$ and $p_{t-1}$ come from? Are they sinusoidal positional encodings? 1. Eq 9: Can you explain in words what this is doing? 1. Fig 5: Why is it easier to learn 5 random q than 1 random q? It would be helpful to be reminded that K is the number of triggers per sequence, not the number of token types used as triggers overall. 1. 310: What is $E[x|y = k]$? What is $\hat{p}_W$? Clarity issues: 1. 43: At this point, I was wondering if you would supplement your empirical analysis of the training dynamics with a mathematical explanation. It would be useful to point out earlier that you will do this. 1. Eq 1: Having $x_t$ on both sides of the equation is confusing. I'm not a big fan of the $:=$ notation. 1. How would the induction head mechanism work in the presence of multiple instances of a in the past context? 1. 120: These sentences seem out of order and are hard to read. It's not clear if $b$ is a variable or just a tag to distinguish $\pi_b$. I don't understand what $\pi_o$ is. I was confused by the difference between $\pi_b$ and $\pi_o$. 1. 123: What is the variable $n$? 1. 124: How is $K$ chosen? Is it randomly sampled or set to a constant value for the experiments? 1. Using i and j as variables for token types instead of input positions is confusing, especially since k is used as an index 1. 126: What is the first token in the sequence? How is that sampled? 1. 126: It would be helpful to provide an intuitive explanation in words of the process that this equation implements. 1. Where is $\pi_u$ used? (I see it is used at 131, but this seems out of order.) 1. 129: What is the tiny Shakespeare dataset? 1. 133: Referring to training details in Section 5 out of order makes this part hard to read. Could you move the training details earlier? 1. 134: Is this on the training data, or on a held-out test set? Do you use a validation set? 1. 134: It would be easier to read this results if they were in a table. 1. Fig 2: The labels for the axes on the first plot are missing. It would be helpful to label the y axes. For the left plot, which version of the dataset was this trained on? Is there a reason you can't show the first layer for both models? I don't understand the significance of the red and green boxes. What are the "previous occurrences of the query"? It would help to highlight the triggers and outputs in the axis labels. 1. 134: At this point I was confused about the significance of using fixed vs. random triggers. But after re-reading 124-126, I understand that for random triggers, the set of triggers is randomly sampled for every sequence, so *every* token type has the potential to be a trigger. The only way they are distinguished as triggers in the training data is that the same bigram appears twice in the same sequence. But this could also happen by chance for non-trigger tokens, right? It would be helpful to add a discussion of this to the main text. 1. 153: What does this notation mean? 1. 159: What does O mean? 1. Eq 5: What does z range over? 1. I didn't realize until Eq 6 that $w_E(z)$ means the $z$th column of $E$. Could you define this notation beforehand? 1. 213: I see that this section answers some of my questions above. It would be useful perhaps to introduce or mention this earlier in some way. 1. 238: Just to clarify, you're using a cross-entropy language modeling loss function, right? I see this is partly answered at 254. 1. 247: In the equation for $W_{*}$, what are $v_j$ and $u_i$? 1. Eq 9: What is $M$? How are the $(i, j)$ picked? 1. 307: This seems to answer one of my earlier questions.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

As mentioned above, I think the authors should discuss the limitations of their simplifying assumptions about the transformer architecture more.

Summary

Given the blackbox that large language models are, this paper tries to use a small simplified view of a 2 layer transformer model, and a synthetic task to understand how global and in-context language statistics are learned by the transformer model. By freezing specific layers, The authors show how the memory recall and in-context accuracy varies.

Strengths

The paper presents a nice setup with the simplified transformer model, and the synthetic task for analysis of transformer architecture. The current analysis are great, but just the start and the foundation laid could be useful for future explorations. The proposal around how the global memory and the in-context induction work is supported (partially) through the experiments by freezing various layers. With the experiments, they are able to show that the global bigram statistics are learned faster and the induction heads for in-context bigrams are developed in subsequent training steps. Additionally, the conclusion that having diverse training distribution leads to better generalization is great. So this is a good started to understanding the transformer LLM blackboxes, and the results may apply on not just text, but other modalities as well.

Weaknesses

- The results are still very preliminary. The transformer model is simplified, and so is the dataset. It is unclear whether a larger transformer model would use a similar paradigm for learning. Essentially, it might be possible that the results do not extend to larger models and additional experiment needs to be done to verify that. - Transformer models have a lot of other components, and it would be nice to study how they impact the memory of these models.

Questions

- In the multi-head attention, was any experiment perfomed to understand whether for a fixed model dimension, is it better to have larger number of heads or smaller number and how that affects memorization? - Was there any experiment performed where we just have attention (no feed forward layer), and how that affects global memorization?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Discussions could mention that this analysis may not extend to larger versions of the model, or provide empirical proof that it does.

Summary

This paper studies the dynamics of how induction heads emerge in LLMs during training. The authors describe a simple synthetic task to test their hypotheses on, outline a plausible implementation of an induction head based on associative memory, and describe both empirical and theoretical observations.

Strengths

* The paper is well-written and generally easy to follow. Motivations, conclusions, and analysis are crystal clear. * The authors study an interesting question: what mechanism do LLMs implement induction heads with, and how might it emerge during training? * The methodology is clean and nice! I like the inclusion of "memory recall probes," as they are a direct measure of what you're claiming. The results are also appropriately framed in the context of the merits & limitations of the experimental setup.

Weaknesses

* Can you better explain the significance of the results? My interpretation is that you framed the $W$ matrices as associative memories and found an order in which they seem to be learned. I'd love to hear what you think the "so what" of this is! What interesting things can we do with this understanding? What more can we learn? Can we make the induction heads better? Implement them manually? Do you expect the learnings here to help us understand larger, more complex models? * It's sort of implied that the framing of induction heads as associative memory is new. What are the other competing mental frameworks? How does this one compare? * It'd be interesting to see experiments that vary the dimensionality of the vectors/matrices. In my experience, superposition is rampant, and LLMs often try to stuff a lot into a not-so-large $d$. How do you expect the findings to change when you can't assume near-orthonormality? Other than a degradation in performance? Minor things: * Figure 1 is nice! It'd suggest labeling the words in the caption with colors or letters, so it's immediately clear what refers to what.

Questions

* Did you check that the actual weights learned by your toy model match the solution you made in any meaningful way? Maybe exactly, modulo a simple transformation or something? Did you find any interesting surprises? * Did you have a chance to test any hypotheses about superposition experimentally, even preliminary results? If so, what did you find? Any unexpected things? Negative results ok! I'd imagine it's quite relevant, and is a better representation of what really happens in transformers.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

n/a

Summary

Authors perform an in-depth study of the toy case of learning associate recall task using causal Transformers with the aim to understand the emergence of in-context learning abilities during training. Informally, they propose a modified bigram distribution where after sampling a sequence, for a special set $Q$ of pre-determined "trigger" tokens, at every occurance of a token $q$ from $Q$, they replace its following token with the token that followed $q$ at its first occurrence in the sequence. I.e. $[...q\ r ....q\ s] \mapsto [...q\ r ....q\ r]$. Hence, for $q$, the model is required to determine $r$ first from the context and, at every following occurance of $q$ output $r$. The authors prove that a simplified 2-layer Transformer (the only non-linearity is softmax) can learn this task via gradient descent and verify this empirically.

Strengths

1. The problem is well-motivated as large language models demonstrate impressive in-context learning abilities and hence it is beneficial to understand how this ability develops. 2. The theoretical statements are non-trivial, interesting and not straightforward to prove.

Weaknesses

1. The associative recall task has been studied in the context of in-context learning of attention-based models before and its not fully clear what novel contributions the authors have made. 2. Given the simplicity of the toy task/model the practical implications and the generality of the findings are unclear (other tasks / architectures).

Questions

.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

There is no discussion of the limitations.

response 1 to rebuttal

Thank you for addressing some of my concerns - after going through your responses to my review (and to other reviews) I am increasing the score.

Response

Thank you for your response. Weakness 1, MIQ1, MIQ4, OQ3, OQ4: I do agree that simplicity is beneficial for analysis, but the simplified version needs to be tethered to unsimplified transformers in some way, or else the findings of the paper are no longer relevant. I still believe the paper needs at minimum a longer, focused discussion of each of these limitations and how the simplified and unsimplified versions connect. The additional empirical results you have mentioned in the rebuttal will be helpful in arguing this. MIQ2: How close was my original interpretation? And there is no limit to the number of vectors that can be in superposition, not 2 as suggested in the figure, right? MIQ3: Some experimental results demonstrating that the embeddings first converge quickly to an optimum would be helpful here, and in addressing Weakness 1. OQ1: > While it is possible that a transformer with very large MLP layers could memorize all sequences in order to guess the correct output token after a trigger, this would likely be intractable since the number of such possible sequences grows exponentially with sequence length. This is a good point and worth mentioning in the paper. OQ3: Can you explain this in more detail? OQ6: Thanks. This ties into Weakness 2. OQ7: Thanks, this is worth mentioning in the paper. OQ8: Ok, somehow I wasn't able to find this definition in the paper.