The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains

For two states, the transition matrix can be represented as $\mathcal{P}=\begin{pmatrix}1-\alpha&\alpha\\\beta&1-\beta\end{pmatrix}$, where $\alpha,\beta$ are iid random variables from the uniform distribution on $[0,1]$. The stationary distribution is $\pi = \begin{pmatrix} \frac{\beta}{\alpha + \beta}, & \frac{\alpha}{\alpha + \beta} \end{pmatrix}$, which is feasible to analyze. While there exist results such as Chafai et al that characterize the stationary distribution in the limit as the number of states approaches infinity, we do not believe analogous results exist for any constant number of states greater than two.

We can prove a version of Lemma 3.1 for any number of states (k) as long as we consider a different distribution for transition matrices: specifically, a mixture of the distribution where the unigram strategy is optimal, and the distribution where the unigram strategy is as bad as guessing randomly. The proof only requires different versions of lemmas C.1 and C.2, and actually adds more intuition to our stage-wise learning story. Specifically, in the first phase of learning, the contribution to the gradient from the unigram-optimal distribution dominates, but in the second phase, the other component (from the distribution where unigrams are useless) is dominant. We would like to add this variation of Lemma 3.1 to our camera-ready version.

Thank you for critically reading our work and helping us improve its presentation. We greatly value your in-depth review!

We will first discuss your most significant concerns with our work:

In-Context Learning

You're right that we don't spend much time in our paper discussing what "in-context learning" means, and indeed, there isn't a single consensus definition! The term "in-context learning" was (we believe) introduced in the GPT-3 paper ([Henighan et al., 2020]) to refer specifically to the phenomenon where LLMs have improved few-shot performance with increasingly many examples in their context. Xie et al., meanwhile, does indeed explore scenarios where there is a distribution shift between training and inference; however, in the broader literature, ICL is not consistently tied to distribution shifts (see, e.g. [Garg et al., 2022] for an influential study of ICL that does not require a distribution shift in its definition). Meanwhile, [Elhage et al., 2021], which was roughly concurrent with Xie et al., and the follow-up work [Olsson et al, 2022] -- use ICL very broadly to refer to LLMs being better at predicting tokens that appear later in a sequence. The way we use the term is somewhat narrower than this; as we state in the Introduction: we think of ICL as when language models are "incorporating patterns from their context into their predictions." In other words, it is when models learn from patterns in their context; few-shot learning being a special case of this.

The task we focus on, ICL-MC, is an in-context version of a classic learning task, Markov chain inference. Each test-time sequence has zero probability of appearing during training, because of the continuous nature of the distribution over chains. During pre-training, the transformer learns to perform the learning task in-context. We don't only focus on the case where there is a distribution shift between train and test -- as discussed above, we think ICL is interesting even in the absence of distribution shift. (Though we do explore distribution shifts in the paper: see Figures 4 and 11.)

We would be happy to elaborate on our interpretation of ICL in the camera-ready version of the paper.

Answers to questions

1. This choice was made to facilitate a theoretical analysis of the observed phenomena. In particular, results such as Lemma C.2 become feasible only when the number of states is 2, or other simplifying assumptions are made. However, in the experiments, we have verified that our insights transfer to Markov Chains with a larger number of states (see, for example, Figure 7 in the Appendix). We will add a discussion of Markov chains with more states, including a reference to figure 7.

Based on your concerns, we have decided to add back in a proof we decided not to include, which means that by making only small changes to the submitted version, we are able to extend our main results to any number of states. This requires using a distribution like that in experiment on the left of figure 4. Due to lack of space, we defer the details of this to a comment.

2. Thank you for pointing out this line as being out of place. The line (edited for clarity) will be moved to the discussion of limitations. It intended to communicate the difficulties with the distribution of the stationary distribution for $k>2$. Due to lack of space, we defer further elaboration to a comment.

3. We chose this to make the setting closer to the next token prediction seen in practice. The marginal distributions of the tokens generated from a markov chain approach the stationary distribution exponentially fast no matter the initial distribution.

