Transformers Learn Variable-order Markov Chains in-Context

Authors: Ruida Zhou, Chao Tian, Suhas Diggavi
Primary area: foundation or frontier models, including LLMs
Keywords: In-context learning; Variable-order Markov chain; Context Tree Weighting

Abstract:
Large language models (LLMs) have demonstrated impressive in-context learning (ICL) capability. However, it is still unclear how the underlying transformers accomplish it, especially in more complex scenarios. Toward this goal, several recent works studied how transformers learn fixed-order Markov chains (FOMC) in context, yet natural languages are more suitably modeled by variable-order Markov chains (VOMC), i.e., context trees (CTs). In this work, we study the ICL of VOMC by viewing language modeling as a form of data compression and focusing on small alphabets and low-order VOMCs. This perspective allows us to leverage mature compression algorithms, such as context-tree weighting (CTW) and prediction by partial matching (PPM) algorithms as baselines, the former of which is Bayesian optimal for a class of priors that we refer to as the CTW priors. We empirically observe a few phenomena: 1) Transformers can indeed learn to compress VOMC in-context, while PPM suffers significantly; 2) The performance of transformers is not very sensitive to the number of layers, and even a two-layer transformer can learn in-context quite well; and 3) Transformers trained and tested on non-CTW priors can significantly outperform the CTW algorithm. To explain these phenomena, we analyze the attention map of the transformers and extract two mechanisms, on which we provide two transformer constructions: 1) A construction with $D+2$ layers that can mimic the CTW algorithm accurately for CTs of maximum order $D$, 2) A 2-layer transformer that utilizes the feed-forward network for probability blending. These constructions can explain most of the phenomena mentioned above. One distinction from the FOMC setting is that a counting mechanism appears to play an important role. We implement these synthetic transformer layers and show that such hybrid transformers can match the ICL performance of transformers, and more interestingly, some of them can perform even better despite the much-reduced parameter sets.

Decision: reject
Reviewer scores: 3, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1730691335768

thread_type:
official_review

<thread_object>
creator: reviewer_p9pl

<content>
summary:
This article examines how transformers learn variable-order MC. It extends previous research on the learning of fixed-order MC.

strengths:
1. The authors found that transformers can learn to compress VOMC in-context, approaching the optimal CTW algorithm for appropriate CTW-priors.
2. The authors found that transformers do not require high network complexity. Even 2-layer transformers can perform well.

weaknesses:
1. Research on fixed-order MC helps us understand how the in-context learning capabilities of LLMs emerge. This is an important finding. Building on this, what key discoveries does this paper offer regarding variable-order MC? The authors' findings, including (1) transformers can learn VOMC in-context capabilities, and (2) transformers are more powerful than CTW, do not seem to bring significant new contributions or explanations.
2. To my understanding, the authors investigated the expressive capabilities of transformers. However, this has already been established. Previous work has found that even transformers with only one encoder layer and three decoder layers are Turing complete. Considering this, what stronger discoveries does the author's research on expressive capabilities bring compared to Turing completeness?
3. Why compare with CTW? Is it a sufficiently powerful framework?
4. What practical insights can be gained from the study of VOMC?

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712210616

thread_type:
official_review

<thread_object>
creator: reviewer_hvxk

<content>
summary:
The paper explores the in-context learning (ICL) capabilities of large language models (LLMs) by introducing a novel perspective of viewing language modeling as data compression. The study finds that transformers can effectively learn to compress VOMC in-context, outperforming traditional compression algorithms under certain conditions. Key phenomena observed include the resilience of transformers to layer variations and their ability to surpass CTW algorithms, particularly when non-CTW priors are used. The paper also provides theoretical frameworks to explain these observations.

strengths:
- The paper presents an innovative perspective on understanding ICL by framing it as a context information compression task, which broadens the theoretical understanding of transformer capabilities.
- It offers empirical evidence and theoretical constructions that explain the observed phenomena (like Fig 6&7), providing insights into how transformers can mimic and even surpass traditional compression methods.

weaknesses:
- The work attempts to interpret ICL, but for readers familiar with ICL/LLM, the introduction of transformer attention (Sec 2.1) is somewhat redundant. Conversely, there is insufficient background information on CTW and related work, which are crucial for understanding the paper's context.
- The paper is challenging to follow without consulting the appendix, like Appendix A/B, which is essential for understanding the introduction and foundational concepts.
- There is a lack of real-world ICL example analyses, especially considering that the data used in the paper is different from the real data used in LLM ICL. Incorporating such examples, especially if the real ICL attention pattern mirrors those in Figures 6 and 7, would provide more robust validation.

