Large Language Models as Markov Chains

Authors: Oussama Zekri, Ambroise Odonnat, Abdelhakim Benechehab, Linus Bleistein, Nicolas Boulle, Ievgen Redko
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, markov chain, in-context learning, generalization bounds, convergence analysis

Abstract:
Large language models (LLMs) have proven to be remarkably efficient, both across a wide range of natural language processing tasks and well beyond them. However, a comprehensive theoretical analysis of the origins of their impressive performance remains elusive. In this paper, we approach this challenging task by drawing an equivalence between generic auto-regressive language models with vocabulary of size $T$ and context window of size $K$ and Markov chains defined on a finite state space of size $\mathcal{O}(T^K)$. We derive several surprising findings related to the existence of a stationary distribution of Markov chains that capture the inference power of LLMs, their speed of convergence to it, and the influence of the temperature on the latter. We then prove pre-training and in-context generalization bounds and show how the drawn equivalence allows us to enrich their interpretation. Finally, we illustrate our theoretical guarantees with experiments on several recent LLMs to highlight how they capture the behavior observed in practice.

Decision: reject
Reviewer scores: 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1730659717014

thread_type:
official_review

<thread_object>
creator: reviewer_gi2n

<content>
summary:
This paper theoretically characterizes the inference mechanism of large language models (LLMs) as a finite-state Markov chain, providing insights into LLM properties based on the transition kernel of the associated Markov chain. Specifically, by modeling an LLM as a Markov chain, the authors show that the output distribution of the LLM converges to a specific stationary distribution and that the convergence rate depends on certain parameters in the LLM. They also derive generalization bounds for the pre-training and in-context learning of LLMs and experimentally verify the behavior of some recent LLMs using these theoretical results.

strengths:
* This paper introduces a unique approach by modeling LLMs as finite-state Markov chains. This approach provides a clear explanation of the inference process and convergence properties of LLMs to a stationary distribution and also provides a new perspective on understanding LLM behavior theoretically. This approach is expected to expand our understanding of the overall model structure of existing LLM schemes, such as autoregressive models with self-attention mechanisms.
* A major contribution is the derivation of generalization bounds for pre-training and in-context learning, applying the concentration inequality from Markov chain theory to gain a quantitative understanding of LLM's learning characteristics. This increases confidence in the accuracy and inference capabilities of LLMs and is also helpful for comparing different LLMs.

weaknesses:
* As noted above, while modeling LLMs as a Markov chain provides a new perspective, there is no new development in Markov chain theory itself. The author's theoretical framework is limited to an extension of existing Markov chain theory, and it remains unclear to what extent this framework fully captures the complex structure and dynamic properties unique to LLMs. Due to the multilayered structure and self-attention mechanisms within LLMs, it is unclear to what extent these can be accurately presented using Markov chains.
* The effect of temperature parameters on convergence rate is analyzed, but the influence of other hyperparameters and model structure, such as depth, number of heads, and layer-specific settings, on convergence and inference performance, is not examined in detail. Consequently, the paper provides insufficient guidance for optimal parameter design in LLMs.

questions:
* I believe that Fig.3(d) on line 243 might be a typo for Fig.3(c). Setting that aside, what is the meaning of Fig.3(c)? If the parameter $K$ is only assigned to Llama and GPT, there seems to be no additional information added to Proposition 3.3. Could the authors explain how to evaluate the specific value of $\epsilon$ in LLMs?

* Fig.4(d) presents the temperature dependence of $\epsilon$, but how was this obtained? I may have missed something, but I would appreciate clarification from the authors. It is obvious that convergence is faster at high temperatures, so I am curious about the functional form of $\epsilon(T)$. 

* The probability distribution formula on line 327 seems redundant.

* In evaluating the generalization bound, it is assumed that the pre-training data is generated from the Marton coupling. Is this condition essential for obtaining the bound?

* In Fig.5(right), the context length $N_\text{icl}$ dependence of the risk is presented for several values of $t_\text{min}$. While the results for small $N_\text{icl}$ deviate from the scaling law, the theoretical prediction is $\sqrt{t_\text{min}/N_\text{icl}}$. Did the authors attempt to plot the results as a function of $t_\text{min}/N_\text{icl}$ or $N_\text{icl}/t_\text{min}$ to check for a universal curve, indicating a crossover?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730728689591

thread_type:
official_review

<thread_object>
creator: reviewer_ugpy

<content>
summary:
In this paper, the authors first enstablish the equivalence between autoregressive models and first-order Markov chains with number of states proportional to $T^K$, where $T$ is the alphabet size and $K$ is the context window, proving results about the stationary distribution of such chains. In the second part of the paper, the authors prove pre-training and in-context generalization bounds for LLMs, validating them with several experiments.