We have ran experiments for other somewhat natural starting distributions (such as uniform) and found no meaningful differences.

4. Yes. We will add the following mathematical definition: the unigram strategy's probability of the token at position $t+1$ being $j$ is $\frac{1}{T}\sum_{i=0}^T \mathbb{1}[x_i=j].$

5. The bigram strategy is computed by counting the frequency of pairs of states in the context of each sequence. This corresponds to the formula in line 157.

The graphs on the left in figure 3 show the KL-div loss using the instance-dependent ground-truth transition probability matrix. Because the context size is long (100 tokens), and we average over 1024 test samples, these two measures are almost identical (as we would expect).

6. No, the KL div between the unigram and bigram strategies is not due to finite sampling. The unigram strategy is a flawed strategy (it ignores the information we get from the relative order of the tokens), for almost all markov chains it will get higher expected loss than the bigram strategy, even for small context sizes.

7. The bigram strategy is a near optimal strategy (it approaches the optimal strategy exponentially fast as the context length grows). This is true regardless of the model used.

8. We are not entirely sure how to interpret this question and would appreciate clarification.

Sampling the first token by the stationary or uniform distribution does not meaningfully affect any empirical observations. The marginal distributions of the tokens after the first will always approach the stationary distribution exponentially fast.

Empirically, the model starts by representing the unigram strategy and later the bigram strategy. In both cases it is not obvious that the model needs to fit the unigram strategy at any point (after all, the unigram strategy performs worse than the bigram strategy), and so it must be only due to the inductive biases of the model and optimizer that this happens.

9. No. The transformers were initialized with the default Pytorch initializations (Gaussian with mean 0), we will clarify this in section B.

10. This is a typo.

We hope we have managed to address your concerns and would be grateful if you adjusted your score accordingly.

4.1. We use a dimension of $3k$ to have a simple and intuitive construction. In practice, models can learn with a far smaller internal dimension. We can add discussion of this to the main text.

4.2. We apologize for the confusing notation with regard to the matrix definitions, it will be improved. $v^{(1)}$ is of dimension $t\times d$, where $d=3k$. We will make sure that the dimensions of each submatrix used to define $v^{(1)}$ and all other matrices are all specified in the notation.

4.3. $x_i$ refers to the token at position $i$. $e$ is used to represent the $1$-hot embeddings, $e_i$ is a vector that is $0$s besides at position $i$ where it is $1$. Similarly, $e_{x_i}$ is all zeros except at position $x_i$ where it has a $1$. $e_{i-1,j}$ is the $j$th index of the vector $e_{i-1}$, hence $e_{i-1,j}=\mathbb{1}[i-1=j]$. We will do an extra pass over all of this notation to clarify and improve it.

4.4. and 4.5. We thank you for pointing out these typos. It should say: $\text{softmax}(\text{mask}(A))\_{i,j}\approx \frac{\mathbb{1}[x_{j-1}=i]}{\sum_{h=1}^i \mathbb{1}[x_{h-1}=i]}$

Fixing these typos, we get the result $Attn_2(e)\_{i,j+2k} =\frac{\sum_{h=1}^{k}\mathbb{1}[x_{h-1}=x_i]\mathbb{1}[x_h=j]}{\sum_{g=1}^i\mathbb{1}[x_{g-1}=x_i]}$ Which is more correct, since this is the bigram probabilities, instead of just the bigram statistics. That is, this is the empirical approximation of $P_{x_i, j}$.

Doubly stochastic transition matrices are those for which the stationary distribution is the uniform vector, or those for which the unigram strategy and the uniform strategy get the same loss. If we want to observe the inductive biases of the model when there is no signal encouraging the unigram strategy, then sampling from doubly stochastic transition matrices is natural.

Finally, we also considered the distribution where each row in the transition matrix is the same, resulting in the unigram and bigram strategies having the same loss.

The mathematical definition of the mixed distribution in the graph on the left of figure 4 is as follows: with 75% chance, choose a uniformly random doubly stochastic transition matrix, otherwise make each row of the transition matrix the same vector, chosen from a flat Dirichlet distribution.