questions:
- How is the compression rate computed, and why can this rate exceed 1 for the PPM algorithm, as shown in Figure 4?
- How much of the training data in Section 3.1 is constructed, is the data sufficient enough for training, and what are the details of the training process, like the learning rate?
- What does the term "segment index" refer to in Figures 4 and 5?
- Typo: Line 356: "Fig 7" should be corrected to "Fig 6"
</content>
</thread_object>

</thread>
<thread>
creation date:
1731252121344

thread_type:
official_review

<thread_object>
creator: reviewer_oxgr

<content>
summary:
This paper explores the in-context learning (ICL) capabilities of transformers for variable-order Markov chains (VOMC), viewing language modeling as a data compression problem. The authors benchmark transformer models against compression algorithms like context-tree weighting (CTW) and prediction by partial matching (PPM). Key findings include: (1) transformers effectively learn VOMC in context, outperforming PPM; (2) ICL performance is relatively insensitive to layer count, with even a two-layer transformer performing well; and (3) transformers can surpass CTW on non-CTW priors. The authors propose two transformer architectures to mimic CTW and explain observed behaviors, with a focus on the role of counting mechanisms in ICL.

strengths:
1. The study’s attempt to connect data compression techniques with transformer-based models is innovative and provides new insights into LLM behavior in sequence modeling.
2. The derivation of two specialized transformer structures based on CTW and PPM is well-motivated.

weaknesses:
1. The experimental setup in the paper is limited, lacking comprehensive benchmarks for language modeling. More diverse and practical tasks, such as question answering and natural language understanding, should be used to confirm the robustness of the proposed transformer models.
2. The experiments were not validated on larger-scale pre-trained language models, such as models of the scale of GPT or LLaMA, which limits the study's practical relevance. Additionally, there is a lack of discussion on how these findings could be applied to mainstream LLM tasks.
3. The paper's readability could be improved.

questions:
please refer to weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1731339177522

thread_type:
official_review

<thread_object>
creator: reviewer_m9fp

<content>
summary:
This paper investigates in-context learning ability of auto-regresive transformers on data generated by variable order markov chains (VOMC). Previous work Edelman 2024 investigated fixed order markov chains (~ngram languages), Akyurek 2024 investigated uniform hidden markov models (probabilistic regular languages), and this is in the same line but learning a different class of probabilistic languages (i.e. probabilistic models). The paper experimentally shows the solutions of transformers are better than the Context Tree Weighting algorithm which is the optimal learning algorithm for a subset of VOMC problems. The paper then goes into proofs showing that the Transformer is theoretically capable of running CTWs.

strengths:
- The paper is math heavy but it has high readability.
- The paper finds ICL on VOMC less dependent on number of layers, which is interesting.
- Experiments guide theory to construct a manually crafted Transformer that gets better performance in this VOMC dataset

weaknesses:
- The paper does not explain why studying VOMC has any practical importance — what do we learn about Transformers? What is the additional information gained after Edelman 2024, Akyurek 2024 as they studied in-context learning of very similar probabilistic models.

- The paper has a very similar structure to Akyurek 2022, Akyurek 2024, Edelman 2024 with a slight change to the class of probabilistic models used. This hurts the novelty aspect of this paper.

- There is a long theory section for two different construction, I would keep only 4.3. I would move them 4.2 to the appendix and display the main theorems only.

questions:
I think this paper is sound! But it might be of interest to a specific, purely theoretical ML audience, and I do not think it contributes significantly to the existing work. Therefore, unfortunately I lean towards reject, and I suggest authors to consider a motivation and show that their findings can have some practical importance, otherwise the writing, theory and analysis seems complete to me.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Based on the reviewers' feedback, I recommend rejecting this paper at this time. While the paper presents an interesting analysis of transformers' in-context learning capabilities for variable-order Markov chains through a data compression lens, several critical issues emerge. First, the paper's novelty is limited, as R1 points out that it largely follows the structure and approach of previous works Akyurek 2022, Akyurek 2024, Edelman 2024 with a slight change to the class of probabilistic models used. Second, the practical significance and insights gained from studying VOMC remain unclear, especially given that transformers' expressive capabilities are already well-established. Third, the experimental validation is limited, lacking comprehensive benchmarks and real-world ICL examples that would demonstrate practical relevance. While the mathematical treatment is sound and the paper is generally well-written, these fundamental concerns about novelty, practical relevance, and limited experimental validation make it difficult to justify acceptance.