strengths:
The paper is well written. I was not able to check all the proofs but the math seems sound to me and the results are interesting.

weaknesses:
My main concern is that it is not clear to me what is the connection between the first part (equivalence between autoregressive models and Markov chains) and the second part (generalization bounds) of the paper. As it is, the paper looks like two different papers glued together.

questions:
I have a few questions that would make it easier for me to understand the results of the paper and its impact:
1) What is the connection between the first and second part of the paper? Are the results of the first part important for the derivation of the generalization bounds? If so, what are the key steps when the equivalence proved in the first part is useful in the second part? What is the intuition?
2) I struggle to understand the implication of the results of the first part for LLMs. In particular, what does it mean for LLMs to converge to their stationary distribution? Does that mean that the model's loss has converged to some minima? Does that mean that the LLM inference performance cannot be improved further?
3) I was curious about the choice of the TV distance in the risk definition (2). Would it be easy to move to some more popular loss choice for LLMs such as the cross-entropy loss? Intuitively, it shouldn't be too hard since cross-entropy loss is basically KL divergence, which in turn is connected to TV distance through Pinsker's inequality.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730837700893

thread_type:
official_review

<thread_object>
creator: reviewer_o34n

<content>
summary:
This work studies the learning theoretical properties of auto-regressive model ( which the author claim to be LLM ) on Markov Chains.
In particular, they prove the convergence of MC and the generalization guarantees of learning the parameter $\Theta$. They also carried out numerical experiments on it.

strengths:
The strength is the theoretical results look pretty solid. And the author writes this paper in a very clear fashion. making it very readable. Also the results follow from the uniform concentration bound of MCs, presented in a rather classical learning theoretic (a.k.a. VC style) fashion.  I believe this result is relevant to the learning theory community although I cannot decide how much technical innovations are made here. The reviewer believes that this result adapt previous results by many other learning theory researchers who develop theory on estimating the transition kernel of MC given single path observations.

weaknesses:
The major weakness is the connection to LLMs. The first thing is whether the class F is sufficiently large to include some instances of LLMs. Right now it seems only takes a single input $v_{I}$ instead of $v_{I-1},v_{I-2},\ldots,v_{I-k}$. Though it may not be feasible anymore to learn MC due to the Markovian properties after we make this change.  However, real world language might not be treated as Markovian, which is a very stylized assumption here. The reviewer believes this will affect the impact of this work in the LLM community.

Another huge concern is that $F$ is regardless of the structure of the model itself. It can either be an RNN, Transformer, or a LSTM, and the result continues to hold. It can even be non-NN model. This does not actually reveal why LLM is a powerful machine since it highly depends on the NN as its backbone. Therefore these limitations significantly affect its contribution to the modern LLM theory community.

Hence, the reviewer suggests that instead of connecting such work to LLM, the author might be able to use the Auto-regressive model as the title, which is deemed more proper. And this model might include LLM as a special instance. At the current stage, the reviewer does not think LLM community will put a lot of focus on this work due to its over-general problem setup.

questions:
Overall this work is very well written, which is the reason why the reviewer gives a 5 rather than a 3. The reviewer believes that this work is not very relevant to the LLM community but more relevant to the traditional learning theory (a.k.a. VC theory) community (given that the author changes the way this work is presented by replacing most of the LLMs to autoregressive models and maybe draw connections in a milder way).
</content>
</thread_object>

</thread>
<thread>
creation date:
1731688643198

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comment

comment:
We thank all the reviewers for thoroughly and carefully reading our paper. We are happy to hear that they acknowledged the **novelty and quality** of our study (Reviewers o34N, ugPy, Gi2N), its **solid and interesting theoretical foundations** (Reviewers o34N, ugPy, Gi2N), and the **clarity of the writing** (Reviewers o34N, ugPy, Gi2N).

### **General comment**
Before giving individual replies to all reviewers, we would like to provide three important clarifications to clear out any potential misunderstanding.

- **1)** In our work, we employ the term LLM to define a deep transformer-based model trained on non-iid data whose inference is based on the next-token prediction principle. The latter implies that such a model transitions from **a sequence of tokens** to **a sequence of tokens**. Hence, in Section 3, the Markov chain formalization transitions between states that are **sequences of tokens** and **not single tokens**. This defines a Markov chain on a state space of size $|\mathcal{V}_K^*| = T(T^{K} - 1)/(T-1)$. The vast majority of existing LLMs fall into our definition suggesting that our results apply to them. We also note that this **does not rely** on any Markovian assumption on the language. 

We propose to add this definition in Section 1 and avoid using the term auto-regressive models which raised some confusion.