To create our minimal model, we started from a two layer attention-only disentangled (see Elhage et al) transformer (using relative positional embeddings) and iteratively simplified parts that empirically did not affect the training dynamics. In our construction, the first layer only attends to positional embeddings, and the second layer ignores positional embeddings, so we set the first layer key matrix and the positional embeddings from the second layer to zero. Then we set the value matrices and query matrices to the identity in both layers. In the experiments, all optimizations such as layer norms and weight decay were removed, and the optimizer used was SGD. With all of these changes, the overall training dynamics were not changed much at all, depending on hyper parameters they could speed up training, but the same loss curve and phases were observed. To make analysis of gradient descent on the minimal model feasible, the softmax on the second layer had to be removed, which did make the phase transition less sharp.

Thank you for the in-depth review and feedback!

You point out that there is some misalignment between our experimental story and our proofs, and you see that as our paper's major weakness. We agree that there is misalignment, due to the difficulty of analyzing SGD on multi-layer Transformers. However, we believe our analysis is very relevant to the experimental observations about simplicity bias and phase transitions. In particular, Lemma 3.1 exhibits simplicity bias by showing that the parameters after the first phase compute a solution that is a weighted version of the unigram solution, and the second phase computes a mixture of unigram and bigram. Also, the magnitude of the signal for the second phase is a factor $T$ lower than the first step, indicating that the second phase requires more time for "signal accumulation". We agree that this does not explain the sharpness of the phase transition, but it does provide theoretical intuition for the existence of multiple phases. We are happy to add these aspects of the story to the discussion following Lemma 3.1 in the text; for these reasons, we don't think the extent of misalignment is a serious weakness in the paper.

We address your various minor concerns below:

1. Thank you for pointing out this line as being out of place. A version of the line (edited for clarity) will be moved to our discussion of limitations. The line intended to communicate the difficulties with understanding the stationary distribution when $k>2$.

2.1 and 2.2. The goal of the minimal model is to capture the dynamics of the full transformer, while being amenable to theoretical analysis. Essentially all analyses of multilayer transformers have relied on simplifying out many parts of the transformer. While simplifying the transformer we carefully checked that the dynamics we wished to understand were preserved as much as possible. Due to space limitations, we have deferred a more in depth explanation of the choices in the minimal to a comment.

2.3. This is correct. Unfortunately, the analysis does not seem to be feasible with softmax in both layers.

3 and 9. Generating each row of the transition matrix iid from a flat Dirichlet distribution is equivalent to uniform distribution over all transition matrices. We chose to focus primarily on this distribution because it is the most natural.

The only mention of doubly stochastic matrices in the theory (lemma C.3) was mistakenly left in after the surrounding context was removed.

Due to space constraints, we defer definitions of the other distributions to a comment.

4. Thanks for your careful proofreading. Due to space constraints we defer the specific changes and improvements made in response to a comment.

5. The trained transformers used an internal dimension of 16. The specific parameters after training always vary randomly due to initialization and training samples. Figure 8 shows that the transformers are implementing the same algorithm as our construction, with each layer performing the same functions.

