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

Authors: ezra  edelman
nikolaos  tsilivis
benjamin l edelman
eran  malach
surbhi  goel
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 19:04:38

Abstract: Large language models have the ability to generate text that mimics patterns in their inputs. We introduce a simple Markov Chain sequence modeling task in order to study how this in-context learning capability emerges. In our setting, each example is sampled from a Markov chain drawn from a prior distribution over Markov chains. Transformers trained on this task form \emph{statistical induction heads} which compute accurate next-token probabilities given the bigram statistics of the context. During the course of training, models pass through multiple phases: after an initial stage in which predictions are uniform, they learn to sub-optimally predict using in-context single-token statistics (unigrams); then, there is a rapid phase transition to the correct in-context bigram solution. We conduct an empirical and theoretical investigation of this multi-phase process, showing how successful learning results from the interaction between the transformer's layers, and uncovering evidence that the presence of the simpler unigram solution may delay formation of the final bigram solution. We examine how learning is affected by varying the prior distribution over Markov chains, and consider the generalization of our in-context learning of Markov chains (ICL-MC) task to $n$-grams for $n > 2$.

Main paper: 
 introduction: 
 Large language models (LLMs) exhibit a remarkable ability to perform in-context learning (ICL) from patterns in their input context [12, 16] . The ability of LLMs to adaptively learn from context is profoundly useful, yet the underlying mechanisms of this emergent capability are not fully understood.In an effort to better understand ICL, some recent works propose to study ICL in controlled synthetic settings-in particular, training transformers on mathematically defined tasks which require learning from the input context. For example, a recent line of works studies the ability of transformers to perform ICL of standard supervised learning problems such as linear regression [3, 20, 26, 41] . Studying these well-understood synthetic learning tasks enables fine-grained control over the data distribution, allows for comparisons with established supervised learning algorithms, and facilitates the examination of the in-context "algorithm" implemented by the network. That said, these supervised settings are reflective specifically of few-shot learning, which is only a special case of the more general phenomenon of networks incorporating patterns from their context into their predictions. A few recent works [8, 42] go beyond the case of cleanly separated in-context inputs and outputs, studying in-context learning on distributions based on discrete stochastic processes.The goal of this work is to propose and analyze a simple synthetic setting for studying ICL. To achieve this, we consider n-gram models [11, 15, 37] , one of the simplest and oldest methods for language modeling. An n-gram language model predicts the probability of a token based on the preceding n -1 tokens, using fixed-size chunks (n-grams) of text data to capture linguistic patterns. Our work studies ICL of n-gram models, where the network needs to compute the conditional probability of the next token based on the statistics of the tokens observed in the input context, rather than on the statistics of the entire training data. We mainly focus on the simple case of n = 2; i.e., bigram models, which can be represented as Markov chains. We therefore consider ICL of Markov chains (ICL-MC): we train two layer attention-only transformers on sequences of tokens, where each sequence is produced by a different Markov chain, generated using a different transition matrix (see Figure 1 (left)).By studying ICL-MC, we are able to replicate and study multiple phenomena that have been observed in ICL for LLMs, and identify new ones. We demonstrate our findings using a combination of empirical observations on transformers trained from scratch on ICL-MC and theoretical analysis of a simplified linear transformer. Our key findings are summarized below:(1) Transformers learn statistical induction heads to optimally solve ICL-MC. Prior work studying ICL in transformers revealed the formation of induction heads [18] , a circuit that looks for recent occurrence(s) of the current token, and boosts the probabilities of tokens which followed in the input context. We show that in order to solve ICL-MC, transformers learn statistical induction heads that are able to compute the correct conditional (posterior) probability of the next token given all previous occurrences of the prior token (see attention patterns in Figure 2 ). We show that these statistical induction heads lead to the transformer achieving performance approaching that of the Bayes-optimal predictor.(2) Transformers learn predictors of increasing complexity and undergo a phase transition when increasing complexity. We observe that transformers display phase transitions when learning Markov chains-learning appears to be separated into phases, with fast drops in loss between the phases. We are able to show that different phases correspond to learning models of increased complexity-unigrams, then bigrams (see Figure 1 )-and characterize the transition between the phases. We also consider the n-gram generalization of our setting, where the next token is generated based on the previous n -1 tokens, and observe a similar multi-stage learning process.(3) Simplicity bias may slow down learning. We provide evidence that the model's inherent bias towards simpler solutions (in particular, in-context unigrams) causes learning of the optimal solution to be delayed. Changing the distribution of the in-context examples to remove the usefulness of in-context unigrams leads to faster convergence, even when evaluated on the original distribution.(4) Alignment of layers is crucial. We show that the transition from a phase of learning the simple-but-inadequate solution to the complex-and-correct solution happens due to an alignment between the layers of the model: the learning signal for the first layer is tied to the extent to which the second layer approaches its correct weights.In-Context Learning. In [13] , the authors discuss how properties of the data distribution promote ICL. Xie et al. [42] suggest a Bayesian interpretation of ICL and studies how ICL emerges when the training distribution comes from a Hidden Markov Model (HMM). Abernethy et al. [2] study the ability of transformers to segment the context into pairs of examples and labels and provide learning guarantees when the labeling is of the form of a sparse function. The work of Bietti et al. [8] studies the dynamics of training transformers on a task that is reminiscent of our Markov chain setting but has additional complexities. Instead of drawing a fresh Markov chain for each sequence, in their task all sequences are sampled from the same Markov chain; after certain 'trigger' tokens, the following 'output' token is chosen deterministically within a sequence. Thus, successful prediction requires incorporating both global bigram statistics and in-context deterministic bigram copying, unlike in our setting where the patterns computed by statistical induction heads are necessary and sufficient. As in our work, the authors identify multiple distinct stages of training and show how multiple top-down gradient steps lead to a solution.Induction Heads. Elhage et al. [18] relates ICL with the formation of induction heads, subcomponents of transformers that match previous occurrences of the current token, retrieving the token that succeeds the most recent occurrence. Reddy [34] studies the formation of induction heads and their role in ICL, showing empirically that a three layer network exhibits a sudden formation of induction heads towards solving some ICL problem of interest. Bietti et al. [8] study the effect of specific trigger tokens on the formation of induction heads.Phase Transitions. It has been observed in different contexts that neural networks and language models display a sudden drop in loss during their training process. This phase transition is often related to emergence of new capabilities in the network. The work of Power et al. [32] observed the "grokking" phenomenon, where the test loss of neural networks sharply drops, long after the network overfits the training data. Chen et al. [14] shows another example of a phase transition in language model training, where the formation of specific attention mechanisms happen suddenly in training, causing the loss to quickly drop. Barak et al. [7] observe that neural networks trained on complex learning problems display a phase transition when converging to the correct solution. Several works [25, 27] attribute these phase transitions to rapid changes in the inductive bias of networks, while Merrill et al. [29] argue that the models are sparser after the phase change. Schaeffer et al. [35] warn that phenomena in deep learning that seem to be discontinuous can actually be understood to evolve continuously once seen through the right lens. Simplicity Bias. Various works observed that neural networks have a "simplicity bias", which causes them to "prioritize" learning simple patterns first [5, 39] . The work of Kalimeris et al. [23] shows that SGD learns functions of increased complexity, first fitting a linear concept to the data before moving to more complex functions. [36] shows that the simplicity bias of neural networks can sometimes be harmful, causing them to ignore important features of the data. Chen et al. [14] demonstrate the effect of simplicity bias on language tasks that require understanding of syntactic structure. Abbe et al. [1] provide a theoretical framework for understanding how the simplicity of the target function can govern the convergence time of SGD, describing how simple partial solutions can speed up learning; in contrast, in our setting, the unigram solution appears likely to be a distractor which delays learning of the correct solution.Concurrent works. In parallel to this work, there have been a number of papers devoted to the study of similar questions regarding in-context learning or Markov chains: Akyürek et al. [4] empirically compare the ability of different architectures to perform in-context learning of regular languages. Their experiments with synthetic languages motivate architectural changes which improve natural language modeling in large scale datasets. Hoogland et al. [21] observe similar stage-wise learning behaviors on transformers trained on language or synthetic linear regression tasks. Makkuva et al. [28] study the loss landscape of transformers trained on sequences sampled from a single Markov Chain. Perhaps closest to our work, Nichani et al. [30] introduces a general family of in-context learning tasks with causal structure, a special case of which is in-context Markov chains. The authors prove that a simplified transformer architecture (similar to the one we introduce in Section 2.2) can learn to identify the causal relationships by training via gradient descent, and also characterize the ability of the trained models to adapt to out-of-distribution data. The focus of our work, instead, is on the different stages of training and how they relate to specific, well-defined, strategies. 
