Transformers on Markov Data: Constant Depth Suffices

Summary

This paper attempts to provide a possible explanation the capability of transformer architecture for efficient next token prediction of (stationary) $k$th order Markov data. The main result is a constructive proof that a $3$-layer transformer with a single head per layer can emulate conditional $k$ grams, and it necessarily uses the non-linearities in the transformer architecture like layer-normalization. The paper also provides lower bounds on representational powers of the transformers for Markov data.

Strengths

(1) The results presented in the paper serve as a crucial step towards understanding the prediction capabilities of transformers for $k$th order Markov processes with finite state spaces (of small to moderate sizes). (2) The paper provides a novel explanation of the (possible) way in which a layer-normalization (LN) might be used in modelling $k$-grams. To my knowledge, the non-linearities are discarded in typical prior work(s). (3) The lower bound is provided for single layer transformer (Theorem 5) is a strong evidence signifying of importance of multiple layers. (4) The paper is well-written and easy to follow.

Weaknesses

(1) In spite of solving a general and interesting problem of generative modelling using Markov structure, the evaluation is very limited. For example, I believe that the paper does not have empirical validation of $k$-gram modelling by the transformer when the size of the state space $S$ is large. (2) The paper does not address or comment how the training data can affect the performance during testing. This step of training might be very significant when the size $S$ of the state space is high. It would be nice to see how the mechanism works when $S$ is of order of tens or a couple hundreds (my understanding is that Figure 2 is only for binary data, $S=2$). Typo in Architecture 1: $\textbf{x}_i^{(1)} = \texttt{Emb}(x_n)$ instead of $\textbf{x}_n^{(1)} = \texttt{Emb}(x_n)$

Questions

(1) I believe that the state space is binary in your evaluations $S=2$. Can you please confirm this? (2) In Table 3 of Appendix F, the embedding dimension is mentioned to be chosen as per grid search over \{16, 32, 64\}. Should this be not the same as $d=6S+3$ as in the proof of Theorem 4 in Appendix C.3? In this case, if $S=2$, then $d=15$ is enforced. Can you provide the explanation for grid search?

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

(1) Although the result sheds light on the abilities of a transformer architecture in modelling conditional $k$-grams, this is only optimal for time-homogeneous stationary Markov data. Hence, the paper does not directly have a huge impact towards understanding how transformers perform on general stochastic sequences (like language tokens, or frames in a video). (2) The constructions in the paper work only for small state spaces (tens to a few hundreds), as the embedding dimension used in the proofs scales with its size.

Summary

This paper studies the learning process and representational power of transformers on the data sequence generated by k-th order Markov processes (or k-gram). Theoretically, this paper proved that (1) an attention-only transformer with $O(\log k)$ layers with one head for each layer can represent the conditional (in-context) $k$-gram. (2) Enhanced with MLP and layer normalization a 3-layer transformer can express the conditional $k$-gram task. They also complement their results with a lower bound for attention-only transformers/a 1-layer transformer with MLP and layer-norm with some assumptions. Detailed empirical validations are conducted to corroborate the construction (1).

Strengths

1. This paper improved the construction using $k$-head, 2-layer in previous papers on the in-context $k$-gram data by constructing a $\log k$-depth, 1-head transformer. The proof technique is based on the binary-tree-like aggregation procedure of the previous $k$ tokens information, which greatly improves the memory cost of the transformer parameters. Empirically, the construction can be partially validated by experiments showing that a 3-layer transformer can learn an in-context 8-gram. Also, a lower bound result is included (with some empirically validated assumption) to make this bound tight. 2. The construction of a 3-layer transformer with the MLP and LayerNorm is quite novel. The technique is based on the unique ternary representation of integers. The ternary representation is to maintain the uniqueness of $v_i$ after normalizing the embedding vector. The role of the non-linearity in the transformer in this $k$-gram mechanism can be crucial for improving the representation power of the model, and this result can serve as an implication.

Weaknesses

The two main construction results improve the previous expressivity results. However, the second theorem (Theorem 4) still lacks empirical evidence showing that the constructed solution is the minimizer that Adam/GD converges to. The previous experiments only show that $\log k$-depth transformer can somehow be learned. It is important to figure out whether the solution is only for construction or it can be actually learned in some way.

Questions

1. Is it possible to add a set of experiments for the 3-layer transformer with LayerNorm and MLP trained on the k-gram model when $k>8$? Or showing that it is hard to obtain from optimization?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See weakness.

Summary

