Transformers on Markov data: Constant depth suffices

We thank the reviewer for the feedback and constructive criticism. We have fixed the typo pointed out. Below we address the main weaknesses discussed in the review:

[W1,W2] Experiments for and learning dynamics under larger state spaces?

In the attached rebuttal pdf, we ran experiments for Markov models on state spaces of size $|S| = 5$. Here too, we show that the main phenomenon discussed in the paper ($2/3$-layer $1$-head transformers learn $k$-th order Markov processes for $k > 1$) holds true. In particular, we show that a $2$-layer $1$-head transformer learns order-$3$ Markov processes. We were unable to run this experiment for $k=4$ keeping the state space size fixed since the model longer to converge than our computational allowance would permit. When the state space becomes even larger, say on the order of tens or a couple of hundreds, it becomes much harder to train the models on reasonably large values of $k$ within the computational budget in our availability.

[Q1] Binary state space

You are correct to note that we run our experiments on binary data (to be able to scale to the largest possible values of $k$ on a computational budget).

[Q2] Precise dependency of embedding dimension

The experiments we run are more in line with the case of the standard transformer architecture, which corresponds to Theorem 4, where the embedding dimension scales as $O(S)$ (where the constant is implicit). In the common rebuttal, we discuss how to theoretically reduce the dependency on the embedding dimension at the cost of a small increase in the loss incurred by the model.

On a separate note, there are usually a few reasons to expect that there would be some differences in theoretical and empirical findings especially when it comes to the precise dependency on quantities like the embedding dimension, etc. Empirically, transformers are trained from a random initialization and in a way where the model may be more / less efficient at realizing certain mechanisms compared to a theoretical construction. We would argue that it is in fact promising that the embedding dimension is of a similar order as what is predicted in our theorems, as it is often the case that the theoretical and empirical findings are off by several orders of magnitude.

[L1] Time inhomogeneous processes

This is a great question, and something we hope to study as the next step. Time homogeneity of the data model (such as with $k$-th order Markov processes) implies that it is possible to write down the estimate of the next-symbol conditional distribution using an unweighted empirical estimate. When the data process becomes time inhomogeneous, the statistical estimators would look more akin to weighted conditional empirical estimates (under assumptions on the nature of time inhomogeneity). The constructions we consider in the paper do have the ability to translate to these kinds of settings as well. However, to understand the limits of how transformers do this is a great direction for future research, and something we hope to explore going forward. It is plausible that transformers need more depth to be able to capture these kinds of processes, depending on the nature of the time inhomogeneity.

[L2] Constructions only apply for small state spaces

In commercial (for which metadata is available) and larger open source models, the vocabulary size has usually been in the range 30K-250K. However, the embedding dimensions have been of the scale of 1000-20,000, which is one order of magnitude smaller. It is worth pointing out that our constructions, while optimized for the dependency on the depth and the number of heads, can also be improved on the depth dependency. In the common rebuttal, we discuss how to improve the dependency of the embedding dimension from $O(S)$ to $O(\log (S)/\epsilon^2)$ where $\epsilon$ is a parameter which now captures an approximation error. Thus, even when the state space is larger, for practical ranges of $\epsilon$, the dimensionality implied by this result should be much smaller than $O(S)$.

We thank the reviewer for the positive comments, and constructive criticism. Below we address the main weakness pointed out, as well as the question asking about experiments for $k > 8$.

[W1] Evidence that ADAM/GD converges to theoretical construction

This is a great question. In Fig. 3 of the rebuttal pdf attached, we plot the attention patterns learnt and observe some sort of approximately exponential decay. The theoretical constructions also match this kind of exponential decay (as seen from the attention pattern in eq. (50) in the proof of Theorem 4). While this is promising, there is certainly more to explore about how and what the transformer learns. In particular, it is interesting to see how these attention patterns evolve over the course of learning before converging at these kinds of exponential patterns. This may reveal some more about how transformers learn, and get us one step closer to a theoretical understanding of the learning dynamics of transformers when exposed to Markovian data.

[Q1] Experiments for $k > 8$?

This was also pointed out by other reviewers, and is indeed an important point. We would be happy to include something which provides more clarification around what happens when $k=8$, however in the form we currently consider, it is not within our computational budget to scale this to $k=9,10$ or anything higher. A longer discussion as to why this kind of computational “phase transition” occurs is provided in the common rebuttal above.

We thank the reviewer for the detailed comments.