setup: 
 In this section, we describe our learning problem and present the neural network architectures that we will use for learning.ICL-MC Task. Our learning task consists of sequences generated from Markov Chains with random transition matrices. The goal is to in-context estimate the transition probabilities from sampled sequences, in order to predict the next state. Formally, each sample sequence is generated by a Markov Chain with state space S = {1, . . . , k} and a transition matrix P sampled from a prior distribution, with x 1 drawn from some other prior distribution (potentially dependent on P), and the rest of x = (x 1 , . . . , x T ) drawn from the Markov Chain. We primarily focus on the case where each row of the matrix is sampled from the Dirichlet distribution with concentration parameter α, i.e. P i,: ∼ Dir(α). We want to learn a predictor that, given context x 1 , . . . , x T , predicts the next token, x T +1 . Note that this is an inherently non-deterministic task, even provided full information about the transition matrix, and as such it can better capture certain properties of language than previous in-context learning modeling approaches, such as linear regression [20] . We focus on the case of the flat Dirichlet distribution, with α = (1, . . . , 1) ⊤ , that corresponds to uniform transition probabilities between states. We draw the initial state x 1 from the stationary distribution π of the chain (which exists almost surely). We primarily consider the case where the number of states k is 2 or 3.In subsection 3.3, we consider the generalization of this setting to n-grams for n > 2. Instead of the distribution of x T being determined by x T -1 , we let it be determined by x T -n+1 , . . . , x T -1 , according to a conditional distribution P which is uniform over the possible states3 .We adopt the Bayesian interpretation of in-context learning [42] , in which a prior distribution is provided by the training data, and, at test time, the model updates this prior given the in-context sequence. In this framework, we focus on two strategies for Bayesian inference: a (suboptimal) unigram strategy which assumes tokens in each sequence are i.i.d. samples (and counts the frequency of the states in the sequence so far), and the bigram strategy which correctly takes into account dependencies among adjacent tokens (and counts frequency of pairs of tokens).1st strategy: Unigrams. Since we initialize the Markov chain at its stationary distribution (which exists a.s.), the optimal strategy across unigrams is just to count frequency of states and form a posterior belief about the stationary distribution. Unfortunately, the stationary distribution of this random Markov chain does not admit a simple analytical characterization when there is a finite number of states, but it can be estimated approximately. At the limit of k → ∞, the stationary distribution converges to the uniform distribution [10] .2nd strategy: Bigrams. For any pair of states i and j, let P ij be the probability of transitioning from i to j. On each sample x, we can focus on the transitions from the i-th state, which follow a categorical distribution with probabilities equal to (P i1 , . . . , P ik ). If we observe the incontext empirical counts {c ij } k j=1 of the transitions, then P ij is given by: (P i1 , . . . , P ik ) |x ∼ Dir(k, c i1 + α 1 , . . . , c ik + α k ), where α 1 , . . . , α k are the Dirichlet concentration parameters of the prior. Hence, each P ij has a (marginal) distribution that is actually a Beta distribution:P ij |x ∼ Beta c ij + α j , j α j + N i -α j -c ij ,where N i is the total number of observed transi-0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 Second Layer t=3 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 t=92 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 t=121 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 First Layer 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2 0 1 2 1 0 2 2 0 0 0 1 2 1 1 2Figure 2 : Attention patterns that correspond to the last token of the sequence for a transformer trained to perform ICL-MC. The intensity of each blue line signifies the strength of the corresponding attention value. As the model gets trained, we observe that the attention weights mimic the construction of Proposition 2.1. Specifically, at the end of training (right), each token in the first layer is attending to the previous token. In the second layer, the last token, a "2", is attending to tokens that followed "2"s, allowing bigram statistics to be calculated. See also Figure 9 for full attention matrices during the course of training.tions from state i. As such, our best (point) estimate for each state j is given by: E[P ij |x] = cij +αj N + i αi . For the uniform Dirichlet, α = (1, . . . , 1) ⊤ , it is E [P ij |x] = cij +1 Ni+k .We are interested in investigating how transformers [40] can succeed in in-context learning this task. We focus on attention-only transformers with 2 layers with causal masking which is a popular architecture for language modeling. Given an input sequence x, the output of an n-layer attention-only transformer 4 is:EQUATIONwhere E ∈ R T ×d is an embedding of x ∈ R d , P ∈ R d×k is a linear projection to the output logits, and Attn(x) is masked self attention with relative position embeddings [38] , which is parameterized byEQUATIONIn general, transformers often contain an MLP module, but for this task they are not necessary (see Appendix A and Figure 10 for additional experiments with transformers with MLPs). During training, we minimize the loss:EQUATIONwhere θ denotes the parameters of the model and l is the cross entropy loss. Notice that we provide supervision in all positions, as standard in language modeling.We now show how a two-layer transformer of the above architecture can represent the optimal bigrams solution.Proposition 2.1 (Transformer Construction). A single-head two layer attention-only transformer can find the bigram statistics in the in-context learning Markov chain task.Intuitively, the first layer of the transformer copies the previous token at each position, and in the second layer each token sums the embeddings of all the tokens whose output from the first layer matches itself. The full proof can be found in Appendix B.1. (for the transformer, these are for the first layer, and averaged over all tokens). The graph on the right shows the KL-divergence between the outputs of the models and three strategy. The lower the KL-divergence, the more similar the model is to that strategy.expressive enough to be able to represent such a solution, but also simple enough to be amenable to analysis. Let e xi denote the one-hot embedding that corresponds to the state at position i ∈ [T ], and let E be the R (T +1)×k one-hot embedding matrix. Then the model is parameterized by W ∈ R k×k and v ∈ R T +1 and defined as:EQUATIONwhere mask (•) is a causal mask, and Softmax(M ) i,j = exp(Mi,j ) T t=0 exp(M T ,j ) . Notice that the role of W is to mimic the attention mechanism of the second layer and the role of v is that of the relative positional embeddings. This model can be seen as a simplified version of a two-layer linear attention-only transformer. See also Appendix B.2 for a discussion. Fact 2.2. Both the bigrams strategy and the unigrams strategy can be expressed by the minimal model with a simple choice of weights.• Bigrams: Let v = (0 c 0 . . . 0) ⊤ and W = I k×k , then f (E) T,s = T t=2 1 {x t = s} 1 {x t-1 = x T } + O kT 2 exp(c) .• Unigrams:For v = (0 . . . 0) ⊤ , W = 1 k 11 ⊤ , we have f (E) T,s = T t=1 1 {x t = s}.See Section B for the proofs.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: The purpose of this work is to introduce a novel task to study in-context learning in transformers. Since it focuses on a theoretical understanding, certain assumptions must be made. We mention in the introduction that our theoretical results apply to a simplified model of a transformer, while we precisely define our experimental setup.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
