Summary
This work introduces a simple “in-context learning” Markov chain modelling task and studies how transformer models tackle it. The task consists of inferring the transition probability matrix of the Markov chain, which is sampled from a prior Dirichlet distribution. The authors consider the task to pertain to “in-context learning” because each data point (that is, each sequence in the dataset) is sampled from a different transition probability matrix. The prior over the matrices is however global and thus shared by both training and test datasets. More in detail, the authors focus on Markov chains with only two states, and sample each sequence from the stationary distribution of the process.
Given this dataset, the authors investigate the training dynamics of a two-layer attention-only transformer model both empirically and theoretically. The theoretical investigation involves the introduction of a simplified model, which is defined to mimic the main features of the transformer model. The authors report three different training stages, during which the model very quickly changes from predicting random transitions, to predicting samples similar to the stationary distribution of the chain, and lastly, predicting transition from the inferred transition probability matrix.
The authors further connect their findings with previous work.
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.
One 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.
Putting 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:
- how much does the proposed learning task really concern “in-context learning”? especially in the context of large language models, which the authors use as motivation.
As framed by e.g. Xie et al. (2022) who focus on LLMs, "in-context learning" deals with sequences which have low probability under the training distribution. The authors do not really elaborate in their understanding or interpretation of "in-context learning", nor do they explain how it relates to "in-context learning" in the setting of language models. In short, I believe the paper contributes more to our understanding of transformers than to our understanding of in-context learning.
A second major issue is that the authors do not explain their reasoning behind nor the limitations of their dataset/task definition.
- First, the authors consider Markov chains with only two states and do not comment on why they restricted their study to that case nor how their findings extend to Markov chains with more states, if at all.
- Second, the authors choose the initial distribution of the chain to be its stationary distribution but, again, did not explain why. Note that a stationary Markov chain corresponds to a very special case, and it's not very clear how the observations made by the authors extend beyond it. Indeed, choosing a Markov chain with more states, initialised from a distribution far from stationarity, can generate long sequences which are “out-of-equilibrium”. One can't but wonder whether one would still find multiple training stages in this case. See e.g. question 8 below.
Despite its merits, I think the issues above require some revision, or more detailed explanation, before the paper can be published.
*Other comments*:
The comment in line 205 “We observe that training a 1-layer transformer fails to undergo a phase transition or converge to the right solution”, could be better highlighted as evidence for the emergence of induction heads. It feels somewhat buried in between other comments and statements. Also there’s a typo in line 239 and another in line 280.
*References:*
- An explanation of in-context learning as implicit Bayesian inference. Xie et al. (2022)
Questions
1. Why did you consider Markov chains with two states only?
2. Can you please elaborate on your comment in line 145, page 4: “the stationary distribution of this random Markov chain does not admit a simple analytical characterization when there is a finite number of state”?
3. Why did you choose the stationary distribution as the initial distribution?
4. Is the ground-truth unigram strategy then computed by estimating the stationary distribution with the count histogram? Is this the one used to compute the KL distance in Figure 3?
5. Similarly, do you use the bigram strategy of section 2.1 to compute the corresponding KL distance in Figure 3? Or do you use the instance-dependent ground-truth transition probability matrix instead?
6. Is the difference in KL between the unigram and bigram strategy due to finite sampling?
7. Why does the bigram strategy have a smaller loss? is it because of the inductive bias of the model?
8. In the setting with an initial condition far from the stationary distribution, would the model first fit the step-dependent marginal distribution of the chain (a unigram strategy) and later fit the bigram strategy of section 2.1? Isn't it simpler to fit the latter first (as opposite to the stationary case)?
9. Is the 2-layer transformer model initialised to the matrices in Appendix C.1?
10. The sequence length is first labelled by $t$ (e.g. after eq. 1) and then by $T$ (e.g. before eq. 4). Is this change of notation intentional? Or am I misunderstanding something?
Limitations
Yes, they authors did address (some of) the limitations of their method.