Understanding Transformers via $N$-gram Statistics

This paper is an attempt to answer the following

Question: How does a transformer-based large language model (LLM) make use of its context when predicting the next token?

Our approach proceeds via studying the statistical properties of training data. This is perhaps the most natural place to start even though it is not exhaustive (e.g. it does not include in-context learning (Brown et al.,, 2020)). The reasons to understand LLM behavior in terms of the statistics of their training data are plenty. First, the functional form of how LLMs use their training data is not well-understood (though there has been progress on understanding memorization (Nasr et al.,, 2023; Carlini et al.,, 2023)). Second, the over-reliance of LLMs on training data statistics leads to brittleness (e.g. the “reversal curse" (Berglund et al.,, 2024)) and the perpetuation of dataset biases (Gallegos et al.,, 2024). Understanding the nature of this statistical dependence can lead to improved and more informed dataset curation and training methods. Finally, in various scenarios, the performance of LLMs on downstream tasks are found to be correlated with frequency of relevant training data (Razeghi et al.,, 2022; Elazar et al.,, 2023; Kandpal et al.,, 2023; Kang and Choi,, 2023). A better understanding of this phenomenon would allow better steering of models towards desired performance levels.

We can think of the complexity of an LLM next token prediction (regarded as a probability distribution over tokens) along two axes: form and selection. Form refers to the functional form of the prediction as a function of the context, e.g. whether the prediction is some explicit function of associated training data statistics (see Figure 1). Selection refers to which functional form, chosen from a set of functional templates, suitably describes the transformer prediction (supposing the choice set is sufficiently rich). As a first nontrivial step, one might hope that an approximate model for an LLM is that each of its next token predictions can be roughly described by simple statistical rules from the context (simple form) even if the mechanism for its rule selection remains hidden (complex selection)¹¹1It is important to emphasize that we seek a descriptive approximation of a transformer, rather than an explanatory one. A description merely requires that we can provide a post-hoc, per-instance approximation of transformer predictions in terms of an available rule; an explanation means we provide reasons for and thus can predict in advance why and when a particular rule approximates transformer predictions. Hence, we make the distinction between form (description) and selection (explanation).. This paper is an attempt to see how far this perspective can be pushed, and fortuitously we obtain additional insights for understanding LLM behavior along the way. The statistical rules we consider, which are based on $N$-grams, are defined in Section 4, with Figure 1 showing some examples.

We perform our main investigations on the TinyStories (Eldan and Li,, 2023) dataset, with supporting experiments on Wikipedia to confirm our results remain robust at larger scales. The use of TinyStories is for practical reasons: its small size makes training models and aggregating $N$-gram statistics computationally efficient, yet it is complex enough to capture basic natural language statistics (those occurring in simple children’s stories).

Below is a summary of our observations and contributions:

1. 1.

   (Approximation-Variance Association) We observe an approximation-variance association indicating that next token LLM predictions that have low variance (across different training runs²²2Different runs have different dataset shuffles.) tend to be well-approximated by $N$-gram rules. (Section 5)

2. 2.

   (Curriculum Learning Dynamics) By grouping our $N$-gram rulesets in terms of complexity (as measured by the amount of context they use), we discover the various ways in which the learning dynamics of LLMs implement a statistical type of curriculum learning, in which easier rules are eventually supplanted by more complex ones. (Section 6.1)

3. 3.

   (Overfitting Criterion) Based on our analysis of approximating LLM predictions by $N$-gram rules, we propose a simple and novel procedure for detecting overfitting of LLMs during training. The procedure makes no use of holdout data and it makes quantatively precise the intuition that overfitting corresponds to a model memorizing long context at the expense being able to generalize through making use of subcontext. (Section 6.2)

4. 4.

   (Approximation Strength) We study how well LLM predictions can be approximated by our $N$-gram rulesets, noting that significant gains in top1-accuracy occur as we increase ruleset complexity and diversity, whereby we achieve up to 78% top-1 accuracy on TinyStories (Table 2). We also visually ground these approximations with concrete examples (Figure 5), which may form the basis for dataset attribution methods in future work. (Section 7)

Rule extraction methods for neural networks have been studied in quite different settings, e.g. (Jacobsson,, 2005; Mcmillan et al.,, 1991). Some recent works have performed $N$-gram analyses for large-language models in the setting of in-context learning (Akyürek et al.,, 2024) and associative recall (Arora et al.,, 2023). The “infini-gram" model (Liu et al.,, 2024) compares LLM predictions with the single $N$-gram rule given by retrieving the largest possible matching context from the training data. Our work uses shorter but more sophisticated $N$-gram rules. In (Voita et al.,, 2023), an approach to understanding how LLMs process $N$-grams is carried out at the level of individual neurons. This complements our dataset-based work, which treat models as a black box. In (Edelman et al.,, 2024), the evolution of the type of $N$-gram statistics that transformers learn during training is analyzed in the setting of synthetic Markov chain data, in contrast to our natural language setting. Other works studying the learning trajectory of language models include (Chen et al.,, 2024; Choshen et al.,, 2022). There is a large literature on building more sophisticated $N$-gram models, e.g. (Kneser and Ney,, 1995; Goodman,, 2001). Such models could have been incorporated into our set of rules, but for simplicity we choose not to include them.

We train standard decoder-only transformer models on the TinyStories (Eldan and Li,, 2023) dataset (480M tokens) consisting of children’s stories synthetically generated from GPT-4 using templated prompts. The value of this dataset lies in its linguistic simplicity, whereby it is possible to model language well on the dataset using very small models. Unless stated otherwise, our experiments use a 160M parameter model trained for 4 epochs, which achieves a loss of around 1.11 nats on the validation set. We train for 4 epochs since we use learning rate warmup and cosine learning rate decay and we want to ensure all datapoints receive updates with a high learning rate (this way all $N$-gram statistics have a fair chance of being learned during training). For overfitting experiments in Section 6.2, we train a 1.4B model for 10 epochs. In the Appendix, we include some additional corresponding experiments on Wikipedia (from MassiveText (Rae et al.,, 2022)) with a single epoch of training in order to validate that our results are of a general nature and extend to more complex datasets. For a fixed dataset, the only source of randomness among different runs are different dataset shuffles. Full experimental details are described in the Appendix.

The attention layer within a transformer is in essence a soft context-selection mechanism. The $N$-gram rules we consider will be loosely modeled on this mechanism. Namely, given a context we will form a derived context in which each token will either be kept, discarded, or marginalized, which is meant to mimic positive attention, no attention, and semantic invariance³³3For instance, the next token distribution for the context “… the tired dog” may be insensitive to replacing “tired” with “brown” or “furry”. Statistics which thus marginalize over all extant substitutions for “tired” yield a crude but generally applicable way of capturing semantic invariance. One can imagine an attention mechanism for which there is a many-to-one mapping of keys to a particular value that might implement semantic invariance., respectively. More formally, we proceed as follows:

Given a regular expression⁴⁴4Our regular expressions operate on tokens not string characters, since our contexts are formed out of sequences of tokens. $\alpha$, all contexts from the training data can be retrieved which match the regular expression. This allows us to define a corresponding rule that defines for us a distribution over tokens $t$:

$$R_{\alpha}(t)=\frac{\#\{\alpha t\}}{\#\{\alpha*\}}$$

(1)

where the numerator and denominator are the counts for the $N$-grams from the training data matching the concatenated regular expressions $\alpha t$ and $\alpha*$, respectively, where $*$ is wildcard (single) character match⁵⁵5We use $*$ (i.e. glob notation) instead of the standard . symbol for readability purposes.. (Thus the $N$-grams in the numerator end with $t$ while those in the denominator can end with any token.) Observe that the next-token predictions of a vanilla $N$-gram model are obtained by letting $R_{\alpha}(t)$ vary over all ordinary token sequences $\alpha$ of length $N-1$.

Given $\sigma$, a symbol from the the alphabet $\{*,-,+\}$, consider the following operation which maps a token $t$ to a regular expression:

$$S_{\sigma}(t)=\begin{cases}t&\sigma=+\\ *&\sigma=*\\ \epsilon&\sigma=-\end{cases}$$

(2)

where $\epsilon$ is the empty regular expression. Given now a sequence $\sigma=\sigma_{-N}\cdots\sigma_{-2}\sigma_{-1}$, define $S_{\sigma}$ on a sequence of tokens $C=C_{-N}\cdots C_{-2}C_{-1}$ by tokenwise application of (2) and concatenation⁶⁶6The empty regular expression does nothing under concatenation and does not contribute to the length of the resulting sequence.:

$$S_{\sigma}(C)=S_{\sigma_{-N}}(C_{-N})\cdots S_{\sigma_{-2}}(C_{-2})S_{\sigma_{-1}}(C_{-1}).$$

(3)

Thus (3) defines a regular expression which we can think of as fuzzy matching for a subset of a context $C$ (the fuzziness arising from the presence of wildcard matches). For notational convenience, we assume $\sigma$ is left padded with $-$ symbols, so that we can define $S_{\sigma}(C)$ for $\mathrm{len}(\sigma)<\mathrm{len}(C)$. Finally, define

$$R_{\sigma}(t|C)=R_{S_{\sigma}(C)}(t)$$

(4)

for $C$ with $\mathrm{len}(C)\geq\mathrm{len}(\sigma)$. The collection of (4) for various $\sigma$ defines our $N$-gram rules under consideration⁷⁷7For $\sigma=\emptyset$, we define $R_{\sigma}$ to be the unigram distribution.. Each such rule is a function which maps a context $C$ to a next token distribution. We refer to $S_{\sigma}(C)$ as the rule context for $R_{\sigma}(t|C)$.

As concrete examples, let $\sigma=+-*+$. If $C=C_{-5}C_{-4}C_{-3}C_{-2}C_{-1}$, then $S_{\sigma}(C)=C_{-4}*C_{-1}$ and

$$R_{+-*+}(t|C)=\frac{\#\{C_{-4}*C_{-1}t\}}{\#\{C_{-4}*C_{-1}*\}}$$

(5)

is a rule which yields a next token distribution based on a particular combination of $4$-gram model statistics: it retrieves all three token contexts in the training data whose first token is $C_{-4}$ and last token is $C_{-1}$ and marginalizes over the second token. Likewise, the rules

$$R_{++--}(t|C)=\frac{\#\{C_{-4}C_{-3}t\}}{\#\{C_{-4}C_{-3}*\}}\qquad R_{++**}=\frac{\#\{C_{-4}C_{-3}**t\}}{\#\{C_{-4}C_{-3}***\}}$$

(6)

are respectively a trigram model with context $C_{-4}C_{-3}$ (all other tokens receiving a $-$ are dropped) and a model which uses four tokens of context but marginalizes over the two most recent ones.

When $\sigma$ consists of all $+$ symbols, we get vanilla $N$-gram rules derived from the suffix of $C$. When $\sigma$ consists of $\pm$ symbols, we get vanilla $N$-gram rules derived from subsets of $C$. Varying the length and the entries of $\sigma$ yields the following rulesets⁸⁸8 There is some redundancy among the $\sigma$’s in terms of the rules they determine: for instance, in between any two $+$ consecutive symbols, permuting the order of $-$ and * will determine the same rule. Also in practice, we can assume the first entry of $\sigma$ is a $+$ since marginalizing the first token is equivalent to reducing the context length. From this, it follows that the number of distinct rules in $\mathcal{R}^{all}_{M}$ is $2,5,13,34,89,233,378$, for $M=1,\ldots,7$, respectively. :

	$\displaystyle\mathcal{R}^{\textrm{suffix}}_{M}$	$\displaystyle=\{R_{\sigma}(t|\cdot):|\sigma|\leq M,\sigma_{i}=+\textrm{ for all i}\}$		(7)
	$\displaystyle\mathcal{R}^{\textrm{subgram}}_{M}$	$\displaystyle=\{R_{\sigma}(t|\cdot):|\sigma|\leq M,\sigma_{i}=\pm\textrm{ for all i}\}$		(8)
	$\displaystyle\mathcal{R}^{\textrm{all}}_{M}$	$\displaystyle=\{R_{\sigma}(t|\cdot):|\sigma|\leq M\}.$		(9)

The parameter $M$ controls how much of the context is being used for the rules.

Let $p(t|C)$ denote the next-token distribution of an LLM conditioned on the context $C$ and for notational similarly, write $p_{r}(t|C)$ for $r(t|C)$, where $r$ is one of the rules defined in Section 4. We wish to measure how similar these distributions are (higher similarity corresponds to a better rule description). To that end, we use the variational distance to measure the difference of two distributions (we discuss our choice in the Appendix):

$$d(p,q)=\frac{1}{2}\sum_{i}|p_{i}-q_{i}|.$$

(10)

Since variational distance may be lacking in concrete interpretability, we will sometimes use top1-accuracy to measure similarity, defined by

$$\textrm{top1-acc}(p,q)=\frac{|\mathrm{argmax}(p)\cap\mathrm{argmax}(q)|}{|\mathrm{argmax}(p)\cup\mathrm{argmax}(q)|}$$

(11)

(in general, the argmax of a probability distribution is a set due to potential ties among maximal probabilities). When the argmaxes in (11) are singletons, top1-accuracy just measures agreement between greedy predictions.

Given a context $C$, we want to understand how $d(p(t|C),p_{r}(t|C))$ varies with different rules $r$ and in particular if it can be made small. To that end, we introduce some terminology:

We are interested in determining the optimal rule $p_{r}(t|C)$ (as defined in Table 1) and if it has small optimal rule distance then we regard the rule as being a good description of the corresponding transformer prediction(s) $p_{i}(t|C)$.⁹⁹9In practice, we will only have one model available and our optimal rules are computed per-context and per-model. In this section, we have available five models from five runs for use in computing optimal quantities. As a first step, note there is a distinguished rule

$$p_{\textrm{full}}(t|C)=\frac{\#\{Ct\}}{\#\{C*\}}$$

(12)

whose rule context is the full unmodified context $C$.¹⁰¹⁰10That is, $p_{\textrm{full}}(t|C)$ is the invokation of the rule corresponding to $\sigma=+\cdots+\in\mathcal{R}^{\textrm{suffix}}_{|C|}$ applied to $C$. This is because (roughly) the language-modeling objective aims to make $p(t|C)$ similar to $p_{\textrm{full}}(t|C)$.¹¹¹¹11See Section C for additional discussion. All other rules in our rulesets are “subleading" in that they drop or marginalize over tokens in the context $C$. Our goal is to quantify which rules, either (12) or subleading ones, are optimal rules and what their optimal rule distances are.

Our main finding is an approximation-variance association: contexts with low model-variance tend to have low optimal rule distance. The surprising aspect of this association is the sufficiency of low model-variance (necessity is a given).¹²¹²12Predictions which have high variance cannot be well approximated by a single model-independent rule. We use five runs in our analysis here since approximation by a rule that remains fixed across models yields a fortiori approximation by a per-model rule. We present the case of $7$-gram contexts in Figure 2 to corroborate this association, with additional examples relegated to the Appendix. We sample around six-thousand $7$-grams from the training data, sampling from logarithmically spaced buckets based on counts, and plot various relations between counts, model variances, and rule distances. For simplicitly, we consider the ruleset $\mathcal{R}=\mathcal{R}^{\textrm{suffix}}_{7}$ to limit the number of rules under consideration. Our analysis of Figure 2 can be summarized as follows:

Plot (d) summarizes our approximation-variance association. There is a clear positive correlation and reasonable fit between optimal rule distance vs model variance when using the ruleset $\mathcal{R}^{\textrm{suffix}}_{7}$. This fit weakens due to many outliers if we only consider the single $p_{\textrm{full}}$ rule (a vanilla $8$-gram model) as shown in (b). These outliers correspond to many LLM predictions being poorly approximated by $p_{\textrm{full}}$. The transition from (b) to (d) is a way of visualizing LLMs performing back-off, whereby LLMs rely on statistics from subsets of the context. We include plots (a) and (c) to highlight how replacing model-variance with the count of a context would lead to a much worse fit. We highlight count of the context $C$ because it is the most obvious and naive measure of how well one should expect $p(t|C)$ to match the rule $p_{\textrm{full}}(t|C)$. Nevertheless, the weak correlation between count and distance measures makes sense: an LLM can make predictions based on subcontexts of $C$ and those subcontexts can have very different count statistics than those of $C$.

We believe our approximation-variance association and its corresponding analyses have significance beyond the experiments carried out here since they (i) highlight that naive count-based statistics do not provide the strongest signal in terms of how LLMs leverage dataset statistics (ii) suggest that LLM predictions that have low-variance are likely the ones that are amenable to description (or even explanation) by some underlying dataset statistic (with high-variance predictions being regarded as noise). We leave a more systematic exploration of (ii) to future work.

Finally, addressing our main question from the introduction, we track how well our rulesets describe LLM predictions (in the sense of Section 5) as a whole at inference time. Here, the utility of our $N$-gram rules defined in Section 4 becomes apparent, since on a holdout set, there will be novel contexts and being able to drop or marginalize context tokens aid in being able to retrieve or aggregate corresponding training dataset statistics. In Table 2, we show the average top1-accuracy between the optimal rule from our various rulesets and LLM predictions on 100 random stories from the validation set. Here, we include as baseline $\textrm{backoff}_{M}$, the single rule given by the predictive model which performs “stupid backoff" (Brants et al.,, 2007) using $M$ tokens of context.¹³¹³13That is, $p_{\textrm{backoff}_{M}}(t|C)=p_{\textrm{full}}(t|C_{-k}\cdots C_{-1})$ where $k\leq M$ is the largest value for which $C_{-k}\cdots C_{-1}$ occurs in the training data.

We see significant gains in accuracy at large $M$ when adding additional types of rules (for $M=7$ we gain $6\%$ each time in going from “suffix" to “subgram" to “all"). In the end, we are able to obtain 78% top1-accuracy between the per-prediction optimal rule and the LLM predictions, averaged over all tokens. This is perhaps a remarkably high figure, considering that the top1 accuracy of the model with respect to the ground truth on the validation set is 69.6%. At minimum, we have provided a precise quantification of structure in LLM next-token predictions: they are often matched (as measured by top token prediction) by some simple $N$-gram rule derived from the training data. See Section D.1 for some supplementary analysis.

To ground our rule optimization procedure, we provide Figure 5 which shows side-by-side how LLM predictions compare with ground truth and optimal rule predictions in an example heldout story. For instance, for the target token “climb” in “... Roxy loved to climb”, both the LLM and optimal rule $R_{\alpha}$ predict “play”, where $\alpha=\textrm{``}\texttt{. * loved to}$”. For target token “climb” in “... She climbed”, the LLM predicts “would" whereas the ground truth and $R_{\alpha}$ predict “climb”, where $\alpha=\textrm{``}\texttt{loved to climb * She}$”. In general, optimal rules provide the closest statistical match from the training data to the given LLM predictive distributions (from amongst our rulesets), and their top1-predictions can agree or disagree agree (as indicated by target token color). Additional examples, including those from Wikipedia, are shown in Section D. For interpretability purposes, we re-emphasize that our optimal rules currently only provide descriptions, not explanations. We leave the possibility of the latter for future work.

Our work provides quantitative measures of how well the predictions of transformer-based LLMs are described (i.e. approximated) by simple $N$-gram rules. Such rules were motivated by the simplest token-level operations applied to the context (keep, ignore, or marginalize). The results we obtained in Section 7 imply that, at least on simple datasets like TinyStories and Wikipedia, LLM predictions contain much quantifiable structure insofar that they often can be described in terms of our simple statistical rules. Along the way, we also obtained novel discoveries into the statistical nature of overfitting, the occurrence of curriculum learning, and the relation between model-variance and approximability by $N$-gram rules. Altogether then, our work provides various avenues of progress in understanding how simple dataset statistics are reflected in LLM behavior.

On the other hand, it is intuitively clear that current state-of-the-art LLMs go well beyond invoking $N$-gram rules. A typical request to perform a nontrivial task (e.g. “Write a thirty line poem about mathematics that rhymes") requires a high-level conceptual understanding of language that goes beyond simple literal token-level associations between the context and the training data that we consider here. Nevertheless, one can speculate that an analogue of our work could still apply: in general, an LLM might be performing some high-level rule application, whereby statistics formed out of distributional categories (Pereira et al.,, 1993) instead of individual tokens are leveraged from the context. Formulating a correct and parsimonious set of rules, if it is at all possible, would be a nontrivial challenge to overcome and one which we leave to future work. Addressing such a challenge and being able to promote the descriptive approximations provided here to explanatory ones would provide a next step towards understanding how LLMs work.

The author would like to especially thank Senthooran Rajamanoharan for numerous conversations and an exceptionally discerning eye, which greatly improved the paper from earlier drafts. The author also thanks Jonathan Hale, Marcus Hutter, Matthew McGill, Nick Roy, Avraham Ruderman, and Daniel Tanis for helpful feedback and discussions. Finally, the author thanks Frank Perbet and Daniel Tanis for engineering support.