empirical findings and theoretical validation: 
 In this section, we present our empirical findings on how transformers succeed in in-context learning Markov Chains, we demonstrate the different learning stages during training and the sudden transitions between them, and draw analytical and empirical insights from the minimal model. (right) The values of the positional embeddings (1st layer) that display a curious even/odd pattern. This is before any softmax is applied to the positional embeddings.We focus on attention-only transformers with 2 layers with causal masking and relative positional encodings and train them with the Adam optimizer on ICL-MC (see Section A for experimental details). As can be seen in Figure 3 , all the models converge near the Bayesian optimal solution, suggesting that they learn to implement the bigram strategy. Curiously, however, learning seems to be happening in stages; there is an initial rapid drop and the model quickly finds a better than random solution. Afterwards, there is a long period of only slight improvement before a second rapid drop brings the model close to the Bayes optimal loss. We observe that training a 1-layer transformer fails to undergo a phase transition or converge to the right solution, even if trained for 10x the amount of time -see Figure 7 in the Appendix.Interestingly, as can be seen from the horizontal lines in Figure 3 , the intermediate plateau corresponds to a phase when the model reaches the unigram baseline. We provide evidence that this is not a coincidence, and that after the initial drop in loss, the model's strategy is very similar to the unigram strategy, before eventually being overtaken by the bigram strategy. Some of the strongest such evidence is on the right in Figure 3 , where we plot the KL divergence between model's prediction and the two different strategies. For both the strategies, their KL divergence from the model quickly goes down, with the unigram solution being significantly lower. Around the point of the second loss drop, the KL divergence between the model and the bigram solution decreases, while the other one increases, making it clear that the model transitions from the one solution to the other. This final drop is what has been associated to prior work with induction heads formation [31] ; special dedicated heads inside a transformer are suddenly being formed to facilitate in-context learning. Similar observation hold for Markov Chains with a larger number of states -see Figures 8 and 11 .Mechanistic evidence for solutions found by transformer. To confirm how the two layer attentiononly transformer solves ICL-MC, we inspected the attention in each layer throughout training. Figure 2 shows the attention for a particular input during different parts of training. By the end of training, the attention patterns match that of our construction in Proposition 2.1, with the first layer attending to tokens one in the past, and the second layer attending to tokens that follow the same token as the current one.Varying the data distribution -Unigrams slow down learning. There are several interesting phenomena in the learning scenario that we just described, but it is the second drop (and the preceding plateau) that warrants the most investigation. In particular, one can ask the question: is the unigram solution helpful for the eventual convergence of the model, or is it perhaps just a by-product of the learning procedure? To answer these questions, we define distributions over Markov chains that are in between the distribution where unigrams is Bayes optimal, and the distribution where unigrams is as good as uniform -see Appendix A for more details. As we see in Figure 4 , the transformers that are being trained on the distribution where there is no unigrams "signal" train much faster. And even more tellingly, giving additional "unigram samples" curiously slows down learning. See also Figure 12 in the Appendix that displays how the models perform on different parts of the distribution during training.To abstract away some of the many complicated components from the transformer architecture, we focus our attention now to the minimal model of Section 2.2. We train minimal models of eq. ( 4), starting from a deterministic constant initialization, by minimizing the cross entropy loss with SGD.Full experimental details can be found in Appendix A. Figure 3 (bottom) displays the training curves for the minimal model. Similar to the transformer, learning occurs in two stages and the models eventually converge close to the optimal solution.We now provide theoretical insights on how training progresses stage by stage and how this is achieved by the synergy between the two layers. As it turns out, there need to be at least two steps of gradient descent in order for both elements of the solution to be formed. Proposition 3.1. Let the model be defined as in eq. ( 4) and initialized with W (0) = 0, v (0) = 0. Suppose the transition matrix P ∈ R k×k is sampled from one of the following two types of distribution:1. k = 2, and P is sampled from the uniform distribution over the set of 2 × 2 stochastic matrices. 2. For any constant k and 0 < α < 1, with probability α, sample the matrix P uniformly from a "bigram-only" distribution-the set of k × k doubly stochastic matrices; and with probability 1 -α use a "unigram-only distribution": draw a vector u uniformly from the set {u ∈ R k ≥0 : ∥u∥ 1 = 1} and let P = 1u ⊤ . Then after one step of population gradient descent with step size η > 0,W (1) = Θ(ηT )I + Θ(ηT )11 ⊤ + E and v (1) = 0 where ∥E∥ ≤ O(η log T ).In other words, for large T , the second layer weights are a mixture of the correct solution I and uniform attention 11 ⊤ .Assuming in the first step η 1 = O 1 T 2 , then W (2) has the same structure as W (1) (up to scaling). Furthermore, v1 = Ω(η 2 log T ), and v (2) n ≤ cv(2)1 ∀n ̸ = 1, c < 1 where η 2 is the step size for the second step.If η 2 = Ω(T ), then the output of the model will be a weighted sum of bigrams and unigrams strategy. Formally,f (E) T,s = Θ(η 2 T ) T i=1 1 [x i-1 = x t , x i = s] + Θ(η 2 T ) T i=1 1 [x i = s] + O(log T )Note In the first distribution (uniformly random 2 × 2) or the second distribution with k > 6, at the end of the two steps, the weight on bigrams is greater than that of the weight on unigrams strategy.Proof Overview. The idea of the proof is that a first step of gradient descent with a small learning rate can align the second layer, while a second step can learn to identify the correct relative positional embedding. The identity bias of W in the second layer ensures there is a strong signal in the gradient to look back one in the first layer. Without a bias in W , the gradient for the positional encodings, v, is zero.We get additional intuition from the second distribution (mixture of unigrams and doubly stochastic): in the first step, effectively all of the gradient comes from the examples where the unigram strategy is optimal, while in the second step effectively all of the gradient comes from the examples where the bigram strategy is optimal. Remark 3.2. It is worth noting that, while this is a simplified setting, the analysis goes beyond NTK-based [22] analyses where the representations do not change much and it crucially involves more than one step which has been a standard tool in the analysis of feature learning [6] .We summarize the key theoretical implications: Learning occurs in two phases. Both in the theoretical and experimental models, training has two phases that work at very different speeds. The first phase is fast in both cases; in the theoretical setting, even a single step with step size O 1 T is sufficient for learning the second layer. In the second phase, a much larger step size of Ω( 1) is needed in order to learn the positional encodings (in one step).Second layer is learned first. It has been observed before in a similar bigram learning setting with a two-layer transformer that the model might be learning first the second layer [8] . We also make similar observations in our experiments with the minimal model and the transformers (see Figure 2 ). For the minimal model, the gradient calculations, clearly suggest that starting from a default initialization, it is only the second layer that quickly "picks up" the right solution.Even/odd pattern in positional encodings emerges. We notice in the experiments, that the positional embeddings of both the transformer and minimal model displayed an intriguing even/odd oscillating pattern -see Figure 3 (top, center) and Figure 4 (right). We believe that a careful analysis the gradient of v in the second step will recover this pattern, which is likely related to the moments of the eigenvalues of the transition matrix.Finally, we investigate the performance of transformers on learning in-context n-grams for n > 2; in particular, 3-grams. We train attention-only transformers with three heads in each layer 5 by minimizing the in-context cross entropy loss with the Adam optimizer. As can be seen in Figure 5 (left), the model eventually converges to the Bayes optimal solution. Interestingly, as in the case of Markov Chains, the model displays a "hierarchical learning" behavior characterized by long plateaus and sudden drops. In this setup, the different strategies correspond to unigrams, bigrams and trigrams, respectively. This is presented clearly on the right of Figure 5 , where we plot the similarity of the model with the different strategies and it exhibits the same clear pattern as in the case of n = 2. We leave a more thorough investigation of n-grams for future work.The main limitation of our work is that our analysis relies on a simplified transformer architecture and our learning task is synthetic. Yet, we see the simplicity of our modeling as a positive, since it allows to make rigorous predictions about the mechanisms behind in-context learning abilities of a transformer. On the theoretical front, it would be interesting to extend our analysis to handle higher number of symbols and more complex models for language generation (beyond Markov chains).On the empirical front, it would be worthwhile to understand similar stage-wise learning with natural language data, and use insights from our minimal model to improve formation of induction heads. In particular, it would be great to understand if better data curriculum could remove the undesirable simplicity bias we observe from unigrams. Such simple but incomplete solutions may be commonplace in language modeling and other rich learning settings; for any such solution, one can ask to what extent its presence speeds up or slows down the formation of more complex circuits with higher accuracy.We train transformers of the form (1) with the AdamW optimizer with learning rate 3 × 10 -5 (for 3-grams a learning rate of 3 × 10 -2 was used), batch size 64, and hidden dimension 16. The sequence length of the examples is 100 tokens. The minimal model was trained with SGD, with batch size 64, and learning rate 3 × 10 -4 . We use PyTorch 2.1.2.The data was generated in an online fashion, using numpy.random.dirichlet to generate each row of the transition matrices. At each epoch, we generate a new transition matrix, and, then, a sequence from the induced Markov chain (starting from the stationary distribution of the chain). Both the model initialization (for the transformers) and the data were randomized based on the seed (in a perfectly reproducible manner). Some of the training and model code was based on minGPT [24] . The experiments all measure the outputs of the models at the last token.All of the experiments were performed with a single NVIDIA GeForce GTX 1650 Ti GPU with 4 gigabytes of vram with 32 gigabytes of system memory. Each training run took under ten minutes.The code used can be found at: https://github.com/EzraEdelman/Evolution-of-Statistical-Induction-Heads.Details for experiment on the left in Figure 4 In this experiment, we compared the test loss of the model trained in too different ways. Consider two distributions, each uniform over their support (subsets of the space of 3 × 3 transition matrices for this particular experiment):1. The "doubly stochastic" distribution is uniform over the space of doubly stochastic matrices, that is, transition matrices for which each row or column has entries summing to 1. This is equivalent to the space of Markov Chains with a uniform stationary distribution (that is, π = 1 k 1), which means that the unigram and uniform strategy are the same. 2. The "unigram distribution" is uniform over the space of stochastic matrices P such that for any two rows P i , P j in P , P i = P j . This is equivalent to the distribution over markov chains for which the distribution of the next state doesn't change depending on the previous state. Notice that the unigram strategy is asymptotically optimal on any markov chain from this distribution.We first train models on a number of samples from the doubly stochastic distribution, and plot their training loss as the blue line. Then, we trained models on a random mixture of the two distributions, with 75% of the samples coming from doubly stochastic distribution, and the remaining 25% coming from the unigram distribution. Importantly, we train this second batch of models on more total samples, so that they see the same number of samples from the doubly stochastic distribution, and then plot their test loss in orange with the x-axis being the number of samples from the doubly stochastic distribution seen. This means that every point on the x-axis in the graph, the models in the orange line have seen more samples than seen by those in the blue line, and yet still take longer to converge.We would also like to note that the test loss measured is in the distribution uniform over all stochastic matrices, which the models in the orange line do seem to generalize slightly better to after convergence.In our experiments, we used KL divergence to measure the difference between the probabilities predicted by the model and various other probability distributions. For test loss (depicted, for instance, in Figure 3 showing that a single layer transformer can not achieve good performance on ICL-MC. This result holds for transformers with or without MLPs, and with absolute or relative positional encodings. These graphs show that even trained 8 times longer, there is no notable increase in performance beyond the unigrams strategy (orange line). transition matrices) and is not directly observable by the learner-this is what needs to be learned in-context.For measuring how close the model was to various strategies (for instance, in Figure 3 (right)), we computed the predicted probabilities given by said strategies, and used those as the base distribution. Note that the output of the bigrams strategy (which is Bayes-optimal for our base setting) is different from the aforementioned ground-truth P x T -1 ). Instead, as described in Section 2, it is a Bayesian posterior distribution of the next state given the observed sequence, with the prior determined by the prior distribution of transition matrices. Formally, this corresponds toE[P x T -1 |x 1:T -1 ], where the expectation is taken over the draw of Markov chain transition matrix.In the main text, we mainly show experiments with one seed per experiment, in order to keep presentation simple. Figure 6 justifies this choice: it plots the test loss of two-layer transformers (multiple random seeds) trained on ICL-MC for chains with 3 states. Randomness slightly affects the duration of the plateau, but not the qualitative, two-stage, process of learning.Figure 7 demonstrates the inability of one layer attention only transformers to in-context learn Markov Chains with 3 states.Figure 8 shows that our results extend for Markov Chains with larger state spaces (here, k = 8). As the number of states grows, larger sequence lengths are needed for learning (this is to be expected as a larger transition matrix needs to be estimated in context -roughly, sequence length needs to be Ω(k 2 )). We also know that, in this particular family of Markov Chains, the stationary distribution approaches the uniform distribution as the number of states grows. As a result, we expect the difference between uniform and unigram solutions to be less noticeable.Figure 9 shows the attention patterns for the two layers of a transformer during training.In Figure 10 , we demonstrate that our observations extend to two-layer transformers, which have an additional fully-connected MLP on top of the attention layers.Finally, in Figure 12 we plot the performance of the transformer and the minimal model in various distributions over the course of training.In this section, we present our theoretical results on in-context learning Markov Chains of Section 2.2. Proof of Proposition 2.1. Set the internal dimension d = 3k, and choose e x to be one-hot embeddings-that is, e xi = δ xi , where δ is the Kronecker delta. We will call the parameters of attention layer i, WEQUATIONEQUATIONi,j = (e i W(1)EQUATIONi-j+1 ) ⊤ √ d .(1)EQUATIONi,j ≈ 1[j = i -1] for large enough c. So, for any 2 ≤ i < T, 1 ≤ j < k, Attn 1 (e) i,j+k = e i-1,j . Effectively, the first layer appends the embedding of the previous token after the embedding of the current token, so that the output at position i is approximately (e xi e xi-1 0). The second layer is defined as follows:EQUATIONEQUATIONEQUATIONV =   0 0 I k×k 0 0 0 0 0 0   Note that z = e + Attn 1 (e), then A (2) i,j = (z i W (2) Q )(z j W (2) K ) ⊤ √ d = ce xi (e xj-1 ) ⊤ √ d = c √ d 1[x j-1 = x i ].So, for all j < i, softmax(mask(A)) i,j ≈1[xj-1=xi] i h=1 1[x h-1 =xi] for large enough c. For any 2 ≤ i < T, 1 ≤ j < k, Attn 2 (e) i,j+2k = 3k h=1 1[x h-1 = x i ] i g=1 1[x g-1 = x i ] (zW (2) V ) h,j = k h=1 1[x h-1 = x i ]1[x h = j] i g=1 1[x g-1 = x i ]. This is exactly the empirical bigram statistics (that is, the number of times x i → j appears before position i). In order to make this the output, we setP =   0 0 I k×k   6Setup and notation Our data consists of sequences of length T + 1, x = (x 0 , . . . , x T ), drawn from a Markov Chain with state space S = {1, . . . , k} (i.e., x j ∈ {1, . . . , k} for all j ∈ [T ]), and a transition matrix P . Each row of the matrix is sampled from a flat Dirichlet distribution, i.e. P i ∼ Dir(1), corresponding drawing the row from a uniform distribution over the simplex. Let 6 Technically, the output of this construction is not the log probabilities as generally cross-entropy loss assumes. These can be approximated linearly by setting P = In practice, this approximation can achieve close to Bayes optimal loss. E ∈ {0, 1} (T +1)×k be the one hot embedding matrix of x, that is, E i,xi = 1 and for all s ̸ = x i E i,s = 0. A crucial difference with parallel work which also studies in-context learning of Markov Chains [30] is that the whole sequence is sampled from the Markov Chain (whilst in [30] , the penultimate token is sampled from the uniform distribution).Model We define our model as a simplified sequence to sequence transformer f : R T ×k → R (T +1)×k with f (E) = mask(EW (Softmax(M )E) ⊤ )E. The trainable parameters are W ∈ R k×k and v ∈ R T +1 . We defineM ∈ R (T +1)×(T +1) as M =     v 0 -∞ . . . -∞ v 1 v 0 . . . -∞ . . . . . . • • • . . . v T v T -1 . . . v 0     , that is, for all T ≥ i ≥ j ≥ 0, M i,j = v i-j and if i > j, M j,i = -∞. Furthermore, v = [v 0 , v 2 , . . . , v T ] ∈ R t×T +1. Softmax is defined as follows:Softmax(M ) i,j = exp (M i,j ) T T =1 exp (M i,j ).The logit for symbol s at position T for our model is:f (E) T,s = k u=1 T i=0 i j=0 W x T ,u 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ).(EQUATIONThe model can represent the unigrams and bigrams solutions as following:• Construction for bigrams: v = (0, c, 0, . . . , 0) ⊤ andW = I k×k , then f (E) T,s = T i=0 1 [x i = s ∧ x i-1 = x T ] + O T 3exp(c) . As c tends to infinity, this becomes bigrams.• Construction for unigrams: v = 0 andW = 1 k 1 ⊤ 1, then f (E) T,s = T i=0 1 [x i = s].We prove the above claims.Proof. We will first prove the unigrams construction.f (E) T,s = k u=1 T i=0 i j=0 W x T ,u 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = k u=1 T i=0 i j=0 1[x i-j = u ∧ x i = s] 1 i = 1 k T i=0 i j=0 1[x i = s] k i = T i=0 1[x i = s],which is exactly the unigrams solution.Now consider the bigrams construction. As c grows, the softmax of v very quickly becomes one hot. Formally, by Lemma B.7 in [17] , for any i > 0,exp(v j ) i ℓ=0 exp(v ℓ ) = 1 [j = 1] + O T exp(c) So, f (E) T,s = k u=1 T i=0 i j=0 W x T ,u 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = T i=0 i j=0 1[x i-j = x T ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = 1[x T = x 0 = s] + T i=1 i j=0 1[x T = x i-j ∧ x i = s] exp(v j ) i -1 + exp(c)) = 1[x T = x 0 = s] + T i=1 i j=0 1[x T = x i-j ∧ x i = s] 1 [j = 1] + O 2T exp(c) (Lemma B.7 in [17]) = 1[x T = x 0 = s] + T i=1 1[x T = x i-1 ∧ x i = s] + T i=1 i j=0 1[x T = x i-j ∧ x i = s]O 2T exp(c) = T i=1 1[x T = x i-1 ∧ x i = s] + T i=1 O T 3 exp(c)This simplified model was constructed by taking a two layer transformer with relative positional encodings and simplifying it. Our construction for how transformers would form induction heads (corroborated with experiments such as the viewing of attention patterns in figure 2 ) implies that the MLPs and the value matrices could just be identity functions, and the first layer query matrix, and the second layer positional embeddings were zero matricies, so in the simplified model we froze these parameters to there final states. We also remove the softmax on the attention in the first layer. Despite these changes, the training dynamics, our main interest, stay remarkably similar.Training We analyze gradient descent with the cross entropy loss L T (f, E, x T +1 ) = -k s=1 log Softmax (f (E)) T,s P X T ,s 7For use in the proofs, here we show the calculations of the gradients of the model with respect to the parameters, and the loss with respect to the model.∂f (E) T,s ∂W a,b = k u=1 T i=0 i j=0 1[x T = a ∧ b = u]1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = T i=0 i j=0 1[x T = a]1[x i-j = b ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) ∂f (E) T,s ∂v a = k u=1 T i=0 i j=0 W x T ,u 1[x i-j = u ∧ x i = s] 1[j = a] exp(v a ) i ℓ=0 exp(v ℓ ) -1[a ≤ i] exp(v j ) i ℓ=0 exp(v ℓ ) exp(v a ) i ℓ=0 exp(v ℓ ) = k u=1 T i=0 exp(v a ) i ℓ=0 exp(v ℓ ) i j=0 W x T ,u 1[x i-a = u ∧ x i = s]1[j = a] -1[x i-j = u ∧ x i = s]1[a ≤ i] exp(v j ) i ℓ=0 exp(v ℓ ) = k u=1 T i=0 exp(v a ) i ℓ=0 exp(v ℓ ) W x T ,u   1[x i-a = u ∧ x i = s]1[a ≤ i] - i j=0 1[x i-j = u ∧ x i = s]1[a ≤ i] exp(v j ) i ℓ=0 exp(v ℓ )   = k u=1 T i=a exp(v a ) i ℓ=0 exp(v ℓ ) W x T ,u   1[x i-a = u ∧ x i = s] - i j=0 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ )   = k u=1 T i=a exp(v a ) i ℓ=0 exp(v ℓ ) W x T ,u 1[x i = s]   1[x i-a = u] - i j=0 1[x i-j = u] exp(v j ) i ℓ=0 exp(v ℓ )   ∂L T ∂f (E) T,s = Softmax(f (E)) T,s -P x T ,s B.4 Proof of lemma 3.1Proof. Recall that at initialization, v = 0 and W = 0, implying further that f (E) = 0.First consider the gradient of the loss with respect to W . By chain rule,∂L T (E) ∂W a,b = k s=1 ∂L T ∂f (E) T,s ∂f (E) T,s ∂W a,b = k s=1 (Softmax(f (E)) T,s -P x T ,s ) T i=0 i j=0 1[x T = a]1[x i-j = b ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = k s=1 1 k -P x T ,s 1[x T = a] T i=0 i j=0 1[x i-j = b ∧ x i = s] 1 i + 1 = 1 k 1[x T = a] T i=0   i j=0 1[x i-j = b] 1 i + 1 - k s=1 P a,s 1[x T = a] i j=0 1[x i-j = b ∧ x i = s] 1 i + 1   = 1 k 1[x T = a] T i=0   1[x i = b] - k s=1 P a,s 1[x T = a i j=0 1[x i-j = b ∧ x i = s] 1 i + 1   = 1 k T i=0   1[x i = b ∧ x T = a] - k s=1 P a,s i j=0 1[x i-j = b ∧ x i = s ∧ x T = a] 1 i + 1   = 1 k T i=0   1[x 0 = b ∧ x T -i = a] - k s=1 P a,s i j=0 1[x 0 = b ∧ x j = s ∧ x T -i+j = a] 1 i + 1  Where the last line follows from the Markov property. Now we take the expectation over x, x T +1 conditioned on the transition matrix P ,E x|P ∂L T ∂W a,b = π b T i=0   1 k P T -i b,a - k s=1 P a,s 1 i + 1 P T -i s,a i j=0 P j b,s   = π b T i=0   1 k P i b,a - k s=1 P a,s 1 T -i + 1 P i s,a T -i j=0 P j b,s   = π b π a (T + 1) 1 k - k s=1 P a,s π s + O(log T )Where the last step follows from Lemma B.6. Then, by applying Lemma B.1 (for the uniform over all 2 × 2 transition matrices case) or lemmas B.3 and B.4 (for the mixture of distributions case), there exist positive constants (potentially depending on k, but not T ) A, B such that for all aE x|P ∂L T ∂W a,a = -(A + B)T + O(log T )and for all a ̸ = b,E x|P ∂L T ∂W a,b = -BT + O(log T )The updated W a,b after the gradient step is just-η 1 E x|P ∂L T ∂W a,b (because W is initialized at 0). Choose η 1 = Θ 1T , so that W will be O(1) with respect to T after the first step. For the gradient with respect to v, since W = 0,∂F (E) T,s ∂v = k u=1 T i=a exp(v a ) i ℓ=0 exp(v ℓ ) W x T ,u   1[x i-a = u ∧ x i = s] - i j=0 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ )   = 0 So, ∂L T (E) ∂v = k s=1 ∂L T (E) ∂f (E) T,s ∂F (E) T,s ∂v = 0Completing the first step calculations.After the first step, W = η 1 AI + B1 ⊤ 1 . Now let us bound the output of the model,|f (E) T,s | = k u=1 T i=0 i j=0 W x T ,u 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = k u=1 T i=0 i j=0 η 1 AI + B1 ⊤ 1 x T ,u 1[x i-j = u ∧ x i = s] 1 i ≤ η 1 k u=1 T i=0 i j=0 (A + B) 1[x i-j = u ∧ x i = s] 1 i ≤ η 1 k u=1 T i=0 i j=0 (A + B) 1[x i-j = u] 1 i ≤ η 1 T i=0 i j=0 (A + B) 1 i ≤ η 1 T |A + B|So, using the first order approximation of softmax,∂L T (E) ∂f (E) T,s = Softmax(f (E)) T,s -1 [x T +1 = s] = 1 k + f (E) T,s k - k u=1 f (E) T,u k 2 + O(f (E) 2 T,s ) -1 [x T +1 = s] = 1 k + O η 1 T k (A + B) + O(η 2 1 T 2 (A + B) 2 ) -1 [x T +1 = s] = 1 k + O η 1 T k (A + B) + O(η 2 1 T 2 (A + B) 2 ) -1 [x T +1 = s] = 1 k -1 [x T +1 = s] + O 1 TWhere the last step follows since for the first step, eta = O 1 T 2 . Now, we can begin to analyze the gradients with respect to the parameters. For W , the gradient is approximately the same as in the last step. Notice that ∂f (E) T ,s ∂W a,b does not depend on W , and v is unchanged, so∂f (E) T ,s ∂W a,b is unchanged. Furthermore, ∂f (E) T,s ∂W a,b = k s=1 T i=0 i j=0 1[x T = a]1[x i-j = b ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = T i=0 i j=0 1[x T = a]1[x i-j = b] 1 i ≤ T i=0 i j=0 1 i = TWe will now show that the gradient is approximately the same as in the first gradient step:∂L T (E) ∂W a,b = k s=1 ∂L T ∂f (E) T,s ∂f (E) T,s ∂W a,b = k s=1 1 k -1 [x T +1 = s] + O 1 T ∂f (E) T,s ∂W a,b = k s=1 1 k -1 [x T +1 = s] ∂f (E) T,s ∂W a,b + O 1 T ∂f (E) T,s ∂W a,b = π b π a (T + 1) 1 k - k s=1 P a,s π s + O(log T )Where the last lines follows from the gradient calculations in the first step. Now we will consider the gradient with respect to v. First, notice that the uniform component of W , B1 ⊤ 1, has no affect on the gradient of v:∂f (E) T,s ∂v a = k u=1 T i=a W x T ,u exp(v a ) i ℓ=0 exp(v ℓ ) 1[x i = s]   1[x i-a = u] - i j=0 exp(v j ) i ℓ=0 exp(v ℓ ) 1[x i-j = u]   = k u=1 T i=a mI + B1 ⊤ 1 x T ,u 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   = k u=1 T i=a (A1[x T = u] + B) 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   = A k u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   + B k u=1 T i=a 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   = A k u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   + B T i=a 1 i + 1 1[x i = s]   k u=1 1[x i-a = u] - i j=0 1 i + 1 k u=1 1[x i-j = u]   = A k u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   + B T i=a 1 i + 1 1[x i = s]   1 - i j=0 1 i + 1   = A k u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   By chain rule, ∂L T ∂v a = k s=1 ∂L T ∂f (E) T,s ∂f (E) T,s ∂v a = k s=1 1 k -P x T ,s + O 1 T k u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   = k s=1 1 k -P x T ,s k u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   + O log T TWhere the last step follows becausek u=1 T i=a 1[x T = u] 1 i + 1 1[x i = s]   1[x i-a = u] - i j=0 1 i + 1 1[x i-j = u]   ≤ k u=1 T i=a 1[x T = u] 1 i + 1 = T i=a 1 i + 1 ≤ log TIn expectation over the values of x, conditioned on the choice of P :E x|P ∂L T ∂v a = k s=1 k u=1 1 k -P u,s T i=a π u i + 1 P T -i s,u   (P a ) u,s - 1 i + 1 i j=0 P i-j u,s   + O log T T = k s=1 k u=1 1 k -P u,s T i=a π u T -i + 1 P i s,u   (P a ) u,s - 1 T -i + 1 T -i j=0 P j u,s   + O log T T = (log (T + 1) -log (a + 1)) k s=1 k u=1 π 2 u P u,s π s -(P a ) u,s + O(1)Where the last step follows from Lemma B.7. Then, by applying Lemma B.2 or lemmas B.3 and B.4 (depending on the distribution assumption on P ),E x|P ∂L T ∂v 1 < E x|P ∂L T ∂v a E x|P ∂L T ∂v 1 < 0Therefore, after the step is taken,v 1 = Θ(η 2 log T ) v 1 -v n = η 2 Ω(log T )Finally, we can consider the state of the model after the second step. Assume that the step size for v in the second step is O(T ), and the step size for W is1 T (A+B) f (E) T,s = k u=1 T i=0 i j=0 W x T ,u 1[x i-j = u ∧ x i = s] exp(v j ) i ℓ=0 exp(v ℓ ) = 1 A + B k u=1 T i=0 i j=0 AI + B1 ⊤ 1 + O log T T x T ,u 1[x i-j = u ∧ x i = s] 1 [j = 1] + O 2T exp(log(T )) = 1 A + B k u=1 T i=0 i j=0 AI + B1 ⊤ 1 + O log T T x T ,u 1[x i-j = u ∧ x i = s] 1 [j = 1] + O 1 T = 1 A + B k u=1 T i=0 AI + B1 ⊤ 1 x T ,u 1[x i-1 = u ∧ x i = s] + O(log T ) = A A + B T i=0 1[x i-1 = x T ∧ x i = s] + B A + B k u=1 T i=0 1[x i-1 = u ∧ x i = s] + O(log T ) = A A + B T i=0 1[x i-1 = x T ∧ x i = s] + B A + B T i=0 1[x i = s] + O(log T )Since A and B are constant in terms of T , we can recover the desired statements, completing the proof.Lemma B.1. If P is a uniformly random stochastic 2 × 2 matrix, and π is the stationary distribution of P , thenE π 2 a 1 k - k s=1 P a,s π s = 5 12 - 2 3 log(2) ≈ -0.045 and for any b ̸ = a E π a π b 1 k - k s=1 P a,s π s = - 7 6 + 5 3 log(2) ≈ -0.011Proof. We have:EQUATIONFor the non-diagonal elements, it holds:EQUATIONLemma B.2. If P is a uniformly random stochastic 2 × 2 matrix, and π is the stationary distribution of P , then,E k s=1 k u=1 π 2 u P u,s (π s -P u,s ) = -7/2 + 5 log(2) ≈ -0.034and for any n ̸ = 1E k s=1 k u=1 π 2 u P u,s (π s -P u,s ) ≤ E k s=1 k u=1 π 2 u P u,s π s -(P n ) u,sProof. We have:E k s=1 k u=1 π 2 u P u,s (π s -P u,s ) = 1/6x 4 + (x + y)(6xy(-4x 2 + 2x + 1) + 6y 4 + y 3 (20 -24x)12(x + y) y 2 (12x 2 -12x -3) + log((x + y) 6x 2 (4x 2 +2x-1) (x + y) 6y 2 (4y 2 +2y-1) ))12(x + y)EQUATIONFor the inequality, we have an intuition that doesn't depend on k, notice that:k s=1 k u=1 π 2 u P u,s π s -(P n ) u,s ≥ - k s=1 k u=1 π 2 u P u,s π s -(P a ) u,s ≥ - k s=1 k u=1 π 2 u P u,s α n = - k u=1 π 2 u α n ≥ -α nAs long as α isn't concentrated around 1, then this shows that the magnitude of the RHS is bounded by a term that shrinks exponentially in n. For k = 2, we will find a similar bound, and then show separately that for all n for which the bound fails, the inequality still holds true.k s=1 k u=1 π 2 u P u,s π s -(P n ) u,s = P 1,2 P 2,1 (4P 1,2 P 2,1 -P 1,2 -P 2,1 ) (P 1,2 + P 2,1 ) 3 (1 -P 1,2 -P 2,1 ) nWe can show that for any choice of P 1,2 and P 2,1 on the unit square,P 1,2 P 2,1 (4P 1,2 P 2,1 -P 1,2 -P 2,1 ) (P 1,2 + P 2,1 ) 3 ≤ 1 4To see why this is true, observe that,(4P 1,2 P 2,1 -P 1,2 -P 2,1 ) 2 = 16P 2 1,2 P 2 2,1 + (P 1,2 + P 2,1 ) 2 -8(P 1,2 + P 2,1 )P 1,2 P 2,1 ≤ 16P 2 1,2 P 2 2,1 + (P 1,2 + P 2,1 ) 2 -4(P 1,2 + P 2,1 ) 2 P 1,2 P 2,1 since P 1,2 + P 2,1 ≤ 2 = 16P 2 1,2 P 2 2,1 + (P 1,2 + P 2,1 ) 2 -4P 1,2 P 2,1 ((P 1,2 + P 2,1 ) 2 -4P 1,2 P 2,1 ) = (P 1,2 + P 2,1 ) 2 -4P 1,2 P 2,1 (P 1,2 -P 2,1 ) 2 ≤ (P 1,2 + P 2,1 ) 2Using the above, we haveP 1,2 P 2,1 (4P 1,2 P 2,1 -P 1,2 -P 2,1 ) (P 1,2 + P 2,1 ) 2 ≤ P 2 1,2 P 2 2,1 (P 1,2 + P 2,1 ) 2 (P 1,2 + P 2,1 ) 6 = P 2 1,2 P 2 2,1 (P 1,2 + P 2,1 ) 4 ≤ P 2 1,2 P 2 2,1 16P 2 1,2 P 2 2,1 using (P 1,2 + P 2,1 ) 2 ≥ 4P 1,2 P 2,1 = 1 16 . So, ∥ k s=1 k u=1 π 2 u P u,s π s -(P n ) u,s ∥ ≤ 1 4 |1 -P 1,2 -P 2,1 | n Now, E -1 1 4 |1 -P 1,2 -P 2,1 | n = - 1 4 1 0 1 0 |1 -x -y| n = - 1 4 2 (n + 1)(n + 2) = - 1 2(n + 1)(n + 2)Notice that this decreases in n, and at n = 3, 1 2(3+1)(3+2) = 1 40 = 0.025 which is less in magnitude than the value we proved at n = 1, | -7/2 + 5 log 2| ≈ 0.034. So, solving for n = 2 (verified by a symbolic algebra program)E P 1,2 P 2,1 (-P 1,2 -P 2,1 + 1) 2 • (2P 1,2 P 2,1 + P 1,2 (P 2,1 -1) + P 2,1 (P 1,2 -1)) (P 1,2 + P 2,1 ) 3 = - 413 60 + 149 log (2) 15 ≈ 0.002Which is not only greater than -7/2 + 5 log 2, but positive. Lastly, we simply need to show that the inequality holds at n = 0, and we are done.E P 1,2 P 2,1 (-P 1,2 -P 2,1 + 1) 0 • (2P 1,2 P 2,1 + P 1,2 (P 2,1 -1) + P 2,1 (P 1,2 -1)) (P 1,2 + P 2,1 ) 3 = -E P 1,2 P 2,1 • (2P 1,2 P 2,1 + P 1,2 (P 2,1 -1) + P 2,1 (P 1,2 -1)) (P 1,2 + P 2,1 ) 3 = -7/6 + 5 * log(2)/3 ≈ -0.0114Which is greater than -7/2 + 5 log 2, completing our proof.Lemma B.3. If P is a uniformly random doubly stochastic matrix, then,E π 2 a 1 k - k s=1 P a,s π s = E π a π b 1 k - k s=1 P a,s π sfor all a, b andE k s=1 k u=1 π 2 u P u,s (π s -P u,s ) < E k s=1 k u=1 π 2 u P u,s π s -(P a ) u,sFor all non-negative a ̸ = 1. andE k s=1 k u=1 π 2 u P u,s (π s -P u,s ) < 0Proof. We will use the fact that for doubly stochastic matrices, the stationary distribution is the uniform vector 1 k 1. The first equality follows directly from π a = 1 k = π b . Now we will prove the first inequality.E k s=1 k u=1 π 2 u P u,s π s -(P a ) u,s = E k s=1 k u=1 1 k 2 P u,s 1 k -(P a ) u,s = 1 k 2 - 1 k 2 k s=1 k u=1 E P u,s (P a ) u,s = 1 k 2 - 1 k 2 E [⟨P, P a ⟩ F ]Where ⟨., .⟩ F is the Frobenius inner product. We will first consider the case where a > 1. Notice that it is sufficient to prove that for any doubly stochastic P (excluding the measure zero case of where P = P 2 , where equality is reached), ⟨P, P a ⟩ F < ∥P ∥ 2 F . First, by Cauchy-Schwarz, ⟨P, P a ⟩ F < ∥P ∥ F ∥P a ∥ F We can use strictly less than because Cauchy Schwarz is only tight when P and P a are linearly dependent, and since both P and P a are doubly stochastic, linear dependence implies equality, which is only the case when P = P a . Then, For now assume a > 0, then,= ∥P ∥ F ∥P a ∥ F = ∥P ∥ F ∥P P a-1 ∥ F = ∥P ∥ F ∥ i α i Λ i P a-1 ∥ F ≤ i α i ∥P ∥ F ∥Λ i P a-1 ∥ F = i α i ∥P ∥ F ∥P a-1 ∥ F = ∥P ∥ F ∥P a-1 ∥ FThe third step used the well known Birkhoff-Von Neumann Theorem [9] that any doubly stochastic matrix P is the convex combination of permutation matrices, so P = i α i Λ i for some permutation matrices Λ i and constants α i > 0 with i α i = 1. The inequality step uses the triangle inequality. Induction on positive a yields the desired inequality for positive a.Now consider the remaining case, a = 0,1 k 2 - 1 k 2 k s=1 k u=1 E P u,s P 0 u,s = 1 k 2 - 1 k 2 k s=1 E [P s,s ] = 1 k 2 - 1 k 2 = 0.While at a = 1, for any P that isn't 11 ⊤ , ∥P ∥ F > 1, so1 k 2 - 1 k 2 k s=1 k u=1 E P 2 u,s < 1 k 2 - 1 k 2 = 0.Completing the proof of both inequalities.Lemma B.4. If P is a uniformly random k × k stochastic matrix subject to each row being the same, then,E π 2 a 1 k - k s=1 P a,s π s < E π a π b 1 k - k s=1 P a,s π s < 0 and E π 2 a 1 k - k s=1 P a,s π s E π a π b 1 k - k s=1 P a,s π s ≥ 8 5for all a and b andE k s=1 k u=1 π 2 u P u,s π s -(P a ) u,s = 0For all a.Proof. The equality statement follows from the facts that for such transition matrices, P a = P for all natural a > 0, and that the stationary distribution matches the rows, that is, for any a, b, π b = P a,b ,E k s=1 k u=1 π 2 u P u,s π s -(P a ) u,s = E k s=1 k u=1 π 2 u P u,s (π s -P u,s ) = E k s=1 k u=1 π 2 u P u,s (π s -π s ) = 0Now we will do the inequalities. We will also use the following facts derived from the moments of the Dirichlet distribution,E ∥π∥ 2 2 = 2 k + 1 E ∥π∥ 4 2 = 4(k + 5) (k + 1)(k + 2)(k + 3) So, E π 2 a 1 k - k s=1 P a,s π s = E π 2 a 1 k - k s=1 π 2 s = 1 k E ∥π∥ 2 2 1 k -∥π∥ 2 2 = 1 k 2 E ∥π∥ 2 2 - 1 k E ∥π∥ 4 2 = 2 k 2 (k + 1) - 4(k + 5) k(k + 1)(k + 2)(k + 3)Which is negative for all k ≥ 2. And,E π a π b 1 k - k s=1 P a,s π s =E π a π b 1 k - k s=1 π 2 s = 1 k 2 E 1 k -∥π∥ 2 2 = 1 k 3 - 1 k 2 E ∥π∥ 2 2 = 1 k 3 - 2 k 2 (k + 1)Which is also negative for all k ≥ 2. Finally, notice that2 k 2 (k+1) - 4(k+5) k(k+1)(k+2)(k+3) 1 k 3 - 2 k 2 (k+1) ≥ 8 5For all k ≥ 2.The following lemma is a well known property of stochastic matrices, (see Lemma 3.3.2 Gallager [19] for example). Lemma B.5. Let α = 1 -2 min i,j P i,j . Then, for any i, j(P n ) i,j -π j ≤ α nLemma B.6 and Lemma B.7 both share similar intuitions and proofs. They largely rely on Lemma B.5, which shows that (P n ) i,j approaches π j exponentially fast with respect to n, to show that over the course of summations over n the stationary distribution dominates, allowing us to simplify the expressions. Lemma B.6. Let P be a stochastic matrix with all positive entries, and let a, b be states. Assume that min i,j P i,j is positive and doesn't dependend on T . Then,π b T i=0   1 k P i b,a - k s=1 P a,s 1 T -i + 1 P i s,a T -i j=0 P j b,s   = π b π a (T + 1) 1 k - k s=1 P a,s π s + O(log T ).Proof. Let us bound the magnitude of the difference between the two expressions.π b T i=0   1 k P i b,a - k s=1 P a,s 1 T -i + 1 P i s,a T -i j=0 P j b,s   -π b π a (T + 1) 1 k - k s=1 P a,s π s = π b T i=0   1 k P i b,a -π a - k s=1 P a,s   1 T -i + 1 P i s,a T -i j=0 P j b,s -π s π a     ≤ π b T i=0   1 k P i b,a -π a + k s=1 P a,s 1 T -i + 1 T -i j=0 P i s,a P j b,s -π s π a   ≤ π b T i=0   1 k α i + k s=1 P a,s 1 T -i + 1 T -i j=0 P i s,a P j b,s -π s π a   = π b T i=0   1 k α i + k s=1 P a,s 1 T -i + 1 T -i j=0 ( P j b,s P i s,a -π a ) + π a ( P j b,s -π s )   ≤ π b T i=0   1 k α i + k s=1 P a,s 1 T -i + 1 T -i j=0 P j b,s P i s,a -π a + π a P j b,s -π s   ≤ π b T i=0   1 k α i + k s=1 P a,s 1 T -i + 1 T -i j=0 P j b,s α i + π a α j   ≤ π b T i=0   1 k α i + k s=1 P a,s 1 T -i + 1 T -i j=0 α i + α j   ≤ π b T i=0 1 k α i + k s=1 P a,s α i + 1 T -i + 1 1 -α T -i+1 1 -α ≤ π b T i=0 1 k α i + k s=1 P a,s α i + 1 T -i + 1 1 1 -α ≤ π b T i=0 1 k α i + α i + 1 T -i + 1 1 1 -α ≤ π b T i=0 1 k α i + α i + 1 T -i + 1 1 1 -α ≤ π b 1 + 1 k 1 -α T +1 1 -α + log(T + 1) + 1 1 -α ≤ π b 2 + 1 k + log(T + 1) 1 -α ≤ 2 log T 1 -α = log T min i,j P i,j = O(log T )The last step follows from our assumption, completing the proof.Lemma B.7. Let P be a stochastic matrix with all positive entries, and let a, b be states. Assume that min i,j P i,j is positive and doesn't depend on T . Then, k s=1 k u=1 1 k -P u,s i=a π u T -i + 1 P i s,u   (P a ) u,s - 1 T -i + 1 T -iπ u T -i + 1 P i s,u   (P a ) u,s - 1 T -i + 1 T -i j=0 P j u,s   - k s=1 k u=1 1 k -P u,s T i=a π 2 u 1 T -i + 1 (P a ) u,s -π s = k s=1 k u=1 1 k -P u,s T i=a π u T -i + 1   P i s,u   (P a ) u,s - 1 T -i + 1 T -i j=0 P j u,s   -π u (P a ) u,s -π s   ≤ k s=1 k u=1 T i=a π u T -i + 1 P i s,u   (P a ) u,s - 1 T -i + 1 T -i j=0 P j u,s   -π u (P a ) u,s -π s ≤ k s=1 k u=1 T i=a π u T -i + 1 P i s,u   π s - 1 T -i + 1 T -i j=0 P j u,s   -(P a ) u,s -π s P i s,u -π u ≤ k s=1 k u=1 T i=a π u T -i + 1 P i s,u 1 T -i + 1 T -i j=0 π s -P j u,s -(P a ) u,s -π s P i s,u -π u ≤ k s=1 k u=1 T i=a π u T -i + 1   P i s,u 1 T -i + 1 T -i j=0 π s -P j u,s + (P a ) u,s -π s P i s,u -π u   ≤ k s=1 k u=1 T i=a π u T -i + 1   P i s,u 1 T -i + 1 T -i j=0 α j + (P a ) u,s -π s α i   By lemma B.5 = k s=1 k u=1 T i=a π u T -i + 1 P i s,u 1 T -i + 1 1 -α T -i+1 1 -α + (P a ) u,s -π s α i ≤ k s=1 k u=1 T i=a π u T -i + 1 1 T -i + 1 1 1 -α + (P a ) u,s -π s α i ≤ T i=a 1 T -i + 1 k s=1 1 T -i + 1 1 1 -α + k s=1 k u=1 π u (P a ) u,s -π s α i ≤ T i=a 1 T -i + 1 k s=1 1 T -i + 1 1 1 -α + k s=1 (π s -π s ) α i ≤ T i=a 1 T -i + 1 k T -i + 1 1 1 -α ≤ T i=a 1 (T -i + 1) 2 k 1 -α ≤ 2k 1 -α = k min i,j P i,j = O(1)The last step follows from our assumption, and the fact that k does not depend on T . Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