[W1] Thorough discussion for why contradictions arise in past work

As discussed in the common rebuttal, the main reason for the differences in observations compared to [1] and [2] is the fact that the models were not trained for sufficiently long to observe the increase in test loss. Please refer to Fig. 2 in the attached document where we plot this explicitly on a random $5$-state order-$3$ Markov chain. In the beginning phase of the training (iterations 1-200), the test-loss curves plateau around the bigram performance, but as we continue training, the loss continues decaying slowly. We believe that this difference only appears because the model was not trained for sufficiently long, keeping all other variables unchanged. We will include this plot in the paper.

[Q1] Range of $k$ considered in previous work

The contradiction we mention in the $2$-layer case is empirical. Prior work such as [1],[2] train $2$-layer $k$-head models on $k$-th order Markov chains for $k=1,2,3,4$. However, we are lead to believe that the $2$-layer $1$-head model was not trained for sufficiently many iterations to learn $k=2,3,4$.

Fig. 2 in the attached material showcases this difference when a $2$-layer $1$-head transformer is trained on $k$-th order Markov chains on $5$ states for $k=3$. The initial phase (first ~100 iterations), the model stagnates in performance at the level of the best bigram model. When we continue training, the loss of the model starts decreasing further and breaks past this bigram phase to learn higher order conditional $k$-gram models.

Finally, it’s worth noting: while we do not have a column for the $2$ layer standard transformer in Table $1$, we believe that the main $3$-layer construction can also work for $2$-layer transformers with $2$ heads in the first layer. This is because in the $3$-layer construction we consider, the first two layers of the model essentially do not interact with each other (and operate on different parts of the embedding space), they can be implemented across two different heads of the first layer. We will add this into the paper to clarify that there exists a contradiction even theoretically.

Formally: For $2$-layer transformers with $2$ heads, prior work argues that these models can learn up to $k \le 2$. Our work shows that these models can represent much higher values of $k$ theoretically (via an extension of Theorem $4$).

[Q2] Clarifications regarding Fig. 2

We have included the plot for $k=1$ in Fig. 1 of the attached document in the paper (this experiment trains $k$-th order Markov chains on $k$-head $2$-layer transformers). We will include this plot in the paper as well. There is a gap in the test loss which decreases when we move to $3$-layers and $1$-head. The reasons for this small but non-negligible test loss gap are unclear - while our theoretical construction works for the $3$ layer case, it may be possible that with $2$ layers and $1$ head transformers are unable to exactly represent the conditional $k$-gram model. It may also be possible that ADAM/GD are unable to find this solution, even if it is exactly representable. Understanding the limits of how $2$-layer $1$-head models learn is an interesting question to look into.

On a separate note, It is worth mentioning that the best achievable test loss will increase as we keep the sequence length fixed and increase $k$. In the learning setup we consider, the training and test data do not come from the same Markov process, since the transformer is trained on data from randomly chosen Markov processes and tested on data from another randomly chosen Markov process. So the transformer has to do some form of in-context learning. When the sequence length is fixed and $k$ increases, the in-context learning problem becomes harder. The model has to learn $S^k$ conditional distributions (for each possible prefix of $k$ states). So as $k$ grows, the model has less data for each of these distributions which makes the estimation task harder. This is a major reason for why the $3$-layer transformer’s loss increases a little bit when dealing with $k=8$. We were computationally bottlenecked to train on sequence lengths of up to $500$ and at this scale, the model has just about enough data to learn each of the $256$ conditional distributions. Fixing the sequence length as $500$, the models should not be expected to have small test loss as $k$ grows to $9$ and beyond.

The reason we do not train on $k > 8$ is discussed in the common rebuttal. The sequence length to consider would have to be at least around $1000$ or longer for $k=9$ and $2000$ for $k=10$. The model also takes longer to train, and we were unable to train for long enough for the model to reach convergence. In the future, we hope to be able to test out the results in this paper on longer sequence lengths (thereby allowing larger values of $k$) while keeping the depth of the model and the number of heads fixed.

[Q3] Additional suggestions

We shall plan to incorporate the notational changes and fix the typos as suggested by the reviewer. With the additional page in the final version, we shall also include a discussion on related work on the role of softmax attention and depth which were present in existing work. Finally, we will also include a discussion on related works which consider techniques similar to the composition technique we considered in Theorem $3$.

References

