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 2024 poster (2024)
Keywords: Deep Learning Theory
Transformers
Induction Heads
Phase Transitions
In-context learning
Markov Chains
Primary area: optimization_for_deep_networks
TL;DR: We characterize the emergence of induction heads in transformer models using the synthetic task of learning markov chains in-context.

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

The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains
Ezra Edelman"
Nikolaos Tsilivis' New York University! nt2231 anrlieeuuu.... edu
University of Pennsylvania ezrae0cis upenn. edu
Benjamin L. Edelman
Eran Malach Harvard University emal mmalach@g.oo harvard. edu
Surbhi Goel University of Pennsylvania surbhigècis upenn edu
Harvard University
bedelman@g. harvard.
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 order to study how this in-context learning capability emerges. our setting, each example is sampled from Markov chain drawn from a prior distribution over Markov chains. Transformers trained on this task form 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 's layers, and uncovering evidence that the presence of the simpler unigram solution may delay formation 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 2.
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 this emergent capability are not fully understood.
In effort to better understand ICI., some recent works propose study ICL in controlled synthetic settings-in in particular, training transformers on mathematically defined tasks which require learning from the input context. For example, recent line of works studies the ability of transformers to perform ICL of standard supervised learning problems such as linear regression B 20 26
41I. 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 only a special case of the more general phenomenon of networks incorporating patterns from their context into their
Equal Contribution
i Work done while visiting Harvard University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Transformer KL LL-Divergence: 3 Symbols
0.25
Uniform
K.-DviDistrbuutenl/mmddd
Unigram
Bigram
00100200202200
0.15
A
0.10
1011 01110001101111
0.00
25
1000
Number of Examples Seen (Thousands)
Figure 1: (leff) We train small transformers to perform in-context learning of Markov chains (ICL- MC). Each training sequence is generated by sampling a transition matrix from prior distribution, and then sampling a sequence from this Markov chain. (right) Distance of a transformer's output distribution to several well-defined strategies over the course of training on our in-context Markov chain task. The model passes through three stages: (1) predicting a uniform distribution (blue region), (2) predicting based on in-context unigram statistics (orange region), (3) predicting based on in-context bigram statistics (green region). Shading based on the minimum of the curves.
predictions. few recent works 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 this work is to propose and analyze a simple synthetic setting for studying ICL. To achieve this, we consider n n-gram models [T [ (3]. one of the simplest and oldest methods for language modeling. An n-gram language model predicts the probability of token based on the preceding n 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 different Markov chain, generated using different transition matrix (see Figure TT (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 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 transformers revealed the formation of induction heads [18] 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 D We show that these statistical induction heads lead to the transformer achieving performance approaching that of the Bayes-optimal predictor.
(2) Transformers learn predictors increasing complexity and undergo a phase transition when increasing complexity. We observe that transformers display phase Iransitions when learning Markov chains-learning appears to be separated into phases, with fast drops loss between the phases. We are able to show that different phases correspond to learning models of increased then bigrams (see Figure]] I 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 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 solution to the cunplex-and-corees solution happens due to an alignment between the layers of the model: the learning signal for the first layer tied the extent which the second layer approaches its correct weights.
1.1 Related Work
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 sparse function. The work of Bietti et al. B studies the dynamics of training transformers on a task that 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 sequence. Thus, successful prediction requires incorporating both global bigram statisties 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 solution.
Induction Heads. Elhage et al. 18 relates ICL with the formation of induction heads, sub- components 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. S study the effect of specific trigger tokens on the formation of induction heads.
Phase Transitions. 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. [4 shows another example of phase transition 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 phase transition when converging to the correct solution. Several works 25 27 attribute these phase transitions to rapid changes 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 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. M provide a thcoretical 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. (2 observe similar stage- wise learning behaviors transformers trained on language synthetic linear regression tasks. Makkuva et al.
3

---

28 study the loss landscape of transformers trained on sequences sampled from a single Markov Chain. Perhaps closest to our work, Nichani et al. [0] introduces 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 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-distribttion data. The focus of our work, instead, is on the different stages of training and how they relate to specific, well-defined, strategies.
2 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 distribution, with drawn from some other prior distribution (potentially dependent on P),, and the rest of æ = (x1, (#i,...,,2) XT) 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. Pt. ~ Dir(ox). We want to learn a predictor that, given context æ IT, predicts the next token, TT+1- Note that this is inherently -on-mttermiiistic 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 a = (I...., 1) that corresponds to uniform transition probabilitics between states. We draw the initial state xı 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.
{1, ... }} and transition matrix sampled from a prior
In subsection 3.3 we consider the generalization of this setting to rigrams for n >> 2. Instead of the distribution of 2t being determined by 2T-11 we let it determined by ær mn11.+..... ; according to a conditional distribution P which is uniform over the possible stated
2.1 Potential Strategies for (Partially) Solving ICL-MC
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 posterior belief about the stationary distribution. Unfortunately, the stationary distribution of this random Markov chain does not admit simple analytical characterization when there is a finite number of states, but it can be estimated approximately. At the limit of k + oo, the stationary distribution converges to the uniform distribution [T].
2nd strategy: Bigrams. For any pair of states and j. let P.j be the probability of transition- ing from i to j. On each sample æ, can focus on the transitions from the i-th state, which follow a categorical distribution with probabilities equal to (Pa...., P.x). ). If we observe the in- context empirical counts (Ctz ), of the transitions, then Pij is given by: (P.i...., la ~ Dir(k, en1 + a1,.... Gik + (gk,, where a1,.... ag are the Dirichlet concentration parameters of the prior. Hence, each Pıj has (marginal) distribution that is actually a Beta distribution:
Puylæ Beta (cg + aj, Σ, aj N₂ aj ci). where N; is the total number of observed transi-
'In particular, for each tuple of n 1 tokens, we sample the vector of conditional probabilities for the next state from flat Dirichlet distribution.

---

Attantion Patterns
12
Socond Lowd
Hest Laya
Figure 2: Attention patterns that correspond to the last token of the sequence for a transformer trained perform ICL-MC. The intensity of each blue line signifies the strength of the corresponding atten- tion 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 the first layer 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: [|]y/ = N42, on For the uniform Dirichlet, a = (I,...,1) 1),, it IE [Pgla0]
eutay
N
2.2 Architectures: Transformers and Simplifications
We are interested in investigating how transformers [40] can succeed in in-context learning this task. We focus attention-only transformers with layers with causal masking which is popular architecture for language modeling. Given an input sequence æ, the output of an n-layer attention-only transforme
TF(E) = Po (Allm, + 1) ..... (Allm + E,
(1)
where E E RTxd is an embedding of æ E R P E R is linear projection to the output logits, and Attn(x) (x) is masked self attention with relative position embeddings (38). which parameterized by Wo, WK. Wy Rgdx Txd.
A. (s,Wo)(s,Wa WK + us --)++ Vd
Attn(z) = softmax(mask( (
In general. transformers often contain MLP module, but for this task they are not necessary (see Appendix and Figure additional experiments with transformers with MLPs). During training, we minimize the loss:
T l ().
I(O)=
E. map P=)
Σ
0).
where 0 denotes the parameters of the model and / 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). 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 Appendix B.I
Simplified Transformer Architecture. As we see from the construction, there are two main ingredients in the solution realized by the transformer; (Ist layer) the ability look one token back and (2nd layer) the ability to attend to itself. For this reason, we define minimal model that is
For simplicity, assume embedding and hidden dimension equal, but they be different in general.
5

---

Test Loss
Positional Encoding
KL Diviiiitrrbuuiinn (mmadcl
Transformer KL-Olvergence: Symbols
0.3
[-64
tw3
Uniform
WM wwww MM
KL-Dv Loss
0:15
0:20
0.00
0.05
0.2
Bigram
Unigram
t-64
0.10
0.00
0.1
0.05
t-127
0.00
25
50
75
100
125
25
50 Position
75
100
0.0
Number Examples
(Thousands)
Number of Examples Seen (Thousands)
Distance between model predictions
candidate strategies
Test Loss
Positional Encoding
Minimal Model layers
0.01
0:20
t=76
0.25
KL-Diy Loss
0.15
0.20
Uniform
0.025
38
0.15
Unigram
0.000
0.10
Bigram
0.05
0.5
0.03
0.00
20.0
40.0
60.0
50
75
0.00
0.0
20,0
40.0
60.0
80.0
Training Sequences Seen (Thousands)
Positian
Training Sequences Sean (Thousands)
Figure 3: A two layer transformer (top) and a minimal model (bottom) trained on our in-context Markov Chain task. A comparison of the two layer attention-only transformer and minimal model (4) (with e and W initialized to 0). The graphs on the left are test loss measured by KI.-Divergence from the underlying truth. The orange line shows the loss of the unigram strategy, and the green line shows the loss of the bigram strategy. The middle graph shows the effective positional encoding (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 ex; denote the one-hot embedding that corresponds to the state at position E [2], and let E be the [((++)x one-hot embedding matrix. Then the model is parameterized by W E R kxk and v E I T+ and defined as:
10
f(E) = mask(E5W(SfftmuxxM))))) M =
01
10
0
E (T+++)((T+))
UT
UT
(4)
where mask (.) is a causal mask, and Softmax(M)i, 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 transformer. See also Appendix B.2 for a discussion.
E. exp(Mr.,)) Notice that the role of W is
(xp(M
simplified version of two-layer linear attention-only
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
and W
Lacks then f(E)T.s =
ET, 1x = [x (I--1 xTI +
kT exple)
Unigrams: For e = (0
0) W
1
we have f(E)T,. = =I 1 {xz = s).
See Section B for the proofs.

---

Test Loss Origi nal Distribution
VINCOM
25%4 Seuby Steshanti: 25%4
0.30
beudry Stothatt
KL-DIv Loos
0.25
0.02
0.15
0.03 1-
0.10
0.80 0-
Doubly Stochastic Sequences Seen (1000s)
20
320
Figure 4: (left) Unigrams slow down optimization: Comparison of two-layer attention only transform- ers trained on two distributions; one with uniformly random doubly stochastic transition matrix and another with mixture of the doubly stochastic and unigrams distribution (see Appendix [A.] for details). We see that in absence of unigrams "signal" the model minimizes the loss (evaluated on the full distribution) much faster. (center, right) Training of the minimal model on ICL-MC with k = 2 states: (center) The heatmap of the second layer (W matrix) that learns to be close to diagonal. (right) The values of the positional embeddings (1st layer) that display curious even/odd pattern. This is before any softmax applied to the positional embeddings.
Empirical Findings and Theoretical Validation
In this section, we present our empirical findings on how transformers succeed in in-context learn- ing 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.
3.1 Transformers In-Context Learn Markov Chains Hierarchically
We focus on attention-only transformers with layers with causal masking and relative positional encodings and train them with the Adam optimizer on ICL-MC (see Section 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 second rapid drop brings the model close to the Bayes optimal loss. We observe that training a I-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 phase when the model reaches the unigram baseline. We provide evidence that this is not 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 3T; special dedicated heads inside a transformer are suddenly being formed to facilitate in-context learning. Similar observation hold for Markov Chains with larger number of states see Figures and 11
Mechanistic evidence for solutions found by transformer. To confirm how the two layer attention- only 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 tokens that follow the same token as the current one.
Cooee: https://giiee http:///ww...eeeaaaaaaaaaoaooaoa Inrrell-Innntnnno--eee Heads.
7

---

Varying the data distribution Unigrams slow down learning. There are several interesting phenomena in the learning scenario that we just described, but 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 il perhaps just by-product of the learning procedure? To answer these questions, we define distributions over Markov chains that are in between the distribution where unigrams Bayes optimal, and the distribution where unigrams is as good uniform see Appendix Afror more details. As we see in Figure | the transformers that are being trained on the distribution where there 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.
3.2 Theoretical Insights from the Minimal Model
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 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 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(O) O, v())) O. Suppose the transition matrix P E BXxk is sampled from of the following two types of distribution:
1. 2, 2, and P is sampled from the uniform distribution over the set of x 2 stochastic matrices.
2. For any constant is and < a < 1, with probability a, sample the matrix P uniformly from "bigram-only distribution the of k x k doubly stochastic matrices: and with probability 1 - & use a tunigram-only distribution": draw a vector u uniformly from the set {u TR ||u||] = 1} and let P lu
Then after one step of population gradient descent with step size 1) > O,
W(2) = I
+ E and (1) =
where || S On log T). other words. for large T, the second layer weights are a mixture of the correct solution and uniform attention 11
Assuming in the first step 1/1 = (+t), then (2)) has the same structure as ( (up to scaling). Furthermore,
(2) 2(12 log T), and (2) cu(2) (2) Vn + 1,051
where 1/2 is the step size for the second step.
n2 (T)), then the output of the model will be weighted sum of bigrams and unigrams strategy. Formally,
T
T
f(E)T, = (( Σ [x:-1 It. x4 = 8) O(neT) Σ [z = +
In the first distribution (uniformly random 2 2) or the second distribution with & > 6, at the end of the two steps, the weight on bigrams is greater than that of the weight on unigrams strategy.
Note
Proof Overview: The idea of the proof is that first step of gradient descent with 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 strong signal in the gradient to look back one in the first layer. Without a bias W, the gradient for the positional encodings, v, zero.
8

---

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 optimal, while in the second step effectively all of the gradient comes from the examples where the bigram strategy 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 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 (+) is sufficient for learning the second layer. In the second phase, much larger step size of (2(1) needed in order to learn the positional encodings (in one step).
Second layer is learned first. It has been observed before in similar bigram learning setting with 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 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 (top. center) and Figure 4 (right). We believe that a careful analysis the gradient of in the second step will recover this pattern, which is likely related to the moments of the eigenvalues of the transition matrix.
3.3 Beyond Bigrams: n-gram Statistics
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 laye by minimizing the in-context cross entropy loss with the Adam optimizer. As can be seen in Figure (left), the model eventually converges the Bayes optimal solution. Interestingly, as in the case of Markov Chains, the model displays a "hierarchical lcarning" 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 3 where we plot the similarity of the model with the different strategies and it exhibits the same clear pattern in the case of 2. We leave a more thorough investigation of n- n-grams for future work.
4 Conclusion and Discussion
In this work, we have introduced a simple learning problem which serves 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.
The main limitation of our work is that our analysis relies 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 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 lcarning with natural language data, and use insights from minimal model to improve formation of induction
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 single head in each layer suffice.

---

Transformer Trigram: 3 Symbols
(Disttiiitiibuttoon [[motel)
Transformer Trigram KL-Divergence: Symbols
0.25
Uniform
0.25
Unigram
KL-Div Loss
0.20
Bigram
0.20
0.15
Trigram
0.15
0.10
0.10
Unigram
0.05
Bigram
0.05
0.00
25
50
75
100
125
25
50
75
100
125
Number of Examples Seen (Thousands)
Number of Examples Seen (Thousands)
Figure 5: Three-headed transformer trained on In-Context Learning 3-grams (trigrams), with context length 200. (left) Loss during training. The model hierarchically converges close to the Bayes optimal solution. (right) KI. divergence between the model and different strategies during training. As we observe, there are stages of learning, each of them corresponding to a different algorithm implemented by the model.
heads. In particular, it would be great understand if better data curriculum could remove the undesirable simplicity bias we observe from unigrams. Such simple but incomplete solutions may be commonplace 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.
Acknowledgements. EE thanks Alan Yan for helpful conversations. NT acknowledges support through the National Science Foundation under NSF Award 1922658. NT would like to thank Boaz Barak, Cengiz Pehlevan and the whole ML Foundations Group at Harvard for their hospitality during Fall 2023 when most of this work was done. BE acknowledges funding from the ONR under award N00014-22-1-2377 and the NSF under award IIS 2229881. SG acknowledges support through the Open AI SuperAlignment Fast Grants.
References
[1] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning neural net- works: laap complexity and saddle-to-saddl dynamics. The Thirty Sixth Annual Conference on Learning Theory, pages 2552-2623. PMLR, 2023.
[2] Jacob D. Abernethy, Alekh Agarwal, Teodor V. Marinov, and Manfred K. Warmuth. mecha- nism for sample-efficient in-context learning for sparse retrieval tasks. CoRR, abs/2305.17040, 2023.
[3] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211 15661, 2022.
[4] Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architectures and algorithms. CoRR, abs/2401..12973 2024, doi: 1044850/ARXXV.2400.29999 URL https://doi. org/10..//0 158550/arXiv> 2401 12973
[5] Devansh Arpit, Stanislaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization deep networks. In International conference on machine learning. pages 233-242. PMLR, 2017.
[6] Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neurul Information Processing Systems 35: Annual Conference Neural Informa- tion Processing Systems 2022, NeurIPS 2022, New Orleans, IA, USA, November 28 December 2022, 2022.
10

---

[7] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21764 2022.
[8] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: memory viewpoint, 2023.
[9] Garrett Birkhoff. Three observations on linear algebra. Univ. Nac. Tacuman, Rev. Sen. A, 5: 147-151, 1946.
[10] Charles Bordenave, Pietro Caputo, and Djalil Chafai. Circular law theorem for random markov matrices. Probubility Theory and Related Fields, 152, 08 2008.
[11] Peter F. Brown, Peter V. deSouza, Robert L. Mercer, Vincent J. Della Pietra, and Jenifer C. Lai. Class-based n-gram models of natural language. Comput. Linguist.. 18(4):467-479, dec 1992. ISSN 0891-2017.
[12] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[13] Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emer- gent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878-18891, 2022.
[14] Angelica Chen, Ravid Shwartz-Ziv, Kyunghyun Cho, Matthew L. Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in mlms, 2023.
[15] Noam Chomsky, Three models for the description of language. IRE Transactions on information theory, 2(3):1 13-124, 1956.
[16] Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey for in-context learning. arXiv preprint arXiv:2301.00m 100344, 2022.
[17] Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms, 2022.
[18] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
[19] Robert Gallager. Discrete Stochastic Processes (Draft of 2nd Edition). MIT OpenCouseWarce 2011.
[20] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What transformers learn in-context? A case study of simple function classes. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Anmual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
[21] Jesse Hoogland, George Wang, Matthew Farrugia- Roberrts, Liam Carroll, Susan Wei, and Daniel Murfet. The developmental landscape of in-context learning. CoRR, abs/2402.02364 2024. doi: 10448500/AXX1.2240000225 URL http:t:t/tii..orreeeee 11tp:/...rr//1.....0..... 400.00.8855/8880 /ato/arxivvv/.. .i.vvv. 220222220006444
[22] Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kernel: Conver- gence and generalization in neural networks. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò
and Roman Garnett, editors, Advances in Neurul Information Processing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NewrIPS 2018, December 3-8, 2018, Montréal, Canada, pages 8580- 2018. URL https: httpt/tppeppp.oeedns....oo nwwr wwwpuripsscc/pppeee wew.rips.cc/ppppp//2111 httppppppp22211////a2 Salleeffs4ee20886666611122222222 http:///22255aaattattttooooo Abstract.html

---

[23]
Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances in neural information processing systems, 32, 2019.
[24] Andrej Karpathy. Mingpt. https:// http///.tttubb... httgg/gtttttt..aaaieeeeeee
2023.
[25]
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking the transition from lazy rich training dynamics. CoRR, abs/2310.061 10, 2023. doi:
10,48550/ARXIV.233 13.0061111111111 10. https::t::///o:: http:://doo..rrg1/0// 485500 arxiv
[26] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pages 19565-19594. PMLR, 2023.
[27] Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon S. Du, Jason D. Lee, and Wei Hu. Dichotomy of early and late phase implicit biases provably induce grokking. CoRR, abs/2311.18817 2023.
[28] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. CoRR, abs/2402.04161, 2024. doi: 10.48550/ARXIV.240022 041611.
[29] William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking as competition of sparse and dense subnetworks. CoRR, abs/2303. 1873, 2023.
[30] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402. 14735, 2024.
[31] Catherine Olsson, Nelson Elhage. Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https//traasfoomerirrrciss wttp:pppoo200//mooootteeeeeeeeee heads/index.html.
and-inductiionn
[32] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting small algorithmic datasets. arXiv preprint arXiv:2201.02117 2022.
[33] Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vard- han Makkuva. Transformers markov data: Constant depth suffices. arXiv preprint arXiv:2407. / 16886, 2024.
[34] Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task, 2023.
[35] Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models a mirage? CORR, abs/2304. 15004, 2023.
[36] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateck Jain, and Pranceth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573-999 95859,5552222 2020.
[37] Claude Elwood Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3):379-423, 1948.
[38] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position repre- sentations. In Marilyn A. Walker, Heng Ji, and Amanda Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT, New Orleans, Louisiana, USA, June 1-6, 2018, Volume 2 (Short Papers), pages 464468. Association for Computational Linguistics, 2018
[39] Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the garameter-functionn map is biased towards simple functions. arXiv preprint arXiv:1805.08522. 2018
[40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von
12

---

Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30; Anmual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008, 2017.
[41] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310008391, 2023.
[42] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022.
13

---

A Experimental Details and Additional Experiments
A.1 Experimental details
We train transformers of the form ① with the AdamW optimizer with learning rate 3 x 10 (for 3-grams a learning rate of 3 10tt 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 10-4 4. We use PyTorch 2.1.2.
The data was generated in an online fashion, using numpy.random.dirich dirichlet to generate each row of the transition matrices. At each epoch, we generate new transition matrix, and, then, 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 perfectly reproducible manner).
Some of the training and model code was based minGPT 24 The experiments all measure the outputs of the models at the last token.
All of the experiments were performed with 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: thub. cow.u...oommeeeeeeeeeeeeeeeeeee. http://w......aaaaaaaooaaonnaaaa
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 x 3 transition matrices for this particular experiment):
1. The "doubly stochastic" distribution uniform over the space of doubly stochastic matrices, that is, transition matrices for which each row or column has entries summing to 1. This equivalent the space of Markov Chains with uniform stationary distribution (that is,
E 1), which means that the unigram and uniform strategy are the same.
2. The "unigram distribution" uniform over the space of stochastic matrices such that for any two rows P., P, in P. P. Py- This is equivalent to the distribution over markov chains for which the distribution of the next state doesn't It change depending on the previous state. Notice that the unigram strategy is asymptotically optimal 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 the blue line. Then, trained models 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 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 stochas- tie matrices, which the models in the orange line do seem to generalize slightly better to after convergence.
Note on KL-divergence 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 (left),, this other distribution comes from the appropriate rows of the transition matrices used generate the test examples.
Formally, let f(x1TT-1) be the softmax distribution of the transformer's output, given the input sequence 11:T-1. In our standard setting, measured
(P ((
where Par-1 1 is the true distribution of the next state ar given the previous state, under the true Markov chain P. Note that P varies from sequence to sequence (it is drawn from prior over
14

---

Transformer: 3 Symbols
Test Loss Distribution For 10 Random Seeds
0.30
0.25
0.20
KL-Div
0.15
0.10
0.05
25
50
100
125
Number of Examples Seen (Thousands)
Figure 6: In distribution test loss for 10 two layer attention only transformers, with random seeds 0, 1, (randomness affects initialization and the training data). The training dynamics are consis- tent for each model, though the exact position of the phase transitions changes.
Single Layer Transformer: 3 Symbols
Single Layer Transformer: Symbols
0.30
0.30
0.20
KL-Div Loss
0.20
D.Div
0.15
0.10
0.10
0.05
0.05
204
612
1020
204
408
612
816
1020
Number of Examples Seen (Thousands)
Number of Examples Seen (Thousands)
Figure 7: Graphs of test loss showing that single layer transformer can not achieve good performance on ICL-MC. This result holds for transformers with without MLPs, and with absolute relative positional encodings. These graphs show that even trained times longer, there is no notable increase performance beyond the unigrams strategy (orange line).
transition matrices) and is not directly observable by the this is what needs to be learned in-context.
For measuring how close the model was to various strategies (for instance, in Figure (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 Par--). Instead, described Section 2 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 to E(Pay-, [11:-1)), where the expectation is taken over the draw of Markov chain transition matrix.
A.2 Additional experiments
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 states.
Figure g shows that our results extend for Markov Chains with larger state spaces (here, & = 8). As the number of states grows, larger sequence lengths are needed for learning (this is to expected as larger transition matrix needs to be estimated in context roughly, sequence length needs to be !2(A²)).
15

---

Transformer KL-Divergence: Symbols
0.25
0.20
Uniform
0.15
Unigram
Bigram
0.10
0.05
0.00
25
100
125
Number of Examples Seen (Thousands)
Figure 8: Our results extend to more symbols than transformer and strategies over training. This required a sequence length greater than 100 (200 in this case) for the difference between unigrams and uniform to be large enough for the unigram phase to be visible (in either case there plateau before the final drop in test loss).
or = 3. The KL-divergence between
Transformer: Symbols
Test Loss
t=121
0.30
3-3
2-92
121
First Layer
1.0
1.0
1.0
0.25
KL-0iv Loss
0.15
1.0
1.0
Second Layer
0.10
0.5
25 50
75
100
125
Examples (Thousands)
Figure 9: A two layer attention-only transformer trained with cross entropy loss ICL-MC. The heatmaps on the right represent part of the attention for the transformer at various time steps, specifically the values of the matrix A from (2) The top are showing A from the first layer, and the bottom row from the second layer.
We also know that, in this particular family of Markov Chains, the stationary distribution approaches the uniform distribution as the number states grows. As result, we expect the difference between uniform and unigram solutions to be less noticeable.
Figure shows the attention patterns for the two layers of transformer during training.
In Figure 10
we demonstrate that our observations extend to two-layer transformers, which have an additional fully-connected MI.P on top of the attention layers.
Finally, in Figure [2] we plot the performance of the transformer and the minimal model in various distributions over the course of training.
B Proofs
In this section, we present our theoretical results on in-context learning Markov Chains Section
2.
16

---

Distance between model predictions and candidate sir strategies
Test Loss
Positional Encoding
KL DiviDitriiutioonn [model]
Relative Transtormer with MLPs 2 layers
4-4
[m96
t=191
-0.005
0.000
Ay H Mypvyy Muy
0.3
Unigram
0.025
0.2-
Bigram
KL-Div
0.000
Why
0.1
10
t=191
0-
192.0
0.0
Training Sequences Seen (Thousands)
Position
Training Sequences Seen (Thousa nds)
Figure 10: A two layer relative position encoding transformer with MLPs trained on ICL-MC with k=3 symbols. Notice while slightly noisier, the overall trend and observations made regarding the attention only transformer still hold.
Two Layer Transformer: 3 Symbols
Positional Encoding
Test Loss
t=3
1-92
H11
0.01
Transformer KL- (L-Divergence: 3 Symbols
Vaitorm
0.00
Laigram
0.2
grram
0.025
KL- KL-Div
0.35
0.1
0.000
t= 121
0.12
1
0.05
100
125
50 Position
100
Number of Examples Seen (Thousands)
Number Exampias Soan [Thousnnes)
50
Distance between model predictions
candidate strategies Minimal Model
Test Loss
Positional Encoding
3-38
0.01
[=1
Uniform
KL-Div Loss
0.2
0.00
0.20
Unigram
Bigram
0.025
0.15
0.1
0.000
0.10
0.5
1-76
0.05
0.0
20.0
40.0
60.0
80.0
0.0
25
50
75
100
0.00
0.0
20.0
Training Sequences Seen (Thousands)
Position
Training Sequences Seen (Thousands)
Figure 11: A comparison of the two layer attention only transformer and minimal model for k symbols.
B.1 Transformer Construction
Proof of Proposition 2.7 Set the internal dimension d 3k, and choose ea to be one-hot embeddings- that is, ez; 82, where the Kronecker delta. We will call the parameters of attention layer i, w (i) w(( w(( v(0) Let
I k
cx
0
0
0
0
WO =
0
W =0
w(
0
Q
0
0
0
0
0
0
So,
10
Vd
Notice that 100 ] = i- 1). So, (ettmaxmmsskA)))) ~ 1 -- 1] for large enough c. So, for any 2 i < T,1 و < k Allm (e)is+* ei- 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 (Cx, 6 0).
17

---

Two Layer Transforme: Symbols
Minimal Model: Symbols
Test Loss Distributions
Distributions
0.30
Unigram Uniform Stationary
0.25
Unigram
Uniterm Stationary
KL-DIv Loss
0.20
Loes
0.20
0.15
0.15
0.10
0.05
0.05
25
50
75
100
125
20.0
40.0
60.0
Number Examples Seen (Thousands)
Training Sequences Seen (Thousands)
Figure 12: two layer attention-only transformer (top) and minimal model (4) (bottom), trained on the main task with ICL-MC with cross entropy loss, test loss measured by KL- KL-Divergence from the underlying truth (labels based on transition probabilities, not samples). The distributions test loss is measured in are (from left to right) in-distribution, distribution where each token is sampled i..d., and a distribution over uniformly random doubly stochastic transition matrices (that is, setting where the stationary distribution is the uniform distribution, which implies that unigram based guesses are as good as guessing uniform probability). For both models, the distribution test loss quickly drops the level of the unigram algorithm.
The second layer is defined as follows:
cI bock
0
0
(2) 0 W
0
0
(2)
1
(2(2)
0 0
0
0
0 0
Note that = = + Attni (e), then
(2,1112) (2)
(2,
ce2, (ez)
1
Q
|| [x
xr,).
Vd
Vd
Vd
So, for all < i, softmax(mask(AA (A))ss 20
1 for large enough c. For any 2 2 i <
Eimr 1[2n-1=2.]
T,15j< k.
3k 1 1a mi]]
Σh 1 [ih-1 x; 1 = j
Attng(e)i
(=W(²) hag
2g=1 18gg11 = x.]]
ΣS=1 [(ig-1 = vi]
1=1
This is exactly the empirical bigram statistics (that is, the number of times I; - appears before position i). In order to make this the output, we set P
0
pk
B.2 ICL-MC with Minimal Model
Setup and notation Our data consists of sequences of length T + x (xo, "oo.....,rr).. drawn from Markov Chain with state space S (1,...,k, k) (i.e., æy (1....AA k) for all j (T), and a transition matrix P. Each row of the matrix is sampled from a flat Dirichlet distribution, i.e. the simplex. Let P. Dir(1), corresponding drawing the from uniform distribution
Technically, the output this construction not the log probabilities generally cross-entropy loss
61
assumes. These can be approximated linearly by setting P
change the output from + b.
a a x
In practice, this approximation achieve close Bayes optimal loss.
18

---

E E {0, {0.1)(7+1)x be the one hot embedding matrix of that is, Ei,,, = and for all * a;
E = 0. 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 [00 the penultimate token sampled from the uniform distribution).
Model We define our model a simplified sequence sequence transformer f R Tx B(T+1)x with f() mask(EW (Esk(EW [)))t E. The trainable parameters are W 1
Uo
-00
-00)
V1
20
and E T+1 We define .M& E [(T-I)x(T+1) as M
that is, for all
VT-1
r2i2j2 0, My, = o.- and if j. Ma = -00 Furthermore, v [00. U2 UT! R Softmax defined as follows:
exp (M.,)
Softmax
Σt= exp (Mis)
The logit for symbol at position T for our model is:
&:
T
exp(v;)
I(E),,,
W (
(5)
Expeepprr)
1=0
The model represent the unigrams and bigrams solutions following:
Construction for bigrams:
(0, .) .) and W TExK, then f(E)r,
Σ [
FT + O exp(e)
T
As tends to infinity, this becomes bigrams.
Construction for unigrams: 0 and W 1 1, then f(E)r,, Σ jx; s).
We prove the above claims.
Proof of Fact 2.2
Proof. We will first prove the unigrams construction.
exp(v;)
f(E)T,s =
W u. [x Ax, =x =
Eizo exp(ve)
= =
I
1/. = s]
which is exactly the unigrams solution.
Now consider the bigrams construction. As grows, the softmax of very quickly becomes one hot. Formally, by Lemma B.7 in 17 for any > 0,
exp(u;)
T exp(c)
1 [ 1]+
Ei exp(vr)
19

---

So,
k Σ
T
W2y,u [ = uA = exp(uy)
f(E)r,,
W
Eiz exp(vr)
u=1 1=0 j=0
T
exp(1))
[i-1 = IT A A;
Ei=o exp(1ee
1=0 j=0
exp(1s)
=1
Zo = s)
1 FT Σi-j A 2 =
+ exp(c))
,==
T
( 27
(<xp(e)) (Lemma B.7 in 17
=1 = 20 s
+
FT = = A2;= = [ = 1
i= 1==
T
T
2T exp(c)
1[1 = 2o = s|+
FTT = 1 =|+
+
[r = I az s|0
:= ==
T
T3 exp(c)
12r 2 1-
This simplified model was constructed by taking 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 ( 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 LT(f, E, = Σ e==1 log Softmax (J(FE))T, Pxr.
B.3
Gradient Calculations
For 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.
0f(E)r. DWab
exp(v;)
Σ
1 PTT A b w|1[x;1} u u.
s)
Eizo exp(v1e)
T
cxp(u;)
1 |
-|| = bb s)
Ei=o cxp(1i) Uz)
0 0=
Of(E)r,
Dva
(11
[[ [ijj- a) Σi Ei=o Ei exp(ve)
exp(va)
exp(v;) exp(va)
W W 1 u = =
1aa S
Ei=o exp(ve) Σi = exp(ve)
1=
exp(va) (up(
(Mliio-ua = s}1زi = 1[1 = u u A s] - il
exp(U;)
War. 3=0
Σi exp( (ve)
1=
1==
*In practice, one could use the empirical value of 41 rather than its distribution Pxy... but in full batch gradient descent this is in fact equivalent. This is because conditional on xy and P. x741 is independent of
20

---

exp(Ua)
exp(13)
Warn
1 = s S
Σ
=
8]1(a 1)
0 exp(ve)
Ei exp(ve)
5=
exp(va)
exp(1;) Ei=o exp(ve)
War,
=
Σ
uu Un =
Ei=o exp(1))
exp(Uia)
exp(u),
W Wor = s| 1 = - Σ
[ u
t= 2i=0 exp(v:)
Ei=0 exp(Ux)
j=0
0hr Softmax(f(E))r== Par.s (f(E)..
B.4
Proof of lemma 3.1
Proof. Recall that at initialization, and W 0, implying further that f(E) = 0.
First step.
First consider the gradient of the loss with respect to W. By chain rule,
(Lr(F)
0Lr (f(E)r
aWab
TS(E)T, aWah
k
T i
1 | = [ = exp(1;)
Σ
(Softmax(f(E))T P
Σ Σ
Eizo exp(ve)
== j=0
k
Σ
PZzis
1[2T = a
1 bb - = i+1
1
j=
T
1
PT
a)
[ = b1 1+1
Pa 1 [TT = a
[i:s = b AAx; x1 = 1 =1,741
k
1=0
Y=0
1=0
T
k Σ
i
|
al
Σ
1 = b]
Pa.e
[ET = a Σ
= b -
11
x2 == i+1 =
k
1=
s=1
1=0
1 T
k
1|2, = bA = a]
Σ
Pa.e Σ
1
k
i+1
Y=0
#=1
j=0
T'
k
b Σ
Pa.s
1[20 = -b = -
7+1
Where the last line follows from the Markov property.
Now we take the expectation over x, TT41 conditioned on the transition matrix P.
ØLT
EgiP
WWab
Tb
(p
b
Σ
Pa
i
(2T-1)
(P))
(P')
b.s
1=0
T-i
(P*) h.
Pa T -i+1 (PO)
(pi)
Tb
b,s
5=
k
= =TbRa
(T
Pa. + O(logT)
k
s=1
21

---

Where the last step follows from Lemma B.6 Then, by applying Lemma B.1 (for the uniform over all 2 transition matrices case) or lemmas B.3 and B.4 the mixture distributions case). there exist positive constants (potentially depending on k, but not 7' A, B such that for all a
ahr
E
BAT+ +
aWa.a_
and for all # b,
åLT dWabb
En P
BT + O(log T)
ØLr The updated Wab after the gradient step is just - EnI Choose 7Ι O (7), so that W will be (1)) with respect to T after the first step. For the gradient with respect to v, since W = 0, OF(E)T, Dv
[Wa,b (because W is initialized at 0).
exp(va)
[x u A = =s] exp(u;)
W Warn Ei=o exp(ve)
1[1;-a = uA =s =||
Σ ==
Σ exp(us)
u=
ima
0
So.
OLr(E)
OLT(E) OF(E)T,
0
Ov
Of(E)r,
Ou
Completing the first step calculations.
Second step. After the first step, W 1/1 (AT + 1)) Now us bound the output of the model,
exp(U;)
((F)r,
W 27.u [ri-j
Et=o Ei-oexp exp(11)
1=
m (AI + B1 11) zr.a 1[x ^
1=0 j=U
k
S
(A B) 1[; =
k T
1=
T
<m
Σ
(A+
B)
j=
< 7 7
So, using the first order approximation of softmax,
(E)
Softmax (f - [IT+I=5)
(S(E)r,
[(F)T,
Σd [(F))T, +0(f(E)
+
[xT41 = s)
1.2
+0
T
(A+ B) + +A B)2) [ = s]
k
k
22

---

T
+0
1 k
(A + B)
+0
1 [xT+1 = s)
1
1 [TT+1 s) + 0
T
Where the last step follows since for the first step, eta - 0 T7
Now, we can begin analyze the gradients with respect to the parameters. For W, the gradient is approximately the same as in the last step. Notice that 0f(E)s. does not depend on W, and is D.a.
unchanged, so aWab
(f())T is unchanged. Furthermore,
T
exp(u;)
[xr = [x) bbR Ei exp(ve)
oWab
==
1
[xT a| bl
==
T
1=
= T
We will now show that the gradient is approximately the same in the first gradient step:
(Lr(E)
ØLr (f(E),,
OWab
0f(E)r, aWab
==1
1 [2T+11 = s) + 0 (+))
(f(F)r,
Σ
=
s=1 k
aWab
(1
(S(E)r,
(J(E)T..
Σ
[xT+1 =s]
UWab
T
oWabb
k Ps + O(log T)
n (T 1)
Where the last lines follows from the gradient calculations the first step.
Now we will consider the gradient with respect to v. First, notice that the uniform component of W, B1 1, has affect the gradient of
Of(F),,
exp(va)
exp(U;)
W
1 [0: = s]
1 u)
Ova
2T.
Dt=o exp(Ur))
=t= exp(ve)
i=a
T
1
(ml + B171) ArrM i+1
1 = 8)
1 u]
11 [1i =u] u]
+
a=1 t=
1=0
T
1
(A1 = u| + B)
1 1/, =s] s]
1 = u]
1 [2 = 2
it
i 1
#=1 imo
jm
T
4 1+1
11 [x; = s]
1 1|21-0 = u] - Σ
[xi_] u1
i
#=1 t=
j=0
T
B
1 1z = s)
1| = u] -
Σ
1 [ = ul
i+1
i +
1=0
23

---

[T
1/ i+1
1
11 [x 8]
1 = u|
Σ
| m)
j=U
+
B
1
[r; s
Σ
1
aiia = - Σ
Σ
[ ml
i+1
j=0
u=1
7
1
[T u|
11 [ s]
1|2 ul
1 | m)
==1
1=0
j=0
T
1
+
B
[r,
=
1
i+1
itl
1=
k
T
[TT
4
11 [z; = s]
1 = w/
Σ
[ii ul
1
i
M=
1=0
j=0
By chain rule,
ØLr
k
0LT Of(E)r,
vvo
Tf(E),,
Ova
s=1
k
Parts + +0
-
1
1
[rT = u|
1 [0 s]
1|2i-a = ul
Σ
i i+1 1 1 = ul
T
1 i
==1 i=a
==
k
Pariis
Σ
k
Σ
T
= u|
[xt = s]
1xiia = u]
Σ
i
[xi-] = u)
+0 (
log
1+1
i i
T
s=1
a
im
j=0
Where the last step follows because
Σ 1
1
[ = u1
1 [ri
s]
1[zi_a u]
[ ui]
5
[x
2ma
Σ
log T
In expectation the values of conditioned on the choice of P:
[dL
log T T
Eg|r
P
(pT-i)..
(Pa)
)
(pi-j)
+0
Ovo
k
su
1.
s=1
a=1
t=o
==
Pass
πu
(Pi)
(Pa) u,
1
T-
(P))
+0
log T T
ke
T
i +
9,8
T
u.6
aml
ima
j=
(log (T + +1 log (a + 1))
2 P u.s
T
(Pa) 0
Where the last step follows from Lemma B.7 (depending on the distribution assumption on P).
Then, by applying Lemma B.2 or lemmas B.3 and B.4
TLL
ØLT
Egip
Ex
Ov
Fc.
24

---

Therefore, after the step is taken,
(12 log 7)
U1 Un (log T)
Finally, we can consider the state of the model after the second step. Assume that the step size for in the second step is ((T), and the step size for W is TTA+B)
T
1[i-j = u uAI== = exp(
f(E)T,
Σ
W 2r
Waral
Ei= exp(Ut)
==1
1=0
==
log Al 1+0
1[x = u = [1]j=1|+0 (exp(log(T))
1|
2T
A+ B
T
1= 5=
llog T
(16- 11 + o ())
ΣΣΣ
A 1++B10+0
A
==
27,
[Al 11 2r.a
1[
A
0 7
A+ B
1=0
A
B
k
T
1
1 1 = = Olog
A+
A+
T
T
+
Σ
1[2 = s) Tlog
A +
B
A
B
1=0
1=0
Since A and B are constant in terms of T, we can recover the desired statements, completing the proof.
B.5 Inequality lemmas for k = 2
Lemma B.1. If a uniformly random stochastic 2 matrix, and π the stationary distribution of P, then
1
5
2
π
PT
log (2) ~ 00.045
k
12
3
==1
and for any b +
k
Pa.oT ]
1
5
E
To Tb
+
log(2) ~ 0.011
k
6
3
1=1
Proof. We have:
k
]
(b 1)2
a(b - 1)
(1- 1- - 1)
π
Po.
Ea,
k
(a+ b - 2)2
2 a b-
atb-2 2
==1
(6 10-1) -12
/'
a(b 1)3
dadb + //' S/' (b- 1--1)((a - 1)²
dadb
dadb
2
J0
2)²
0n Jn
(a+ 2)3
J ba+b- 2)3
1
(1
In 2)
(1
In
(5 8 8ln))
1n
2
12
12
(6)
25

---

For the non-diagonal elements, it holds:
1
Ta TTb
Pa Pa, TT
= Jo J/ [="
(b (b-1)(a- 11)
dadb - (-] { a(b aabb-1()(())
1)² -
dadb + 1.' f (b- 1-))(a- 1
dadh
(a + b b- 22)
(a (a+1) b b- 2)3
(a + b - 2)
3
5
7
In
2 6 (1 2- + In2
In2
4
3
6'
Lemma B.2. a uniformly random stochastic 2 matrix, and π is the stationary distribution of P, then,
k
Pu.a (TT Puss) = - + 5og(2) ~ -0.034
2
and any n +
k
Th Pu,e (Te -u,s)
SE
Ty u,s (To (pm )wis)
Proof. We have:
h
E
Tee Pu. (T,, Pu,a)
u=1
1/6x
(x (6xy( -4x 2x + 1) + 6y² y² 24x)
12(x + y)
(8)
(1222 12r - 3) + log((x y) (40² +2m +2m
() + 1) 000 (4y² 2 12))
+
12(x +
--7/2 +2og((2
For the inequality, have intuition that doesn't 't depend on k notice that:
Σ
Ty Pu.e (To (P")) (Pm)).e)) 2
n Pu.s
To (Pa) u.8
8=1 u=1
2 1 u
8=1 1=
2 -a
As long as a isn't concentrated around 1, then this shows that the magnitude the RHS bounded by term that shrinks exponentially in n. For & = 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
(4.
P1,2 P2.1) (1 P1.2 P2.)
2 Pa a.s
IT& (Pm) )m.))
Teg
(P1,2 P21) )3
=
We show that for any choice of P1,2 and P2,1 the unit square,
P2.2 (4P1, P2,1 P1.2 P2.1)
$
(P1,2 + P2,1)
26

---

To see why this is true, observe that,
(4 P12 P2,1)²
-16P?2P21+00 + (P1,2 P2.) - 8(P12
11,2,)
- $16Pi,2 (P1, - + P2.1) 2 P1,2P21
since PL,2 + P2, S 2
-16Pf,222 -
P2,1) -2.))-4PP.PP2.11
(P1,2 P2,1)² -4 2
< (Pi.2 + 1
Using the above, we have
P1, (4P1. P1, - Pe.1))
P2 12 12
1
(P1. + P2,1)
(212)
(P1. + P2
P..212.111
using (P1,2 P2,1) 2 421
16P 1,2 P2,1
1
16
So.
1
T Pa.s
u.s T5s (Pm) (Pm)a.a)
uu.
11 1
P1.2 - P2.1
4
Now,
1
{" [
I
1 1 P1.2 - P2,11")
= - Jo Ju
|| x yl"
4
4
2
4 (n 1)(n + 2)
2(n + 1)(n +
Notice that this decreases in n, and 3, 13+1113+2] 1 0.025 which is less in magnitude than the value we proved at n = 1,| 7/2 + 5lgg x 0.034. So, solving for n = (verified by symbolic algebra program)
Pi. P2.1 (- P1. P2,1 + 1) (2P1,2 P1. (P2.1 - 1) P2.1 (P1. 1))
413
149 log (2)
+
0.002
(P1,2 P2,1)
60
15
Which 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.
Pi P2.1 (- P21 + 1) (2P1,P2.1 + P1, (P21 P2, (P1.2 -1))
E
(P1,2 + P2,1)
P1,2221 (2P11.11 + P1,2 (P2, 1) + P21 (P12 - 1))
=-B E
(P1,2 + P2.1) 3
55 (o)(())) ~ -0.0114 -7/6 -
Which is greater than -7/2 + 5 log 2, completing proof.
27

---

Lemma B.3. P is uniformly random doubly stochastic matrix, then,
k
E
Ra
2
Pa
Pa.
E
Tb
Σ To
k
==1
for all b and
E
Fu Pu, (Ta Pun))
<E
P
TTo
(1a) u.s))
For all non-negative a + and
Tu Pa Pu,s (T.
Pu.s)
<0
Proof. We will use the fact that for doubly stochastic matrices, the stationary distribution is the uniform vector I.
The first equality follows directly from Te Th. Now we will prove the first inequality.
1
1
E
T.P u.a (T. (Pa) (Pa)).s)
E
P
(pa)
yuss
k2
11,8
E
Puio
1:2
(PP
k2
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 = p2, where equality is reached), (P. Pa) < ||| First, by Cauchy-Schwarz,
P", < P|||||
We can strictly less than because Cauchy Schwarz is only tight when P and Pa are linearly dependent, and since both P and pa are doubly stochastic, linear dependence implies equality, which only the when P = P Then, For assume 0, then,
||P||||| P
=|P|| |PPa 1|
= ||P|| Σ A, ||F
S Σ CaiIPlrr/AA,P-1111 1|
Σ [a||P|ls|| Pa-11
||| |F
The 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 P E, a for some permutation matrices A, and constants a; > with Σ, a; = 1. The inequality step uses the triangle inequality. Induction on positive yields the desired inequality for positive
Now consider the remaining case, 0.
E Pu,s (Pº)
P.
u.e
1.2
28

---

0.
k2
1-2
While at a 1, for any that isn't
11
||
1, so
k
&
1
E
p2
0.
1
u.
122
k2
Completing the proof of both inequalities.
Lemma B.4. If P is a uniformly random k. X k stochastic matrix subject to each row being the same, then,
k
]
1
E
T
Pa
<
E
Ta
Σ
Pa
Pa
<0
k
k
sm
and
E
n
P Ps
-11 8
E
Tu
Tb
k
Σs=1 Pa
for all a and b and
2 Tue
Pu.s Mis
(Tx (Pa)) )u.s
=
For all a.
Proof. The equality statement follows from the facts that for such transition matrices, pa P for all natural a > 0, and that the stationary distribution matches the rows, that is, for any a, b, Tb = Pa.b
2 PuP u, (To (Pa) )w.s)
E
P u,, ( Pu.s)
T P u, (T, Ta)
Now we will do the inequalities. We will also use the following facts derived from the moments of the Dirichlet distribution.
k +
4(k: +
E
[|
(k +
1)(k + 2) 3)
So,
k
1
π
PT.
= E
T
k
k
s=1
s=1
||
113
E
[||
E [[II
4(k
k2 + 1)
k(k + 1)(k: + (k + +3)
Which is negative for all 2. And,
k
Ta Tb
T:
Pa
E
Ta Rb
k
s=1
29

---

1
]
E
|| 2
2²
E
[
1:3
1:2
k3
k2
(1c
+
Which is also negative for all k 2 2. Finally, notice that
4(K++5)
RI(++1)
krk +)++++))
8
5
1
(++++
For all 2.
B.6 Approximation Lemmata
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 - mins, Pt,. Then, for any i,
(P")ss Tjj S a
Lemma and
and Lemma B.7 both share similar intuitions and proofs. They largely rely on Lemma B.5 which shows that (P"),, )., approaches π 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 o, be states. Assume that min,, P.,j is positive and doesn't t dependend on T. Then,
k
T- Σ
Tb
k
(P))
(P')
ba
Σ
Pa.s T T i+1
(P')
8,0
(P
b,s
1=0
1=0
ToJa (T + 1)
Σ Pa
+ Oloog
k
Proof. Let bound the magnitude of the difference between the expressions.
(P))
Σ Pa.e
Pass
(Pi)
Σ
(P))
(T
1
Σ Pa
k
i+
3,a
h
1=0
7--
b
h: (PP)
b,a
Na
Pa.
T-i+1
T
(Pt)
8.0
Σ
(P ),,n na
1=
s=l
j=0
T-
1
T6
K.
PP
bia
Ta
+
P 0,8 T -
i+t
(Pi) s.a
(P)) - TT
1=0
s=1
j=0
T
T-
0 +
P 0,8 T
(pi) 5.a
(P') - Tottal
k
it1
1=0
#=1
j=0
T
T-i
+
P 0,8
((P)) b,s (Pt) a.a Ta) + TT (Ps) Ta)
T
i+1
1=0
1=1
1=
k
T-
Σ
(Pd)
b,s (Pi) - Ta + Ta (P') mal)
T6
k
a +
0.8
T i 1
1=0
1=1
1=0
30

---

(Pi) a² + +
0
Pa.s T -i+1
k
1=0
T-i
Pa
Σ
(a' ao)
k
T i+1
==1
1
1- 1-a --a--11)
P F,
a'
T i-i+1
1- a
k
m
Pa.g a'
T --++11-1
8=1
+
a- + --i+ll-a
k
1
Σ
a
k
T -it 1i+1-11
(1++1)
1
a T+1
log(T Og( +1++1++1
$ Tb
k
1 a
1
2+ 1 g
+ log(( + 1)
1-a
2log T
1 log T
mingj Pa.j
=O T)
The last step follows from our assumption, completing the proof.
Lemma B.7. Let be a stochastic matrix with all positive entries, and let a, be states. Assume that mins,j Pu. is positive and doesn't It depend on T. Then,
k
T
T
Pu.e
Tu
(pi) 8.u
(Pa) Ju,s
(P))
T
--it
i+1
T i
u,s
==a
1=0
k
(log (T + 1) - log (a + +))
2 P
To - (PP)).)))
+ +11)
Ty u,
smm caml
Proof.
First notice that,
k
k
ΣΣ
Pu,s] nh ((Pt)
Nz
nu
k
1
(pa)
a.s
TTa
Pa.s ((pa) ) T
h Σ k
k Σ a=1
(pa)
un5 ns)
Pu,s ((Pu)) - T3a)
4=
a=l
2
(1 - 1)
Σ
Purs ((Pa)) ) TTs)
k
u=l
s=1
2
P uu
:T (pa)
We will bound the distance between Σ
(1 Pa.s)
n3 (IDo) 3x
and
[aLT] Ezjp
u=1
8r
Define = 1 2 mini,y P., as in lemma
B.5
31

---

T
T-
1
Pa. u,
(P') 6,u
(P
k
T
T-
u.e
u=1
i=a
Pu,a Σ
2
1
(Pa)) M1,
NUT
T 1i11
T
k
P 0,5
-+1+1
(Pi) 2,u
(Pa) (Pm)uss
T i i+1
( u.e
-u
(Pe) TT.
T-
Tu
(Pi) 5,
(Pa)
u,
1
Σ
(pi).ss
(Pi)
-u (PPa) -Ta)
T
i+1
T 1it1
u.s
j=0
T-
(P)) 5,0
1
(P))
(PDa) - No) ((P)) TT)
Te
T-i+1
T i
0,6
j=0
T-
TR
(PM) 5,8 T i i+1
Σ
-(P)) (.2))
(PO -(po) To ((D)) u, TT
-
T -i+1
J=0
4 k T
T-i
Tu
1
Σ
axa (PS) (P)) u.o + ((pa) )u. Ta (Pi) 2.23 Tea
T
-11
9,0 T +
120
1=0
k
T
T-i
MM
(Pi) T - 1-it1
Σ
0 + (Pa)
u,s nz a²
By lemma B.5
T
-+1
5=0
TU
(Pi)) 8,8s T -- [+1 1-
T-it1
+ ((pa) Ta) a²)
T i+1
1
((Pa) )a, TTs ) a²'
+
--i+1
T-i+11 11
k i+1
k
Tu ( pa ) Ts a²
+
+1+11
1=0
s=1
T
h Σ
k
+
(Ta T a
i+1
T [
1=0
[=1
1=1
T
k
1
T- i+1
T-i+11
T
(T -i 1)²
1-a
2k
1-a
k
min.j Pa.
=0(1)
The last step follows from our assumption, and the fact that & does not depend on T.
32

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope?
Answer: [Yes]
Justification: All claims made our paper are supported by theoretical proof's or rigorous empirical evaluations.
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 other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals not attained by the paper.
2. Limitations
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 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 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-specifficttinn 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 few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, 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 provide closed captions for online lectures because 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 grounds for rejection, 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.
3. Theory Assumptions and Proofs
33

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Our main paper may have some informal theorem statements for ease of read- ability, we give formal versions of the statements with full proofs in the appendix. We also make sure to clarify the assumptions under which our results hold.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide all training details in Appendix and information on the data generating process and architecture in Section
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- reasonable avenue for reproducibility, which may depend on the
sions provide nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution should either be way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
new model (e.g., a large language model), then there
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilitty In the case of closed-source models, may be that access to the model is limited in
34

---

Justification: Our main results regard the multiphase nature of training, and we include Figure which shows that for the seeds 1,... the model each time has the same patterns in the training curve. The curves were shown individually instead of through error bars since the 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.
5. Open access 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]
Justificattin:
code is shared at https://github. http::///.tiuuu..o//Erreeee.... com/EzraEdelman/ Evolution-of of www.offtttttttaalllaaaaaa Induction Heads
Our
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
piblicc/guiles//Coeetiimssssss oot.isssysnooooyyyoloooooloo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurlPS code and data submission guidelines (https Coaeee//oeeSSeeeseeeeee
//nips. https/cccpppbc///eeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code permitted.
Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparame- ters, 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 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, 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]
main purpose of the loss curves is to show the shape, but because the phase transition doesn't '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.
35

---

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 capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show 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.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: The resources described in Appendix A Only single computer 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 compute required for each of the individual experimental rums well 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 Ethies
Question:
Does the research conducted in the paper conform, in every respect, with the NeurIPS
Code of Ethics /ttp//////ottttt//////////////// http///www..p..oo.o..........aaa
Answer: [Yes]
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 NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure preserve anonymity (c.g., if there is special considera- tion due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
36

---

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 No, they should explain why their work has no societal impact or why the paper does not address socictal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeen 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 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 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, not needed 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 being used 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 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).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, 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 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 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 URL.
37

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released. the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswi toode.. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged reach the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification:
We do not introduce any new assets.
Guidelines:
The answer means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
38

---

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 include any information that would break anonymity (if applicable), such as the institution conducting the review.
39


Reviews:
[{"review_id": "1YUvACaZ28", "reviewer": "Reviewer_oy83", "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.", "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)?\n2. 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": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper joins a line of works that aim at studying transformers with the help of synthetic data generated according to Markov distributions. Even if the model considered is significantly simplified, I believe that the theoretical and experimental insights provided by the paper are intriguing. The paper is also well written and the results seem correct to me, even if the notation can be improved, especially for the proofs in the appendix, that are not easy to follow.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful review. Here we address the questions/concerns raised by the reviewer.\n\n1. > 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)?\n\nWe believe that the analysis would mostly carry over to a more complicated architecture with one exception. Specifically, the softmax in the second layer attention proved intractable to analyze at the same time as the softmax in the first layer. Analyzing softmax is a common difficulty in transformer analyses, and why most theoretical results apply to linear attention only, we are happy to have managed to include the softmax in the first layer. Empirically, the training dynamics of the minimal model are very similar to that of the full model, the primary difference being that the phase transition isn't as sharp in the minimal model. With the softmax in the second layer added in, the phase transition becomes sharp like in the full model, but the tools do not currently exist to analyze this (specifically nested softmaxes on inputs not close to 0 or infinity). \n\nTo create our minimal model, we started from a two layer attention-only disentangled transformer (using relative positional embeddings) and iteratively simplified parts that empirically did not affect the training dynamics. In this context, disentangled transformer (introduced in [Elhage et al](https://transformer-circuits.pub/2021/framework/index.html)) is like a normal transformer, except the output of each layer is appended after the residual. In our construction, the first layer only attends to positional embeddings, and the second layer ignores positional embeddings, so we set the first layer key matrix, $W_k^{(1)}$, to $0$, and removed positional embeddings from the second layer ($v^{(2)}=0$). Then we set the value matrices and query matrices to the identity in both layers. In the experiments, all optimizations such as layer norms and weight decay were removed, and the optimizer used was SGD. With all of these changes, the overall training dynamics were not changed much at all, depending on hyper parameters they could speed up training, but the same loss curve and phases were observed. To make analysis of gradient descent on the minimal model feasible, the softmax on the second layer had to be removed, which did make the phase transition less sharp.\n\nWe believe an analysis that includes this softmax could confirm additional insights into why the phase transition is so sharp, and is a potential direction for future work.\n\n2. > 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?\n\nEmpirically, the absolute positional embeddings add noise, but do not fundamentally change any of the results. The analysis might be a bit messier, but we would not expect it to be fundamentally different. Additionally, many state of the art models use non absolute positional embeddings, including the Llama series of models ([1](https://arxiv.org/abs/2302.13971), [2](https://arxiv.org/abs/2307.09288), and [3](https://arxiv.org/abs/2407.21783)).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D2BYBfaxhX", "reviewer": "Reviewer_Kgpp", "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.", "questions": "What is the key differentiating contribution of the work in comparison with the myriad of other works in this space?\n\nWhat is the conclusion of the work that you believe is transferable to real-world ICL setting. What n-gram statistics would that follow?", "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.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Nice presentation, with clear figures exhibiting the key behavior in question (sudden \"emergence\" of the correct behavior.\n2. Good attempt to form a theoretical/mathematical foundation, which extends to the appendix.\n3. Attempt to understand ICL LLMs using a good toy task.", "weaknesses": "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.\n2. Questionable how well the toy setting can actually transfer to real-world ICL, albeit interesting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and questions. Here we address the main weaknesses and questions raised by you.\n\n## Impact and related works\n>Questionable impact since this topic has been explored heavily in the past years, with other similar tasks and toy models existing, showing similar results... What is the key differentiating contribution of the work in comparison with the myriad of other works in this space?\n\nAs we discuss in our related work section, there are indeed prior works which introduced synthetic settings for exploring ICL:\n- [Garg et al., 2022](https://arxiv.org/abs/2208.01066) (and follow-ups) train models to learn linear functions (and other simple function classes) in a few-shot setting.\n- [Xie et al. 2022](https://arxiv.org/pdf/2111.02080) describe a few-shot learning setting in which each document consists of examples generated by a hidden Markov model.\n- Finally, most related to our work is [Bietti et al. 2023](https://arxiv.org/abs/2306.00802), in which all sequences are generated from a single Markov chain (which doesn't need to be learned in-context), but certain 'trigger' tokens are automatically followed by sequence-specific tokens which need to be learned in-context.\n\nWe believe that our task, in-context learning of Markov chains (and the higher-order generalizations thereof), is a valuable contribution in the context of this rich literature. (No single synthetic task is going to capture all of the scientifically interesting aspects of in-context learning.) In particular:\n- It is very natural and simple to describe.\n- It elicits the formation of *induction heads* in networks trained on the task. (It is a particularly natural setting for studying their emergence.)\n- There are multiple strategies of various levels of sophistication which can be used to solve the task to varying degreees of success (unigram, bigram, etc.), and which networks pass through in stages during training.\n\nMoreover, the results we obtained by studying this task experimentally and theoretically are novel. Our key findings -- the formation of statistical induction heads, the stage-wise phase transitions between in-context solutions (unigram, bigram, trigram, ...) on the way to success, the arguments that simplicity bias may delay the formation of the correct in-context solution, and finding that the second layer of the model is learned before the first layer, rather than the other way around -- are all original to our work.\n\nGiven that you say there are lots of very similar works, **we suspect you may be mistakenly assuming that some concurrent works were actually prior works.** Indeed, at the same time we released our preprint, there were several simultaneous papers (all with pre-prints posted within a single month) that shared aspects of our setup and/or scientific focus. Please see the \"Concurrent Works\" paragraph in our Related Work section for a discussion of these papers and their relation to our work:\n [Akyürek et al., 2024](https://arxiv.org/abs/2211.15661), [Hoogland et al., 2024](https://arxiv.org/abs/2402.02364), [Makkuva et al., 2024a](https://arxiv.org/abs/2402.04161), and [Nichani et al., 2024](https://arxiv.org/abs/2402.14735). There have also been subsequent works in this space, including [Rajaraman et al., 2024a](https://arxiv.org/abs/2407.17686), [Rajaraman et al., 2024b](https://arxiv.org/abs/2404.08335) and [Makuva et al., 2024b](https://arxiv.org/abs/2406.03072).\n\nWe want to make sure that you are following the scientific best practice and judging our work's significance in relation to *prior works*, not concurrent works. If you mistakenly thought the above concurrent works were prior, we understand! There has indeed been a burst of activity on this topic recently. In that case, we hope you will reassess our work in its proper context, and adjust your score accordingly. If you were referring only to actual prior works, we hope you find the above contextualization of our work convincing and reassess accordingly; if not, we would appreciate some elaboration on the works you are referring to.\n\n\n## Real-world relevance\n> What is the conclusion of the work that you believe is transferable to real-world ICL setting. What n-gram statistics would that follow?\n\nN-gram statistics are useful for predicting real-world natural language text. Indeed, historically, language models were often simply n-gram predictors! It is also useful for models to learn in-context bigram (and more generally, n-gram) statistics. For instance, the writing style of a particular document is connected to the n-gram statistics of that document.\n\nThe most concrete connection the toy MC-ICL setting has to real world ICL settings is induction heads, which have been shown to have importance to in-context learning in LLMs in past work ([Olsson et al, 2022](https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html)). In the MC-ICL setting, the emergence of induction heads corresponds to the phase transition in loss, just like the phase transition found in the ICL ability of transformers hypothesized to be caused by induction heads emerging.\n\nAdditionally, compared to other ICL tasks studied in the past (such as linear of logistic regression) we believe that n-grams are a better model for natural language specifically.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iIT595PZOJ", "reviewer": "Reviewer_aVXq", "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.", "questions": "See above.", "limitations": "- Simplified transformer architecture \n- Synthetic task only, not clear if similar phenomena appear in larger models trained on natural language data.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper addresses a relevant topic within a simplified setting, facilitating the interpretability of transformer model solutions and an understanding of their training dynamics.\n1. Utilising Markov chains is an effective approach to studying sequence-to-sequence modelling.\n1. The study seeks to balance empirical investigation and theoretical analysis.\n1. The observation that models may progress from simple solutions, like unigrams, to more complex structures during training is interesting.\n1. It is relevant that transformers can learn an algorithm to estimate the transition matrix in context by gradient descent.", "weaknesses": "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.\n\n1. **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? \n\n1. **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: \n\t1. 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?\n\t1. 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?  \n\t1. 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? \n1. **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? \n1. **Proof of Proposition 2.2:** the proof of this proposition appears disorganised and not easy to follow. In particular: \n\t1. 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. \n\t2. 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$.  \n\t3. 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. \n\t4. 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?   \n\t5. 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.  \n1. **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 ? \n1. **Proof of Proposition 2.3:** The proof for the unigram construction seems to be missing. \n1. **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? \n1. **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? \n1. **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? \n1. **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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the in-depth review and feedback!\n\nYou point out that there is some misalignment between our experimental story and our proofs, and you see that as our paper's major weakness. We agree that there is misalignment, due to the difficulty of analyzing SGD on multi-layer Transformers. However, we believe our analysis is very relevant to the experimental observations about simplicity bias and phase transitions. In particular, Lemma 3.1 exhibits simplicity bias by showing that the parameters after the first phase compute a solution that is a weighted version of the unigram solution, and the second phase computes a mixture of unigram and bigram. Also, the magnitude of the signal for the second phase is a factor $T$ lower than the first step, indicating that the second phase requires more time for \"signal accumulation\". We agree that this does not explain the sharpness of the phase transition, but it does provide theoretical intuition for the existence of multiple phases. We are happy to add these aspects of the story to the discussion following Lemma 3.1 in the text; for these reasons, we don't think the extent of misalignment is a serious weakness in the paper.\n\nWe address your various minor concerns below:\n\n1\\. Thank you for pointing out this line as being out of place. A version of the line (edited for clarity) will be moved to our discussion of limitations. The line intended to communicate the difficulties with understanding the stationary distribution when $k>2$.\n\n2.1 and 2.2. The goal of the minimal model is to capture the dynamics of the full transformer, while being amenable to theoretical analysis. Essentially all analyses of multilayer transformers have relied on simplifying out many parts of the transformer. While simplifying the transformer we carefully checked that the dynamics we wished to understand were preserved as much as possible. Due to space limitations, we have deferred a more in depth explanation of the choices in the minimal to a comment.\n\n2.3. This is correct. Unfortunately, the analysis does not seem to be feasible with softmax in both layers.\n\n3 and 9. Generating each row of the transition matrix iid from a flat Dirichlet distribution is equivalent to uniform distribution over all transition matrices. We chose to focus primarily on this distribution because it is the most natural.\n\nThe only mention of doubly stochastic matrices in the theory (lemma C.3) was mistakenly left in after the surrounding context was removed.\n\nDue to space constraints, we defer definitions of the other distributions to a comment.\n\n4\\. Thanks for your careful proofreading. Due to space constraints we defer the specific changes and improvements made in response to a comment.\n\n5\\. The trained transformers used an internal dimension of 16. The specific parameters after training always vary randomly due to initialization and training samples. Figure 8 shows that the transformers are implementing the same algorithm as our construction, with each layer performing the same functions.\n\n6\\. Thanks for pointing out this oversight, we have added the full proof in (it's rather short and simple).\n\n7\\. That's an interesting question! Unfortunately, the details of the loss landscape for even simplified transformers are hard to precisely characterize with current theoretical tools. Among other difficulties, a barrier to transformer optimization results is analyzing cases where the input to softmax isn't close to zero, or one hot.\n\nOur intuition is that the unigram is not exactly a saddle point: during the plateau period, there is very slow continual improvement in the loss. We suspect the gradient at the unigram solution is simply small relative to the change needed to move towards the bigram solution.\n\n8\\. Line 241 is a holdover from a previous version, and no longer applies to the experiments in the paper, we apologize that it made into this version and will remove it.\n\nOn the left in figure 3, both the minimal model and transformer model reach their minimum loss after seeing around 80,000 training sequences. We do believe though that one shouldn't put too much weight into the fact that these took the same time to train.\n\nWe believe this comparison still makes sense. The minimal model is designed to be able to express everything the transformer expresses when being trained on the task. Experimentally, it goes through the same phases. In the right most graphs in figure 3, both models quickly learn to use the unigram strategy, before eventually adopting the bigram strategy. Since this task does not require the full expressive power of a transformer, it is not too surprising that a model with many fewer parameters is sufficient to gain insight into the training dynamics.\n\n10\\. The most compelling piece of evidence that the minimal model actually goes through these phases are the experimental results shown in figures 3 and 4. Experimentally this also works for standard small Gaussian initializations. In the conclusion, we will add a mention of the limitation of this being a two step analysis. Unfortunately, existing tools are insufficient to analyze dynamics involving softmaxes (specifically when the inputs to the softmax are neither small nor dominated by a single index). For lemma 3.1 we were forced to limit our analysis to two steps. However, as we mention on lines 579-580, the uniform component of $W$ does not contribute to the gradient of $v$, so even if we initialized $W$ to any uniform non-zero initialization, $v$ would still not change in the first step. We chose the specific step sizes to allow the theory to show how the model can go through the unigram and bigram phases, which happens in practice over the course of many small steps.\n***\nWe hope we have managed to address your concerns. If you think we have adequately done so, we would be grateful if you adjusted your score accordingly. Thank you!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DBUIumHe53", "reviewer": "Reviewer_1EuX", "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.\n\nGiven 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.\n\nThe authors further connect their findings with previous work.", "questions": "1. Why did you consider Markov chains with two states only?\n2. 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”?\n3. Why did you choose the stationary distribution as the initial distribution?\n4. 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?\n5. 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?\n6. Is the difference in KL between the unigram and bigram strategy due to finite sampling?\n7. Why does the bigram strategy have a smaller loss? is it because of the inductive bias of the model?\n8. 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)? \n9. Is the 2-layer transformer model initialised to the matrices in Appendix C.1? \n10. 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": "Yes, they authors did address (some of) the limitations of their method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Deep learning architectures are remarkably complex systems and designing both tractable models and modelling tasks that mimic the behaviour of these systems, prior, during or after the training process, is an important approach to unveil their inner workings. This paper attempts to carry one such analysis, by introducing a simple “in-context learning” Markov chain modelling task and studying how a two-layer attention-only transformer model and a simplification thereof solve it, thereby adding another interesting contribution to our understanding of transformer models.\n\nOne strength of the paper is that the authors first demonstrate that both their models can indeed find the optimal solution of their Markov chain task, and then empirically verify this, which gives soundness to their claims.  In particular the results in Figure 3, which demonstrate the similarity in behaviour of both models, are very compelling (see however the questions below). A second strength is that the authors identify different phenomena within their setting which have also been observed before, like multiple distinct stages of training, the latter of which is also connected to induction head formation, or the order in which the layers in the network are learned. Finally, the authors demonstrate that some of their findings are also present in a second-order Markov chain modelling task.\n\nPutting aside some minor details with some notation and content organisation (see below), the paper is overall well written.", "weaknesses": "My first issue with the paper is: \n\n- 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.  \n\nAs 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. \n\nA second major issue is that the authors do not explain their reasoning behind nor the limitations of their dataset/task definition. \n\n- 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. \n- 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. \n\nDespite its merits, I think the issues above require some revision, or more detailed explanation, before the paper can be published. \n\n*Other comments*: \n\nThe 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.\n\n*References:* \n- An explanation of in-context learning as implicit Bayesian inference. Xie et al. (2022)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for critically reading our work and helping us improve its presentation. We greatly value your in-depth review!\n\nWe will first discuss your most significant concerns with our work:\n\n## In-Context Learning\n\nYou're right that we don't spend much time in our paper discussing what \"in-context learning\" means, and indeed, there isn't a single consensus definition! The term \"in-context learning\" was (we believe) introduced in the GPT-3 paper ([Henighan et al., 2020]) to refer specifically to the phenomenon where LLMs have improved few-shot performance with increasingly many examples in their context. Xie et al., meanwhile, does indeed explore scenarios where there is a distribution shift between training and inference; however, in the broader literature, ICL is not consistently tied to distribution shifts (see, e.g. [Garg et al., 2022] for an influential study of ICL that does not require a distribution shift in its definition). Meanwhile, [Elhage et al., 2021], which was roughly concurrent with Xie et al., and the follow-up work [Olsson et al, 2022] -- use ICL very broadly to refer to LLMs being better at predicting tokens that appear later in a sequence. The way we use the term is somewhat narrower than this; as we state in the Introduction: we think of ICL as when language models are \"incorporating patterns from their context into their predictions.\" In other words, it is when models *learn* from patterns in their context; few-shot learning being a special case of this.\n\nThe task we focus on, ICL-MC, is an in-context version of a classic learning task, Markov chain inference. Each test-time sequence has *zero* probability of appearing during training, because of the continuous nature of the distribution over chains. During pre-training, the transformer learns to perform the learning task in-context. We don't only focus on the case where there is a distribution shift between train and test -- as discussed above, we think ICL is interesting even in the absence of distribution shift. (Though we do explore distribution shifts in the paper: see Figures 4 and 11.)\n\nWe would be happy to elaborate on our interpretation of ICL in the camera-ready version of the paper.\n\n## Answers to questions\n1\\. This choice was made to facilitate a theoretical analysis of the observed phenomena. In particular, results such as Lemma C.2 become feasible only when the number of states is 2, or other simplifying assumptions are made. However, in the experiments, we have verified that our insights transfer to Markov Chains with a larger number of states (see, for example, Figure 7 in the Appendix). We will add a discussion of Markov chains with more states, including a reference to figure 7.\n    \nBased on your concerns, we have decided to add back in a proof we decided not to include, which means that **by making only small changes to the submitted version, we are able to extend our main results to any number of states**. This requires using a distribution like that in experiment on the left of figure 4. Due to lack of space, we defer the details of this to a comment.\n\n2\\. Thank you for pointing out this line as being out of place. The line (edited for clarity) will be moved to the discussion of limitations. It intended to communicate the difficulties with the distribution of the stationary distribution for $k>2$. Due to lack of space, we defer further elaboration to a comment.\n\n3\\. We chose this to make the setting closer to the next token prediction seen in practice. The marginal distributions of the tokens generated from a markov chain approach the stationary distribution exponentially fast no matter the initial distribution.\n\nWe have ran experiments for other somewhat natural starting distributions (such as uniform) and found no meaningful differences.\n\n4\\. Yes. We will add the following mathematical definition: the unigram strategy's probability of the token at position $t+1$ being $j$ is $\\frac{1}{T}\\sum_{i=0}^T \\mathbb{1}[x_i=j].$\n\n5\\. The bigram strategy is computed by counting the frequency of pairs of states in the context of each sequence. This corresponds to the formula in line 157.\n    \nThe graphs on the left in figure 3 show the KL-div loss using the instance-dependent ground-truth transition probability matrix. Because the context size is long (100 tokens), and we average over 1024 test samples, these two measures are almost identical (as we would expect).\n    \n6\\. No, the KL div between the unigram and bigram strategies is not due to finite sampling. The unigram strategy is a flawed strategy (it ignores the information we get from the relative order of the tokens), for almost all markov chains it will get higher expected loss than the bigram strategy, even for small context sizes.\n    \n7\\. The bigram strategy is a near optimal strategy (it approaches the optimal strategy exponentially fast as the context length grows). This is true regardless of the model used.\n\n8\\. We are not entirely sure how to interpret this question and would appreciate clarification.\n\nSampling the first token by the stationary or uniform distribution does not meaningfully affect any empirical observations. The marginal distributions of the tokens after the first will always approach the stationary distribution exponentially fast. \n\nEmpirically, the model starts by representing the unigram strategy and later the bigram strategy. In both cases it is not obvious that the model needs to fit the unigram strategy at any point (after all, the unigram strategy performs worse than the bigram strategy), and so it must be  only due to the inductive biases of the model and optimizer that this happens.\n\n9\\. No. The transformers were initialized with the default Pytorch initializations (Gaussian with mean 0), we will clarify this in section B.\n    \n10\\. This is a typo.\n***\nWe hope we have managed to address your concerns and would be grateful if you adjusted your score accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)