open access to data and code: 
 Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: Our code is shared at https://github.com/EzraEdelman/ Evolution-of-Statistical-Induction-Heads. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: The details are available in Appendix A.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: Our main results regard the multiphase nature of training, and we include Figure 6 which shows that for the seeds 0, 1, . . . 9 the model each time has the same patterns in the training curve. The curves were shown individually instead of through error bars since the main purpose of the loss curves is to show the shape, but because the phase transition doesn't occur at a consistent time, adding error bars smooths out the curve making the phase transition look less sharp.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: The resources are described in Appendix A. Only a single computer was used, and none of the training runs took longer than ten minutes.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: We read the code of ethics in its entirety and strongly believe that our research abides by the stated code.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work focuses on understanding the internal mechanisms of Transformer models on synthetic tasks. We do not foresee any direct societal impact of this work. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Our work does not release any data or models that poses safety risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We cite the Github repository we use as the codebase for our Transformer models. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.In particular, for each tuple of n -1 tokens, we sample the vector of conditional probabilities for the next state from a flat Dirichlet distribution.For simplicity, we assume embedding and hidden dimension are equal, but they can be different in general.Conclusion and DiscussionIn this work, we have introduced a simple learning problem which serves as a controlled setting for understanding in-context learning and the emergence of (statistical) induction heads. Through a combination of theoretical analysis and empirical investigation, we identify different stages during learning, which we were able to precisely characterize. These validate similar observations from training large-scale language models.Our experiments with one head per layer did not converge to low loss during training, but follow-up work by[33] showed that with sufficiently long training time a single head in each layer can suffice.In practice, one could use the empirical value of xT +1 rather than its distribution PX T ,s, but in full batch gradient descent this is in fact equivalent. This is because conditional on xT and P , xT +1 is independent of x1, . . . xT -1. 
 
 Appendix:  