[1] Makkuva, Ashok Vardhan, et al. "Attention with markov: A framework for principled analysis of transformers via markov chains." arXiv preprint arXiv:2402.04161.

[2] Nichani, E., Damian, A. and Lee, J.D., 2024. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735.

[3] Edelman, Benjamin L., et al. "The evolution of statistical induction heads: In-context learning markov chains." arXiv preprint arXiv:2402.11004.

We thank the reviewer for their detailed comments. Below we address the main questions and weaknesses:

[W1] Error bounds

We are happy to include a longer discussion of error bounds in the paper. To expand, suppose all the weights in the transformer model are upper bounded by $1$. When the transformer is implemented with a precision of $m$ bits, the new attention weights $\widehat{\text{att}}_{n,i}$ satisfy,

This uses the fact that by truncating vectors to $m$ bits, the approximation error in the inner product of the inner product of two $d$-dimensional vectors. i.e. $\varepsilon$, is at most $d \cdot 2^{-m+1}$. In particular, simplifying assuming $\varepsilon$ is small, we get a multiplicative error of $[ 1 - \Omega (\varepsilon), 1 + \Omega (\varepsilon)]$. A similar analysis works for layer-normalization and results in a similar error scaling. Thus, within a single attention layer, we compute vectors which are within $[1-\varepsilon', 1+\varepsilon']$ where $\varepsilon' = T \cdot 2^{-m-\log(d)}$. Iterating across $L$ layers, the error incurred scales as a multiplicative $1 \pm 2^{-m-L\log(dT)}$. Thus, when the bit-complexity scales as $m \gg L \log (dT)$, the approximation error is a multiplicative polynomially small constant in $T$.

There is only one caveat in case of the constant depth transformer, the weights of the transformer in each of the $3$ layers are upper bounded by $O(2^{k})$, rather than $1$. Thus, the bit precision must now scale as $k + O(\log(ST))$.

We chose to de-emphasize these details in the paper to avoid unnecessarily complicated theorem statements, and to avoid detracting from the (arguably more interesting) other phenomena discussed in the paper.

[W2] Intuition behind Theorem 5

We have rewritten this section to make the intuitions more clear. We have added a new “proof sketch” section, clarified what the $k$-th order induction head does in this context, and emphasized how layernorm allows the transformer realize these $k$-th order induction heads.

In short, layernorm allows changing attention from $\text{att}\_{n,i} \propto \exp (\langle k_i, q_n \rangle)$ to instead look like $\text{att}\_{n,i} \propto \exp ( - \| \hat{k}\_i - \hat{q}\_n \|\_2^2)$ where $\hat{k}\_i$ and $\hat{q}\_n$ are the $L_2$-normalized key and query vectors. This is nice, because if we can realize the key vectors $k_i = \sum_{j=1}^k 2^j e_{x_{i-j}}$ and the query vectors as $q_n = \sum_{j=0}^{k-1} 2^j e_{x_{n-j}}$, the attention is maximized if and only if $\hat{k}_i = \hat{q}_n$, which we can prove occurs only when $x\_{i-1} = x\_n, \cdots, x\_{i-k} = x\_{n-k+1}$ (using the fact that the binary representation of a number is unique). This realizes a $k$-th order induction head.

[W3] Representation vs Training dynamics

In this paper, we focus on the representation capacity, which itself turns out to be a complicated phenomenon. On the optimization side, there are indeed many questions open. Unfortunately, even in the simplest settings (1 layer transformers), the analysis of gradient descent on transformers trained on Markovian data is incredibly complex. The rigorous analysis in [Makkuva et al] is over 70 pages long for this case. Extending to higher depth, is incredibly non-trivial. [Lee et al] discusses this case, however, their analysis is only amenable under very strong assumptions i.e., the “reduced model” of transformers. Even with these assumptions, the analysis is over 60 pages long. We hope to convince the reviewer that while training dynamics is indeed an interesting and timely question, it certainly requires a dedicated effort to be able to analyze, and would constitute a separate work of its own.

[W4] Experiments on higher values of $k$

Discussed in the common rebuttal.

[Q1] Learning dynamics

At a high level, yes, we believe that the k-head mechanism discussed in prior works may be more amenable for analysis. When the depth is constant, the transformer is forced to be more creative in the ways it learns the $k$-th order chain, and this leads to more intricate mechanisms. However, unless simplifying assumptions are imposed, we believe that understanding learning dynamics falls into the realm where current tools in optimization theory are not strong enough to be able to answer.