6. Thanks for pointing out this oversight, we have added the full proof in (it's rather short and simple).

7. That's an interesting question! Unfortunately, the details of the loss landscape for even simplified transformers are hard to precisely characterize with current theoretical tools. Among other difficulties, a barrier to transformer optimization results is analyzing cases where the input to softmax isn't close to zero, or one hot.

Our intuition is that the unigram is not exactly a saddle point: during the plateau period, there is very slow continual improvement in the loss. We suspect the gradient at the unigram solution is simply small relative to the change needed to move towards the bigram solution.

8. Line 241 is a holdover from a previous version, and no longer applies to the experiments in the paper, we apologize that it made into this version and will remove it.

On the left in figure 3, both the minimal model and transformer model reach their minimum loss after seeing around 80,000 training sequences. We do believe though that one shouldn't put too much weight into the fact that these took the same time to train.

We believe this comparison still makes sense. The minimal model is designed to be able to express everything the transformer expresses when being trained on the task. Experimentally, it goes through the same phases. In the right most graphs in figure 3, both models quickly learn to use the unigram strategy, before eventually adopting the bigram strategy. Since this task does not require the full expressive power of a transformer, it is not too surprising that a model with many fewer parameters is sufficient to gain insight into the training dynamics.

10. The most compelling piece of evidence that the minimal model actually goes through these phases are the experimental results shown in figures 3 and 4. Experimentally this also works for standard small Gaussian initializations. In the conclusion, we will add a mention of the limitation of this being a two step analysis. Unfortunately, existing tools are insufficient to analyze dynamics involving softmaxes (specifically when the inputs to the softmax are neither small nor dominated by a single index). For lemma 3.1 we were forced to limit our analysis to two steps. However, as we mention on lines 579-580, the uniform component of $W$ does not contribute to the gradient of $v$, so even if we initialized $W$ to any uniform non-zero initialization, $v$ would still not change in the first step. We chose the specific step sizes to allow the theory to show how the model can go through the unigram and bigram phases, which happens in practice over the course of many small steps.

We hope we have managed to address your concerns. If you think we have adequately done so, we would be grateful if you adjusted your score accordingly. Thank you!

Thank you for your review and questions. Here we address the main weaknesses and questions raised by you.

Impact and related works

Questionable impact since this topic has been explored heavily in the past years, with other similar tasks and toy models existing, showing similar results... What is the key differentiating contribution of the work in comparison with the myriad of other works in this space?

As we discuss in our related work section, there are indeed prior works which introduced synthetic settings for exploring ICL:

• Garg et al., 2022 (and follow-ups) train models to learn linear functions (and other simple function classes) in a few-shot setting.
• Xie et al. 2022 describe a few-shot learning setting in which each document consists of examples generated by a hidden Markov model.
• Finally, most related to our work is Bietti et al. 2023, in which all sequences are generated from a single Markov chain (which doesn't need to be learned in-context), but certain 'trigger' tokens are automatically followed by sequence-specific tokens which need to be learned in-context.

We believe that our task, in-context learning of Markov chains (and the higher-order generalizations thereof), is a valuable contribution in the context of this rich literature. (No single synthetic task is going to capture all of the scientifically interesting aspects of in-context learning.) In particular:

• It is very natural and simple to describe.
• It elicits the formation of induction heads in networks trained on the task. (It is a particularly natural setting for studying their emergence.)
• There are multiple strategies of various levels of sophistication which can be used to solve the task to varying degreees of success (unigram, bigram, etc.), and which networks pass through in stages during training.

Moreover, the results we obtained by studying this task experimentally and theoretically are novel. Our key findings -- the formation of statistical induction heads, the stage-wise phase transitions between in-context solutions (unigram, bigram, trigram, ...) on the way to success, the arguments that simplicity bias may delay the formation of the correct in-context solution, and finding that the second layer of the model is learned before the first layer, rather than the other way around -- are all original to our work.

Given that you say there are lots of very similar works, we suspect you may be mistakenly assuming that some concurrent works were actually prior works. Indeed, at the same time we released our preprint, there were several simultaneous papers (all with pre-prints posted within a single month) that shared aspects of our setup and/or scientific focus. Please see the "Concurrent Works" paragraph in our Related Work section for a discussion of these papers and their relation to our work: Akyürek et al., 2024, Hoogland et al., 2024, Makkuva et al., 2024a, and Nichani et al., 2024. There have also been subsequent works in this space, including Rajaraman et al., 2024a, Rajaraman et al., 2024b and Makuva et al., 2024b.

We want to make sure that you are following the scientific best practice and judging our work's significance in relation to prior works, not concurrent works. If you mistakenly thought the above concurrent works were prior, we understand! There has indeed been a burst of activity on this topic recently. In that case, we hope you will reassess our work in its proper context, and adjust your score accordingly. If you were referring only to actual prior works, we hope you find the above contextualization of our work convincing and reassess accordingly; if not, we would appreciate some elaboration on the works you are referring to.

Real-world relevance

What is the conclusion of the work that you believe is transferable to real-world ICL setting. What n-gram statistics would that follow?

N-gram statistics are useful for predicting real-world natural language text. Indeed, historically, language models were often simply n-gram predictors! It is also useful for models to learn in-context bigram (and more generally, n-gram) statistics. For instance, the writing style of a particular document is connected to the n-gram statistics of that document.

The most concrete connection the toy MC-ICL setting has to real world ICL settings is induction heads, which have been shown to have importance to in-context learning in LLMs in past work (Olsson et al, 2022). In the MC-ICL setting, the emergence of induction heads corresponds to the phase transition in loss, just like the phase transition found in the ICL ability of transformers hypothesized to be caused by induction heads emerging.

Additionally, compared to other ICL tasks studied in the past (such as linear of logistic regression) we believe that n-grams are a better model for natural language specifically.

Thank you for answering my questions. I admit that some, but definitely not all of the papers in mind may be concurrent works. I have read the other reviewer's comments, and in general it seems like while there is an agreement that this is a well-done piece of scientific literature, the impact is questionable, as the toy-setting field particularly when it pertains to ICL is definitely crowded. Finding new angles to attack the problem is a worthwhile pursuit, but at times may distract from perhaps other more important and less-studied issues.

As I see there are better reviewers than me to assess this work, so I will raise my score from a 4 to a 5 while keeping my confidence low, and leave it up to them to reach a consensus.

We thank the reviewer for their helpful review. Here we address the questions/concerns raised by the reviewer.

1. Do you think that an SGD analysis similar to the one carried out for the minimal model would carry over to a more complex architecture, closer to the actual transformer? If not, what are the main issues with working out such a analysis for, for example, the simplified transformer model of Equation (1)?

We believe that the analysis would mostly carry over to a more complicated architecture with one exception. Specifically, the softmax in the second layer attention proved intractable to analyze at the same time as the softmax in the first layer. Analyzing softmax is a common difficulty in transformer analyses, and why most theoretical results apply to linear attention only, we are happy to have managed to include the softmax in the first layer. Empirically, the training dynamics of the minimal model are very similar to that of the full model, the primary difference being that the phase transition isn't as sharp in the minimal model. With the softmax in the second layer added in, the phase transition becomes sharp like in the full model, but the tools do not currently exist to analyze this (specifically nested softmaxes on inputs not close to 0 or infinity).

To create our minimal model, we started from a two layer attention-only disentangled transformer (using relative positional embeddings) and iteratively simplified parts that empirically did not affect the training dynamics. In this context, disentangled transformer (introduced in Elhage et al) is like a normal transformer, except the output of each layer is appended after the residual. In our construction, the first layer only attends to positional embeddings, and the second layer ignores positional embeddings, so we set the first layer key matrix, $W_k^{(1)}$, to $0$, and removed positional embeddings from the second layer ($v^{(2)}=0$). Then we set the value matrices and query matrices to the identity in both layers. In the experiments, all optimizations such as layer norms and weight decay were removed, and the optimizer used was SGD. With all of these changes, the overall training dynamics were not changed much at all, depending on hyper parameters they could speed up training, but the same loss curve and phases were observed. To make analysis of gradient descent on the minimal model feasible, the softmax on the second layer had to be removed, which did make the phase transition less sharp.

We believe an analysis that includes this softmax could confirm additional insights into why the phase transition is so sharp, and is a potential direction for future work.

2. Your model architecture uses relative positional embeddings, as opposed to most state-of-the-art transformer models. Do you think that your analysis would be easily extended to a model with absolute positional embeddings?

Empirically, the absolute positional embeddings add noise, but do not fundamentally change any of the results. The analysis might be a bit messier, but we would not expect it to be fundamentally different. Additionally, many state of the art models use non absolute positional embeddings, including the Llama series of models (1, 2, and 3).