Decision: This paper studies the in-context learning capabilities of Transformers on an in-context learning task, where the model has to a learn a random n-gram model in-context. The paper studies two settings: a simplified Transformer architecture with only attention layers, and a "minimal" model that is more amenable to theoretical analysis. 

 On the positive side, this paper studies an interesting synthetic setup (as studied by the concurrent works mentioned by the authors) that illuminates the inner workings of Transformers. On the negative side:
- the theoretical analysis is only performed on the minimal model, and while the authors argue that the behavioral similarity between the simplified Transformer and the minimal model justifies extending the analysis to the simplified model, this is a bit of a leap
- the synthetic experiments are extremely toy
- there are no real-world experiments (unlike some recent papers which use insights from synthetic tasks to improve actual language models).

Nonetheless, I think the contributions here are enough to justify acceptance.

Reviews: review: paper_summary: paper_summary: In the paper, the authors investigate the phenomenon of in-context learning exhibited by Transformers with the help of a simplified architecture and Markovian synthetic data. The authors show that experimentally the attention layers form statistical induction heads that help the model to implement an add-constant estimator based on the empirical counts of the input. They also show that a simplified transformer architecture can effectively represent such an estimator, and they provide a SGD convergence analysis for a minimal model.
strength_weakness: strength_weakness: The main limitations of the paper are: (1) the transformer architecture is heavily simplified (even if I believe that the results obtained in the paper should extend to more complex architectures); (2) the model used for the SGD analysis is, indeed, minimal. While limitation (1) is not so important to me, I am a bit concerned about limitation (2). I am not sure what is the actual utility of a SGD convergence analysis on such a minimal model. While I think that the representation result of Proposition 2.2 can be extended to a more complex model, I don't think that the SGD analysis of the minimal model would work for a more complex architecture, which would then need completely different and more sophisticated techniques.
questions: questions: 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)?
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?
limitations: limitations: N/A
paper_summary: paper_summary: This paper introduces a task to investigate how in-context learning capabilities are learnt by transformer models. They show that models trained on this task go through a phase transition from which they start by modeling unigram to then acting as a bigram model. The authors further extend their work to the case of n=3 and show similar behavior.
strength_weakness: strength_weakness: 1. Questionable impact since this topic has been explored heavily in the past years, with other similar tasks and toy models existing, showing similar results.
2. Questionable how well the toy setting can actually transfer to real-world ICL, albeit interesting.
questions: questions: What is the key differentiating contribution of the work in comparison with the myriad of other works in this space?