[Q2] $\Omega(k)$ bit precision for learning $k$-th order Markov processes

One may consider an extension of $k$-th order Markov processes to the case where the conditional distribution depends on a sparse subset (of size $p$) of the previous symbols observed. The conditional distribution has a low degree of dependency on the past, even though it is not the immediately preceding $p$ symbols. While these kinds of processes are special case of $k$-th order Markov processes, modeling them as such may necessitate a very large value of $k \gg p$, since the there may be dependencies on symbols which appeared long ago.

Our constant depth constructions work even in this setting, with the bit-complexity scaling with $p$ and not with $k \gg p$. This analysis shows that transformers don’t require high bit-precision as long as the degree of dependency of the conditional distribution on the past is small. In practice, this is often the case.

On a separate note, our constant depth constructions require $\Omega(k)$ bit complexity, but the embeddings are now $O(S)$ dimensional, rather than being $\Theta(Sk)$ dimensional in the attention-only setting. Thus, the total number of bits in each embedding vector is the same in both cases. We believe that it is possible to reduce the bit complexity of the constant depth construction at the cost of making the embeddings $\Theta (Sk)$ dimensional.

We thank the reviewer for the comments and questions about the paper. Below we address the main questions and weaknesses raised:

[W1] Comparing single vs. multi-head

In the attached material, please refer to Fig. 1, which we will add into the paper in the subsequent version (as the new Fig. 1). This plot discusses how the multi-head transformer learns, while the number of iterations to converge to optimality is of the same order, the test-loss these models achieves is nearly $0$. In comparison, the single head $2$-layer model has non-zero but small test-loss, while the single head $3$-layer model has even smaller test loss. At this scale, we did not distinguish between the small loss of the transformer vs. negligible loss incurred by the multi-head setup. It is indeed plausible that with better choices of hyperparameters (optimizing over embedding dimension, for instance), we may see the delta become even smaller.

In our evaluation setup, the transformer does not know what the distribution of samples at test-time are (only that they are Markovian). Thus, any model which achieves low test-loss has to resort to something like an in-context estimate. We were computationally bottlenecked to train transformers on sequence lengths of up to around $500 \approx 2^9$ for the number of iterations required for the test-loss to begin to converge. At this scale, the model has “just” enough data to be approximately able to estimate order-8 Markov chains in context (i.e., having to approximately estimate $2^8 \approx 250$ conditional distributions on $\{ 0,1\}$). At this scale, if were to increase $k$ the model would stop performing well, and more importantly, no in-context estimator could work at this scale. Information theoretically, there is just not enough data in a test-sequence to be able to estimate the next sample distribution, even approximately.

TLDR; When the sequence length is around $500$, $k=8$ is the point where we would start to see degrading performance (for any model, and not just a transformer) by virtue of approaching the regime where we don’t have enough samples to estimate conditional probabilities from the test-sequence accurately.

[W2] Long training to observe contradictory results

This is discussed in the common rebuttal above.

[Q1] Why do 2-layer 1-head transformers learn order-4 Markov chains?

This is a great question. While our construction shows that the standard $3$-layer transformers with $1$ head are able to represent $k$-th order induction heads, it may be possible that at an even smaller scale (say $2$-layer with $1$ head) that transformers are still able to represent the in-context conditional $k$-gram with some amount of error. It is also possible that at an even smaller scale, (say $2$-layer $1$-head transformers for sufficiently large $k$, or if the embedding dimension is reduced even further) that induction heads are no longer representable at all. Understanding this regime presents new avenues of research: here the transformer may still achieve low loss, while “provably” not being able to use the induction head mechanism to achieve this loss. It’s worth pointing out the case of $1$-layer $1$-head transformers which were studied in [1]. The authors show that transformers are still able to learn Markov chains, but under this extreme size constraint, learning does not happen in a transductive way. At this scale, the transformer does not do in-context learning, rather learns the parameters of the Markov chain from which data is generated directly. This is a slight departure from our setting, since learning the parameters is infeasible when the training and test data come from different distributions (such as two different randomly sampled Markov chains).

[Q2] Precise dependency on embedding dimension

This is discussed in the common rebuttal above.

References

[1] Makkuva, Ashok Vardhan, et al. "Attention with markov: A framework for principled analysis of transformers via markov chains." arXiv preprint arXiv:2402.04161 (2024).