-  **2)** We would like to highlight that our **experimental results**, not acknowledged in the reviews, **validate our theoretical findings** using Mistral 7B, Gemma 2B, Llama 3.2, 2 versions of Llama 2, which are all very recent models and commonly used in the LLM community. We believe that this shows the high relevance of this work for the broad LLM community. We would be more than happy to answer any questions on the experimental section. 

-  **3)** All our theoretical results are obtained by considering **a generic transformer-based LLM** (multi-layer, multi-head Transformers typically used in GPTs, Llamas, etc. [1]) **without any simplifying assumption on the architecture** which is rarely done in the literature (e.g., most theoretical work focus on linear attention or simplified one-layer transformers [2, 3 ,4]). In this setup, we consider a challenging setting with **non-IID data** which was also scarcely explored. Although our work does not provide a definite and complete answer, we believe it makes a significant contribution to this field of study. We made these choices to remain as close as possible to the practice of the LLM community (see point 2). This comes with many technical challenges and to derive our proofs, we **extensively rely on the characteristics of the transformer architecture**. We believe that in addition to our experiments on commonly used LLMs, our results and the techniques used can be valuable both to the LLM and the Machine Learning communities.

We hope that the reviewers will take these clarifications into account in their evaluation of our work. We remain open to continuing this constructive discussion for the length of the rebuttal period.


We strongly believe that the paper benefited from the reviews, and we respectfully invite the reviewers to reconsider their evaluation score if they think we addressed their concerns.

### **Updates on the revised paper**

We list below the changes made to the paper (in blue) as suggested by the reviewers:

- Additional definition of LLM following point 1 (l 112-117).
- Explicit formulation of the size of the state space $\mathcal{V}_K^*$ of the Markov chain in Proposition 3.1 (l 177).
- Risk definition, with a more direct link to the Markov chain equivalence. (l311 & l314).
- Updated Figure 5 (p.10) with the evolution of the risk as a function of $\frac{N_{icl}}{t_{min}}$
- Additional results with the KL divergence in the risk instead of the TV distance (Appendix F).
- Replaced the term “auto-regressive models”, which has given rise to some confusion.

[1] Brown et al. Language models are few-shot learners. NeurIPS 2020

[2] Furuya et al. Transformers are Universal In-context Learners. arXiv 2024

[3] Zhang et al. Trained Transformers Learn Linear Models In-Context. JMLR 2024

[4] Oswald et al. Transformers Learn In-Context by Gradient Descent. ICML 2023

[5] Kim et al. Transformers are Minimax Optimal Nonparametric In-Context Learners. ICML 2024 Workshop
</content>
</thread_object>

</thread>
<thread>
creation date:
1733226900927

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Wrap-up of the rebuttal

comment:
As the rebuttal period approaches its end, we would like to thank the reviewers for their contribution to improving our manuscript. We are pleased with the fact that **two reviewers (ugPy and Gi2N) chose to raise their scores**. To summarize, we would like to recapitulate the following
1. No reviewer questioned the novelty of our contributions, which, together with the importance of the studied topic, experiments on recent up-to-date LLMs and lack of technical flaws leaves **no particular unresolved issues**. As it appears now, the scores for our paper are rather borderline but we hope that the reviewers will take into account what is written above to endorse our paper stronger.
2. We would like to reemphasize that our answers to reviewer o34N provide a factual justification for our claims. We hope that the reviewer will acknowledge this post-rebuttal as they haven’t engaged in any further discussion so far.
3. The reviewer o34N’s motivation to decrease their score seems out-of-place as we do not see any negative reaction from reviewers’ ugPy and Gi2N justifying lowering their score. In fact, both of them raised their scores as mentioned above, and highlighted the fact our thorough answers resolved their concerns.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper brings a learning theory viewpoint to the standard equivalence between language models (LLMs or others, of window size K) and Markov chains (of state size |V|^K). It later connects the existence of a stationary distribution of Markov chains to the the inference power of LLMs.

Main concerns after author feedback remained around whether the theoretical framework truly captures transformer-specific properties of LLMs. The work suffers from over-claiming in this aspect, even affecting the title of the paper. The main equivalence in Proposition 3.1 (and results of Section 3) have nothing to do with transformers but holds for any neural network which (in an autoregressive fashion) predicts the next token based on previous K tokens.

As suggested by Reviewer ugPy, it would be better to decouple this into two papers, separating the (in my view not novel results of Section 3) and the more interesting results of Section 4, or at least deemphasize Section 3 as being only a warm-up for the main results in Section 4.
The author feedback phase led to constructive discussion and improvements in clarity, but nevertheless couldn't fully resolve the main concerns.