What is the conclusion of the work that you believe is transferable to real-world ICL setting. What n-gram statistics would that follow?
limitations: limitations: It would potentially be useful to also discuss the limitations of the toy task to explain the full-fledged LLM ICL setting. A lot has been said in the paper about conclusions made, and I understand that this the same with any work employing a toy task, but an attempt for making more direct comparisons between real vs. toy setting may enhance the value of the work.
paper_summary: paper_summary: The paper studies in-context learning with transformer models in a simple Markov Chain sequence modelling task. The authors empirically show the formation of statistical induction heads which correctly compute the posterior probabilities given bigram statistics. Moreover, they observe that during training the model undergoes phase transitions where the complexity of the n-gram model increases. They also propose a simplified theoretical model of a two-layer transformer to analyse these phenomena.
strength_weakness: strength_weakness: The major weakness of this work is that the abstract and introduction suggest a primary focus on analysing and uncovering the mechanisms behind the simplicity bias and the phase transition from simple unigram solutions to more complex ones. However, upon reading the main body of the paper, it seems that the experiments mainly reveal the existence of such biases and transitions, while the theoretical section addresses a different phenomenon: how the two layers are learned in different training phases. The plateau behaviour is particularly intriguing, such as the model learning unigrams, bigrams, trigrams, and so on, but their theory does not seem to describe or explain this observation.