This paper studies the ability of transformers to learn $k^{th}$ order Markov chains. They first conduct experiments showing that transformers with 2 layers and 1 head can learn Markov chains of up to order $k=4$. Similarly, with 3 layers, they can learn Markov chains of order $k=8$. Based on these observations, they show theoretical results about the representation power. Specifically, they present novel constructions to show that attention-only transformers with 1 head can learn $k^{th}$ order Markov chains with $\log_2(k+1)$ layers. On the other hand, 2-layer attention-only transformers require $k$ heads. These results show that increasing the depth is much more beneficial than increasing the number of heads. Next, the authors show that for the full transformer model, a constant depth of 3 suffices to learn $k^{th}$ order Markov chains with an embedding dimension of the order of the vocabulary size. This result reveals the benefit of non-linearities arising from layer normalization. They also present a lower bound on the depth requirement for 1-head attention-only transformers under some assumptions on the attention pattern learned by the model.

Strengths

- Overall, the paper is well-written and easy to follow. The authors present nice intuitions about the constructions and the theoretical results. - The results in the paper showcase the benefit of depth over the number of attention heads in transformers, which contributes to our understanding of transformers. The paper also provides insights into the benefit of non-linearities that arise from layer normalization in transformers, which is interesting, and helps compare attention-only transformers with standard transformers.

Weaknesses

- The main weakness is that some of the statements made by the authors about their results/observations contradicting prior work *lack clarity* and are not supported by a thorough *discussion of the potential reasons for the differences*. I think these statements should be accompanied by more context. Please see the Questions section for queries/suggestions regarding this. - There are some minor typos and grammatical errors that should be corrected. Please see the next section for details.

Questions

- The statement in lines 28-31 needs further clarification: - It’s not clear what range of $k$ is considered in prior works. It seems that the observations are not exactly in contradiction to prior work, since a) 2-layer 1-head transformers can’t seem to be able to learn Markov chains of order $k>4$, and b) the result in Table 1 and Theorem 2 states that for order $k$, the number of heads for 2 layers is $k$, at least for attention-only transformers. Table 1 is missing a column for the number of heads needed for standard transformers with 2 layers. - Has prior work considered training transformers with 3 layers? How does the statement after (ii) contradict prior observations? - Regarding the experiments, particularly Fig. 2: - In Fig. 2(a), the test loss gap seems a bit high. Can the authors include results for $k=1$ in this plot for a comparison? - Can the authors share results when training the 3-layer 1-head transformer on Markov chains of higher order $k>8$? Is there a reason why higher-order Markov chains were not considered for the experiment in Fig. 2(b)? I am curious if there is a gap between the theoretical result about the representation power (which holds for $k>8$), and the training dynamics. - Other suggestions: - In the abstract, it would be good to emphasize the word ‘constant’ in line 14. - I suggest using $V$ instead of $S$ for the vocabulary size. - It would be good to include some discussion on whether similar construction techniques (for Theorem 3) have been used in the literature. - I suggest including some discussion on related work on the role of softmax attention, the benefit of layer normalization techniques, and the benefit of depth in neural networks. - Typos/grammatical errors: - Missing citation in line 23. - Extra ‘studies’ in line 27. Extra ‘them’ in line 132. Extra ‘how’ in line 272. - Please check the phrasing in lines 220-222.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This paper investigates the representation capability of transformer with different number of layers or heads when learning k-th order Markov Process. They provide theoretical results demonstrating that attention-only transformers with O(log2(k)) layers can represent the in-context conditional k-th order Markov Process. This conclusion is supported by empirical results and is novel compared with previous training dynamics results. Adding LayerNorm, they prove that standard transformers with just three single-head layers can represent arbitrary order Markov processes. The paper also presents lower bounds on the size of transformers needed to represent certain attention patterns.

Strengths

1. The writting is good and easy to follow. And the proof is relatively solid. Most of the intuition is clear, like using hierarchical way to construct the intermediate logits which include information of multiple tokens. 2. The main contribution of this paper is that the authors prove there exists some transformers which have fewer head and logarithm-limited layers that can learn the determined k-th Markov Process. And these results are supported by the empirical experiments, and it's more related to the real-world case. (Like we don't need too many heads)

Weaknesses