1. **Unigram strategy:** Can the authors clarify this statement "*the stationary distribution of this random Markov chain does not admit a simple analytical characterization when there is a finite number of states*". In this context, what does "*this Markov chain*" refer to? 

1. **Simplified transformer:** the proposed simplified transformer seems to be able to capture all the phenomenology of a real transformer. Nevertheless, it is not clear to me if the proposed model is a good proxy for an attention-only transformer or not. In particular: 
	1. The first attention layer does not use the interaction between tokens but represents the attention through a learnable matrix that could in principle learn the same structure. Could the authors elaborate on this choice?
	1. The second attention instead captures the interaction between tokens but between the input and the output of the first layer which is unusual. I understand that this is based on the construction of the real transformer where the output of the first attention is copied in the second block of the embedding and therefore can interact with the input but it still means that the simplified version doesn't need to learn this mechanism. Could the author elaborate on this choice?  
	1. The model seems to be composed of a non-linear attention for the first layer (softmax is present) and linear attention for the second. Could the authors elaborate on this particular choice? Is it a way to simplify the analysis while maintaining the properties of the softmax where needed? 
1. **Data generation:** in section 2, it is explained how the transition matrices are generated according to a Dirichlet distribution, nevertheless in the theory and some of the experiments there is the additional requirements for the matrix to be doubly stochastic which is not mentioned when the setup is described. Could the authors clarify this point and highlight the importance of doubly stochastic transition matrices? 
1. **Proof of Proposition 2.2:** the proof of this proposition appears disorganised and not easy to follow. In particular: 
	1. Setting the internal dimension d=3k seems fundamental to ensure that the model has the correct number of dimensions to copy the tokens from one layer to the next. I do not necessarily have a problem with this choice but it would be useful to see it discussed in the main text. 
	2. Is the definition of $v^{(1)}$ correct? if I am not mistaken it appears that the matrix is of dimensions $3k \times 3k$ whereas by the main text, it should be of dimensions $t \times 3k$ could the author clarify this together with the definition of $\delta_2$ and $1_k$.  
	3. There are multiple quantities used in the text such as $e_{x_i}$ or $e_i$ or $e_{{i-1},j}$ could the authors clarify their meaning. 
	4. In the expression $\text{softmax}(\text{mask}(A))_{i,j} \approx \mathbb{1}[j=i-1]$ shouldn't it be  $[i=j-1]$ ? Why is the indicator function comparing the indexes instead of the values $x_i$ ? Moreover, could the authors clarify why it is only the indicator function and not $\frac{1}{\text{count}([x_i=x_j-1])}$ given the softmax?   
	5. In the expression of $\text{Attn}_2(e)$ why is the summation from $h=1$ to $h=3k$ if h is the index of the element in the sequence? Furthermore a new index $g$ appears which is not used anywhere.  
1. **Experiments in Figure 8:** are the experiments in Figure 8 using d=3k ? do the authors observe that the parameters converge to the construction given in Proposition 2.2 ? 
1. **Proof of Proposition 2.3:** The proof for the unigram construction seems to be missing. 
1. **Unigram and bigram constructions:** For both the two-layer transformer and the simplified model it is possible to show constructions for the unigram and bigram models. Nevertheless, besides giving the weights for such constructions the authors do not discuss the relationship between them and the objective functions. Are they stationary points of the dynamics? Does this help explain the plateau observed in the experiments? is the unigram a saddle point? 
1. **Minimal model:** In line 241 the authors state that "*the minimal model converges to the bigram solution spending however significantly less time at the unigram solution*".  By looking at Figure 3 I understand how the minimal model reaches the same Kl value in almost half the time but it also has fewer parameters and some of the structure is already enforced by construction (for example the fact that the second attention is between input and output of the first layer). Could the authors explain why the comparison still makes sense? 
1. **Varying the data distribution:** I find this part unclear. In line 231 the authors state that they define distributions  "*we define distributions over Markov chains that are in between the distribution where unigrams is Bayes optimal, and the distribution where unigrams is as good as uniform.*" can the authors provide a mathematical definition of such distribution?   Are doubly stochastic transition matrices used to create Markov chains with uniform stationary distribution such that the latter would be the only unigram possible? 
1. **Two phases of learning in Lemma 3.1:** The model is capable of reproducing the effect of learning the second layer first, nevertheless I am not convinced that this is a property of the model of the simplified transformer but rather a consequence of the initialization and step size. Could the authors clarify this point?
questions: questions: See above.
limitations: limitations: - Simplified transformer architecture 
- Synthetic task only, not clear if similar phenomena appear in larger models trained on natural language data.
paper_summary: paper_summary: This work introduces a simple “in-context learning” Markov chain modelling task and studies how transformer models tackle it. The task consists of inferring the transition probability matrix of the Markov chain, which is sampled from a prior Dirichlet distribution. The authors consider the task to pertain to “in-context learning” because each data point (that is, each sequence in the dataset) is sampled from a different transition probability matrix. The prior over the matrices is however global and thus shared by both training and test datasets. More in detail, the authors focus on Markov chains with only two states, and sample each sequence from the stationary distribution of the process.

Given this dataset, the authors investigate the training dynamics of a two-layer attention-only transformer model both empirically and theoretically.  The theoretical investigation involves the introduction of a simplified model, which is defined to mimic the main features of the transformer model.  The authors report three different training stages, during which the model very quickly changes from predicting random transitions, to predicting samples similar to the stationary distribution of the chain, and lastly, predicting transition from the inferred transition probability matrix.

The authors further connect their findings with previous work.
strength_weakness: strength_weakness: My first issue with the paper is: 

- how much does the proposed learning task really concern “in-context learning”? especially in the context of large language models, which the authors use as motivation.  

As framed by e.g. Xie et al. (2022) who focus on LLMs, "in-context learning" deals with sequences which have low probability under the training distribution. The authors do not really elaborate in their understanding or interpretation of "in-context learning", nor do they explain how it relates to "in-context learning" in the setting of language models. In short, I believe the paper contributes more to our understanding of transformers than to our understanding of in-context learning. 

A second major issue is that the authors do not explain their reasoning behind nor the limitations of their dataset/task definition. 

- First, the authors consider Markov chains with only two states and do not comment on why they restricted their study to that case nor how their findings extend to Markov chains with more states, if at all. 
- Second, the authors choose the initial distribution of the chain to be its stationary distribution but, again, did not explain why. Note that a stationary Markov chain corresponds to a very special case, and it's not very clear how the observations made by the authors extend beyond it.  Indeed, choosing a Markov chain with more states, initialised from a distribution far from stationarity, can generate long sequences which are “out-of-equilibrium”. One can't but wonder whether one would still find multiple training stages in this case. See e.g. question 8 below. 

Despite its merits, I think the issues above require some revision, or more detailed explanation, before the paper can be published. 

*Other comments*: 

The comment in line 205 “We observe that training a 1-layer transformer fails to undergo a phase transition or converge to the right solution”, could be better highlighted as evidence for the emergence of induction heads. It feels somewhat buried in between other comments and statements. Also there’s a typo in line 239 and another in line 280.

*References:* 
- An explanation of in-context learning as implicit Bayesian inference. Xie et al. (2022)
questions: questions: 1. Why did you consider Markov chains with two states only?
2. Can you please elaborate on your comment in line 145, page 4: “the stationary distribution of this random Markov chain does not admit a simple analytical characterization when there is a finite number of state”?
3. Why did you choose the stationary distribution as the initial distribution?
4. Is the ground-truth unigram strategy then computed by estimating the stationary distribution with the count histogram? Is this the one used to compute the KL distance in Figure 3?
5. Similarly, do you use the bigram strategy of section 2.1 to compute the corresponding KL distance in Figure 3? Or do you use the instance-dependent ground-truth transition probability matrix instead?
6. Is the difference in KL between the unigram and bigram strategy due to finite sampling?
7. Why does the bigram strategy have a smaller loss? is it because of the inductive bias of the model?
8. In the setting with an initial condition far from the stationary distribution, would the model first fit the step-dependent marginal distribution of the chain (a unigram strategy) and later fit the bigram strategy of section 2.1? Isn't it simpler to fit the latter first (as opposite to the stationary case)? 
9. Is the 2-layer transformer model initialised to the matrices in Appendix C.1? 
10. The sequence length is first labelled by $t$ (e.g. after eq. 1) and then by $T$ (e.g. before eq. 4). Is this change of notation intentional? Or am I misunderstanding something?
limitations: limitations: Yes, they authors did address (some of) the limitations of their method.

Comments: title: Link correction
Response to Reviewer Comment
Reviewers, please respond to the author rebuttals.
Stationary Distribution Results
New Lemma for $k>2$
Question 4
Doubly stochastic and 'unigram' distributions
Minimal Model Design
comment: After considering your clarifications, I acknowledge that despite the multiple typos and minor concerns that affect the readability of the paper and the communication of the results, your work contains some insightful observations. These contributions are valuable and, therefore, I have decided to increase my score.
Following up to see what you think about our response to your review. Let us know if you have any further questions or need any further clarifications?
A minor fix to our rebuttal: the reference to Akyürek et al., 2024 should point to [this paper](https://arxiv.org/abs/2401.12973).
I thank the authors again for their responses. I will update my score.
@3: To clarify, in our setting $T$ is taken to be large, while $k$ is effectively a constant. This is similar to natural language settings, where the number of letters or tokens is some constant, but the context can be arbitrarily long. When we wrote that the  "marginal distributions... approach the stationary distribution exponentially fast", we were referring to Lemma C.4, which states that the distribution of the token n tokens after the current has distance to the stationary distributed bounded by $\alpha^n$, for some $\alpha<1$ depending on the specific chain.

@6: In our experiments, we used KL divergence to measure the difference between the probabilities predicted by the model and other probability distributions. For test loss, this other distribution was the appropriate rows of the transition matrices used to generate the test examples.

Formally, let $f(x_{1:T-1})$ be the softmax distribution of the transformer's output, given the input sequence $x_{1:T-1}$. In our standard setting, we measured
$$d_{KL}(\mathcal{P}\_{x_{T-1}} || f(x_{1:T-1}))$$
where $\mathcal{P}\_{x_{T-1}}$ is the true distribution of the next state $x_T$ given the previous state, under the true Markov chain $\mathcal{P}$. Note that $\mathcal{P}$ varies from sequence to sequence (it is drawn from a prior over transition matrices) and is not directly observable by the learner—this is what needs to be learned in-context.

For measuring how close the model was to various strategies, we computed the predicted probabilities given by said strategies, and used those as the base distribution. Note that the  output of the bigrams strategy (which is Bayes-optimal for our base setting) is different from the aforenentioned ground-truth $\mathcal{P}\_{x_{T-1}})$. Instead, as described in Section 2.1, it is a Bayesian posterior distribution of the next state given the observed sequence, with the prior determined by the prior distribution of transition matrices. Formally:
$$ \mathbb{E}[\mathcal{P}\_{x_{T-1}} | x_{1:T-1}] $$
where the expectation is taken over the draw of Markov chain transition matrix.

@8: Setting aside the mixing time issue discussed in @3, your question seems to suppose that all the training sequences are drawn from a single Markov chain, whereas in our work there is a new chain for each sequence. Do you have a more precise formulation of your question in mind? Regardless, we agree it could be interesting to explore what happens when the number of states is comparable to the sequence length, and the prior distribution over chains is crafted to favor chains that mix slowly, though this is beyond the scope of this work.
I appreciate the authors' detailed response. Before I update my score, I have a couple more questions/comments, if possible. I apologize in advance for these additional inquiries; my intention is simply to gain a better understanding of some of the claims made in the paper. I hope these questions also help the authors present their results more clearly.

@3: When you write: 

*"The marginal distributions of the tokens generated from a markov chain approach the stationary distribution exponentially fast no matter the initial distribution"*

do you mean for $k=2$ states? 

Because one can construct transition probability matrices which yield Markov chain that exhibit slow convergence to stationarity. A random walk with periodic boundary conditions is an example, where the mixing times are of the order of $k^2$, for $k$ the number of states. Or am I missing sth?

@6: Could you please remind me how you compute the KL divergence in your experiments? Meaning, you compute the KL wrt. what? 

@8: Let's assume that we are dealing with a Markov chain that exhibits slow convergence to stationarity. In such a case, the marginal unigram distribution changes at every step. Let's also assume one train your model on the *out-of-equilibrium* sequences sampled from said Markov chain. Do you still think the model will first find such a unigram strategy? I hope this rephrasing makes my question clearer.
Thank you for your comments. I will maintain my positive score.
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.
Dear reviewers, this is your AC! 

The authors have responded to your review. Could you please take a look at the rebuttal (and respond/edit your scores if need be)?
(Thank you if you have already done so)

Thanks!
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](https://arxiv.org/abs/0808.1502) 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.
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](https://transformer-circuits.pub/2021/framework/index.html)) 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.