1. This work analyzes the representation capability of transformer of learning Markov Process, however, there is no clear statement of error bounds in the key theorem like Theorem 4 (although there are some helpful discussions like Remark 2). And this may related to some questions like Questions-2. 2. The intuition of the proof of Theorem 5 is still not very clear, maybe adding more examples will be helpful, like how LayerNorm works. 3. Like mentioned in limitation parts, it's more focused on representation capability rather than training dynamics. The latter one may be more difficult for analyzing. More questions in Questions-1. 4. In Figure 2, the authors should better add k=8 and k=16 to further support the conclusion. 5. There is a citation typo in Line 23. 6. Maybe you can explicitly say that learn the conditional k-gram model is equivalent to letting transformer generate the same logits as that model (i.e., Eq.1), otherwise the reviewer maybe confused by the conclusion in the main theorems.

Questions

1. This work is more related to representation capability of transformer, rather than training dynamics, which are the focus of related works [6], [7] mentioned in the papers (Line 114 - 133), so is it possible that the construction of k heads is more theoretically friendly than the that in your paper for analyzing dynamics like gradient flow? And what's the possible way for training dynamics under your structure? 2. It seems that you still need $\Omega(k)$ bit precision for learning $k$-th order Markov Process, but in real-world case we don't need such a high precision. Is this necessary? and is it required by the constant layers transformer setting and not required by the $O(log(k))$ attention-only setting?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors have mentioned their limitations in the end of paper, like they just focus on representation capability rather than training dynamics, which maybe the futural direction.

Summary

The paper studies the representation capacity of transformers in in-context learning of order-$k$ Markov chains. First, the authors theoretically show that $O(\log (k))$ layers are sufficient to represent $k$-th order induction heads in attention-only transformers. The paper also demonstrates the benefit of non-linearities, such as layer-norm, by showing that a slightly modified transformer (modified from the original architecture) with constant depth is also sufficient to represent the in-context conditional distribution.

Strengths

Given the prevalence of in-context learning, I believe understanding the ICL of Markov chains is a very important problem, not only because of the Markovian nature of language but also the simplicity and control it provides in a synthetic setup. This is a good work that builds on recent research on ICL of Markov chains by showing that in a slightly modified transformer, scaling of heads with the order of the chain ($k$) is not necessary to represent $k$-th order induction heads. The lower bound (albeit contingent on the $k$-th order induction head assumption) also provides useful insights. I really enjoyed reading the paper; it was very clearly written and easy to follow. Generally speaking, all the ideas are laid out very clearly, with intuitive explanations that follow before and after the theorems.

Weaknesses

1. It would be beneficial to compare single head results with multiple heads (e.g., in Figure 2a, $2$ heads for $k=2$ and $4$ heads for $k=4$) as a sanity check, given previous works [1] that highlight the multi-head requirement for learning order-$k$ induction heads. I mention this because, to me, it seems like in Figure 2a, the loss isn’t approaching zero unlike in Figure 2b. 2. The work mentions "long training" to observe contradictory results (2-layer single head tfs being able to learn up to order-4 MCs in-context). How long are you training compared to [1]? How significant is this effect, and what impact do you think it has in terms of the underlying mechanism being picked up? I think this needs to be discussed further. [1] Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains, 2024.

Questions

1. What do you think is the reason 2-layer, 1-head transformers can learn up to order-$4$ MC (if they do, see Weakness 1)? What mechanism do you think is inherently being picked up? Is it learning order-$k$ induction heads or some alternate mechanism? I ask because if there is an alternate mechanism, it could be useful for the lower bound (Thm. 6) in establishing if "the lower bound representing $k$-th order induction heads implies an unconditional lower bound." 2. Regarding Theorem 4, even though the embedding dimension requirement is reduced compared to attention-only transformers (Thm. 2, Thm. 3), it still feels quite large compared to what we see in practice. Any comments on this? How tight do you think this is?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have discussed the limitations separately in the Appendix.

Discussion period ending soon: call for response

Dear Reviewer TFSx, We sincerely appreciate the time you have taken to provide valuable feedback for our work. As we are getting closer to the end of the discussion period, could you let us know if our responses above have adequately addressed your concerns? We remain at your disposal for any further questions. If you agree that our responses to your reviews have addressed the concerns you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments! Sincerely, The Authors

Discussion period ending soon: call for response

Dear Reviewer YGsY, We sincerely appreciate the time you have taken to provide valuable feedback for our work. As we are getting closer to the end of the discussion period, could you let us know if our responses above have adequately addressed your concerns? We remain at your disposal for any further questions. If you agree that our responses to your reviews have addressed the concerns you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments! Sincerely, The Authors

Thank you

We are grateful for the reviewer providing valuable feedback and their reconsideration of the score!