Tracing LLM Behavior to the Training Data with
Empirical Next-Token Distributions
July 15, 2026
Abstract
In this paper, we study the connection between an LLM’s output distribution and the data used to train it. Specifically, we study the degree to which an LLM’s next-token distribution agrees with the empirical next-token distribution (ENTD) given the context in the training data. The ENTD is an appealing target because it is the unrestricted global minimizer of the next-token cross entropy loss used for pretraining, as well as an easily interpretable function of the pretraining corpus. We find that for a significant fraction of inputs, the LLM’s distribution agrees with the ENTD almost perfectly, and the average agreement increases with model scale and training compute. Nevertheless, there is a long tail of input sequences where the LLM and ENTD differ significantly, and we examine several possible sources of this discrepancy across the transformer architecture, training procedure, and finite-sample noise in the ENTD estimate itself. More broadly, we hope our findings will encourage more work on “data-centric mechanistic interpretability,” a complement to standard mechanistic interpretability that opens the black box of how model behaviors arise from the data, rather than how they are encoded in the learned weights.
1 Introduction↩︎
The goal of this paper is to enhance our understanding of how LLM capabilities emerge from the training data. It is well-known that the training data are critical for shaping what the model can and cannot do [1]–[4]. In spite of this, a more precise, first-principles-based explanation of where model capabilities “come from” in the data remains elusive. This paper is an attempt to make progress on this front.
To this end, we study how well a trained transformer approximates the empirical next-token distribution (ENTD) of its data. The ENTD is the global minimizer of the next-token prediction cross-entropy loss used in pretraining, and it is easily interpretable as a fully “memorized” model of the data: the predictions for the next token are simply the empirical distribution of next tokens, given the context. The existence of in-context learning capabilities (e.g., induction heads [5], which can make sensible predictions even on semi-random, completely out-of-distribution inputs) means that the ENTD cannot capture the full capabilities of modern LLMs. Nevertheless, we believe that studying the extent to which an LLM approximates the ENTD is an important initial step to understanding the link between training data and trained model.
We therefore examine how the LLM’s predictions relate to the ENTD predictions as training progresses across a wide range of model scales. Our results reveal a complicated picture with deviations between the model and ENTD potentially arising from a variety of sources, including sampling error, optimization bias, optimization variance, and the limitations and inductive biases of the transformer architecture. These results can be viewed as “scaling laws under a microscope”: the clean, steady decrease in average training loss (which is essentially the mean discrepancy between the model and ENTD) conceals a wide range of per-sample behaviors. Understanding the full breadth of these dynamics is critical for enhancing our understanding of LLMs as a whole.
In summary, we make the following contributions:
We study how the discrepancy between an LLM’s output and the empirical next-token distribution of the training data evolves over the course of training. Unlike classical scaling laws, which are equivalent to the mean discrepancy measured in cross-entropy, we study the distribution of discrepancies, revealing a wide range of behaviors from complete convergence to the ENTD to complete deviation.
We examine how different components of the model creation procedure contribute to this discrepancy. These include sampling error, variance in the optimization procedure, and structural properties of the transformer architecture.
We also analyze properties of the data itself which correlate with greater or lesser similarity between the model and the ENTD. In addition to well-established or intuitive effects, we surprisingly find that LLMs struggle to approximate the next-token distributions of frequent examples with a high empirical entropy.
2 Background↩︎
2.1 Standard ENTD↩︎
We focus on the base model obtained after the pretraining stage. Recent evidence shows that many of the final model’s capabilities are learned during pretraining, with common post-training stages improving the sample efficiency of the final model (i.e., the chance that the LLM produces a useful output on the first try, rather than requiring many attempts) rather than fundamentally expanding its reasoning capabilities [6]. Thus, for the sake of simplicity in this initial exploration, we restrict our attention to this first important step.
The pretraining phase updates the model weights to minimize the cross-entropy loss on a next-token prediction task. Specifically, let \(p_\theta(x_{k+1}|x_{1:k})\) be the model’s predicted probability that the token \(x_{k+1}\) is the next token after the string of \(k\) tokens \(x_{1:k}\), where the model parameters/weights are collected in \(\theta\). Pretraining updates \(\theta\) to approximately minimize \[\label{eq:32pretraining32loss} \mathbb{E}_{x_{1:k+1} \sim \mathcal{D}}[-\log p_\theta(x_{k+1} | x_{1:k})],\tag{1}\] where \(\mathcal{D}\) is the training dataset (consisting of a massive number of text documents, represented as strings of tokens) and inputs of variable lengths \(x_{1:k+1}\) are sampled from this corpus.
If \(p_\theta\) can represent arbitrary functions of its input, then we can exactly compute the global minimizer of 1 . It is given by the empirical next-token distribution \(p_{\textrm{ENTD}}\) of the training data, which is defined as \[\label{eq:32ENTD} p_{\textrm{ENTD}}(x_{k+1}|x_{1:k}) = \frac{\textrm{count}(x_{1:k+1}, \mathcal{D})}{\textrm{count}(x_{1:k}, \mathcal{D})},\tag{2}\] i.e., the empirical probability that \(x_{k+1}\) follows \(x_{1:k}\) in the training data. The empirical next-token distribution is, in some sense, the ultimate “stochastic parrot” [7], making predictions which are pure memorization of the training data.
2.2 Training-Aligned ENTD↩︎
During training, the model does not make a prediction on every context present in the training data. Rather, a common procedure (and the one employed by the models we experimented with) is for the data to be loaded in batches consisting of \(B\) examples of consecutive tokens of size equal to the model’s max context length \(C\). If \(x^{(b)}_{1:C}\) is one such example, then the loss is computed only on prefixes of the example, i.e., \[\sum_{b=1}^B\sum_{i=1}^C -\log p_\theta(x^{(b)}_{i+1} | x^{(b)}_{1:i}).\] Thus, the minimizer of the loss for the examples which are actually seen and evaluated during training is not the ENTD, but rather an analog (which we dub the training-aligned empirical next-token distribution or ENTD-T for short) which respects this batch-aligned structure. Specifically, let \(x_{0:N-1}\) denote the full tokenized corpus. Given an input \(\hat{x}_{1:k}\) of length \(k\) and a next token \(\hat{x}_{k+1}\), the ENTD-T is defined as \[p_{\textrm{ENTD-T}}(\hat{x}_{k+1}|\hat{x}_{1:k}) = \frac{\#\{0 \leq i \leq N / (C+1) \: | \: x_{(C+1)i: (C+1)i + k + 1} = \hat{x}_{1:k+1}\}}{\#\{0 \leq i \leq N / (C+1) \: | \: x_{(C+1)i: (C+1)i + k} = \hat{x}_{1:k}\}}.\] Equivalently, this is the fraction of training context-aligned prefixes starting with \(\hat{x}_{1:k}\) where the next token was \(\hat{x}_{k+1}\). (Note that training context alignment is a mod-\((C+1)\) quantity since the need for next-token “targets” means that each example of length \(C\) actually consumes \(C+1\) tokens. This is aligned with the training dataloader for the models we tested.) Similarly to how the ENTD was defined only for inputs which occur somewhere in the training corpus, the ENTD-T is defined only for batch-aligned prefixes, which is a subset of valid inputs to the ENTD.
2.3 Representing the ENTD and ENTD-T Distributions↩︎
Naively computing the ENTD and ENTD-T distributions requires a full-corpus scan to compute the necessary counts. For a realistic modern LLM training corpus, this is computationally infeasible. We make use of \(\infty\)-gram [8], a method which uses suffix arrays to efficiently count \(k\)-gram occurrences in large text corpora for arbitrary lengths \(k\). A simple modification of \(\infty\)-gram also allows us to compute the batch-aligned ENTD-T counts efficiently.
3 Experimental Setup↩︎
3.1 Models↩︎
We use the Pythia model suite [9]. The models cover a wide range of model scales, from 70M to 2.8B parameters,1 and are trained on a large, realistic corpus (the Pile [10]). While these are still far from industrial scales, they are larger and more comprehensive than previous related studies [11]. Importantly, Pythia contains model checkpoints throughout the training process, not just the final weights; and the dataloader is reproducible, allowing us to know not only the full set of examples seen during the entirety of training, but specifically which samples were seen at which points. This will allow us to examine some of the effects of the training itself in detail.
3.2 Evaluation Data↩︎
While there are many reasons why a transformer-based language model is preferred over the ENTD itself in practice, one of the critical ones for our setting is that the ENTD only makes predictions on inputs which are present in the training data. As such, we restrict ourselves to evaluation samples from the training set. While the ENTD can be extended to samples not contained in the training data via a backoff procedure (and other works, e.g., [11], take this approach), the backed-off distributions are not directly connected to the training loss in the same way that the ENTD is. Thus, we choose to sacrifice coverage of examples in order to retain a fully mathematically principled target.
If we simply took a uniformly random sample of examples used in training, the vast majority of samples used for our evaluations would occur only once in the dataset. This is because the Pythia models have a context window of length 2048 and a token vocabulary of order \(10^5\). For inputs of even moderate length (e.g., 20 or more tokens), the odds of this string occurring multiple times in the corpus are slim, and any reoccurrences will tend to be duplicates of a longer string. In such cases, the ENTD and ENTD-T will be a point mass on the single completion seen for this exact string. A uniformly random sample from training examples would contain a fraction of at most \(20/2048 \approx 1\%\) of examples outside of this regime, which does not give a good picture of the breadth of the model’s behavior.
As such, we take an approach similar to [12] and define two stratified samples from the data. Specifically, we define several occurrence and length bins for the sampled training strings, and then uniformly sample an equal number of examples from each (occurrence count, length) bin. In some cases, the total number of samples in the corpus belonging to a given bin was below the desired count; in this case, we simply take all of the samples in the bin. Due to computational constraints, we shard the full Pile corpus, compute the stratified samples per-shard, and then combine the samples across shards to get the final evaluation sample, removing duplicates that may have been chosen in multiple shards. We use this procedure to create one evaluation sample where individual examples can be pulled from anywhere in the corpus, and another evaluation sample where the examples are training context-aligned (and therefore valid for the ENTD-T). Statistics of the final samples are shown in Tables 3 and 4.
4 Large Models Exhibit Strong Positional Invariance↩︎
We first test the hypothesis which motivates the construction of the ENTD-T. Namely, the global minimizer of the loss that the model actually sees during training—where predictions are only made on batch-aligned prefixes—is the ENTD-T, not the ENTD. A natural conjecture is therefore that the model will approximately converge to the ENTD-T, rather than the ENTD.
Figure 2 shows the distribution of TV discrepancies for the model vs. the ENTD (left) and vs. the ENTD-T (right). To provide a fair comparison, these discrepancies are computed on the ENTD-T evaluation sample (which are valid inputs to both the ENTD-T and ENTD). Surprisingly, even though the ENTD-T is aligned with the actual training procedure, the LLM distributions have consistently lower TV with the ENTD, with lower mean and median TV, and visually lighter tails at the upper end of the TV range.
There are two factors at play here. First, there is a sample efficiency issue: namely, because the ENTD-T only estimates empirical counts at batch-aligned positions, it has on average 1/(context window + 1) as many samples to estimate the empirical next-token distribution as compared to the ENTD. In the case of the Pythia models we used, this is a factor of 2049 sample discrepancy, which is significant. Since batch-aligned prefixes are statistically identical to non-aligned samples, the increased sample size (and consequently reduced sampling error) of the ENTD vs. ENTD-T may outweigh the modest gains from batch alignment.
The second factor is that the conjectured batch alignment effect is much weaker than expected. Put another way, the transformer exhibits a strong positional invariance, even though it was not explicitly trained to do so. This can be seen via the following experiment. The ENTD-T is constructed via the empirical counts for strings which are in positions that are at a 0-offset with respect to the training-aligned context windows. We can construct similar empirical distributions by simply changing the offset. We computed the TV distributions for the model vs. the empirical distributions with an offset of 1 (i.e., all strings which start one token after the training-aligned contexts, which gives slight misalignment) and 1024 (i.e., all strings which start 1024 tokens after the training contexts; this is in the exact middle of the training contexts, or maximally misaligned).
The results are shown in Figure 3. Across different model sizes and throughout training, the TV distributions are nearly identical for different offsets. This indicates that the model has learned a strong positional invariance, in spite of the fact that the batch structure of training would not encourage this at first glance.
This is a surprising result. We therefore ran additional experiments to ensure that the positional invariance is real at these larger scales and not due to a bug in our empirical index construction. These results can be found in Appendix 12. In light of these results, all the remaining experiments will measure discrepancies between the LLM and the ENTD.
5 The Softmax Bottleneck↩︎
The final logits for a transformer are determined by dot products with \(d_{\mathrm{model}}\)-dimensional unembedding vectors, where \(d_{\mathrm{model}}\) is the dimension of the model’s residual stream. As such, it can be shown that for any collection of \(n\) inputs to the transformer, if the output logits for these examples are collected into a \(n \times |V|\) matrix (where \(|V|\) is the size of the token vocabulary; the matrix is formed by stacking the logit vectors in rows), then this matrix has rank at most \(d_{\mathrm{model}}\). This is known as the softmax bottleneck [13], [14].
One possible explanation for the deviations between the LLM and empirical next-token distributions is that natural text does not have a clean, low-dimensional structure, and indeed this hypothesis motivated [13] to find ways around it. Nevertheless, the standard softmax is still the most popular choice for modern transformers, and this can perhaps explain more of the model/empirical discrepancy.
To test this, we can find an upper bound on the bias introduced by the softmax bottleneck by finding a fully unconstrained rank-\(d_{\mathrm{model}}\) logit approximation to the matrix of empirical distributions for the evaluation sample. Specifically, let \(x^{(1)},\ldots,x^{(n)}\) be the set of evaluation strings and let \(p_{\mathrm{ENTD}}(x^{(i)})\) be the empirical NTD for \(x^{(i)}\), viewed as a vector in \(\mathbb{R}^{|V|}\), and let \(P \in \mathbb{R}^{n\times |V|}\) be the matrix with rows \(p_{\mathrm{ENTD}}(x^{(i)})^\top\). Let \(U\in \mathbb{R}^{n\times d}\) and \(W\in \mathbb{R}^{|V|\times d}\). Here, the rows of \(U\) can be thought of as the last layer token embeddings for the evaluation samples (before the unembedding), and \(W\) stands in for the unembedding matrix. We then train \(U\) and \(W\) so that they minimize the cross-entropy loss: \[\mathcal{L}(U, W) = \mathrm{CrossEntropy}(\mathrm{Softmax}(U W^\top), \: P),\] where the softmax is taken row-wise and the rows of \(P\) are used as soft targets for the cross entropy. Essentially, this allows us to ask: if the only constraint on the transformer’s output distribution were due to the softmax bottleneck, and there were no constraints on the pre-softmax representations, how well could it approximate the empirical distribution?
Surprisingly, we find that [13]’s hypothesis does not hold: the empirical distributions of the evaluation samples are in fact highly amenable to a low-rank logit approximation, far below the deviation obtained by even the largest models at the end of training. This can be seen in Figure 4, where the TV discrepancies for the unconstrained low-rank logit approximation are far lower than they are even for the largest model at the end of training. Thus, the natural empirical distributions (at least on our evaluation sample) are amenable to a low-rank approximation, but the transformer does not take full advantage of this fact!
There is a caveat to this result, namely, that the best low-rank approximation on the evaluation sample (roughly 84K examples) is not an apples-to-apples comparison with how the transformer was trained, where it makes a total of \((\textrm{context length} = 2048) \times (\textrm{batch size} = 1024) \times (\textrm{training steps} = 143,000) \approx 300B\) predictions. The comparison may not be as unreasonable as it initially appears, however, because the number of samples used to build the full set of empirical NTDs on the evaluation sample (i.e., the sum of the number of corpus occurrences for each evaluation sample) is roughly 254B. Nevertheless, whether this low-rank structure persists as the number of evaluation samples grows is an open question.
6 Optimization Randomness: Inter-model Discrepancy↩︎
Training an LLM is an inherently stochastic procedure, with randomness arising from the ordering of the training data, random weight initialization, etc. It is therefore natural to ask, how much of the deviation between the LLM’s predictions and the ENTD can be attributed to these random fluctuations, i.e., due to variance rather than a systemic bias?
To test this, we made use of the PolyPythias [15] extension of the Pythia models. PolyPythias supplies model checkpoints from different random training runs using the same model architectures on the same overall corpus, but varying training order and random initializations. This allows us to directly measure the effects of randomness from the training process. The results are shown in Figure 5, where we observe a number of interesting behaviors.
First, the shape of the model-empirical vs. model-model discrepancies are noticeably different: the model-empirical discrepancies tend to be bimodal at the extremes, whereas the model-model discrepancies generally have a modal outcome of moderate TV with a second mode of increasing size at very low TV values. This is consistent with the following picture: Given the training set and a fixed model size, the model learns an “expected transformer distribution” plus noise due to the random optimization. The shape of these discrepancy distributions is consistent with a fairly continuous/“well-behaved” error distribution.
Another trend is that larger models converge to a more stable target, and the randomness from optimization has less of an impact. Indeed, for the 70M model, the average effect of the transformer bias (measured by the average discrepancy with the empirical distribution) is smaller than the model distribution variance (measured by average model-model discrepancy). But for the larger two models, the inter-model gap is smaller than the average gap between the models and the ENTD.
7 Stratified Analysis↩︎
7.1 Different Lengths Are Learned Differently↩︎
Figure 6 shows TV distributions stratified by different input lengths (1, 4, 16, and 64) over the course of training. We observe a progressive “stretching” of the TV distributions from shorter to longer prefixes: the modal TV for shorter prefixes is moderate, whereas longer prefixes concentrate more and more at the extremes. One possible explanation for this is the relative importance of these terms to the overall training loss. Unigram discrepancy decreases moderately throughout training for larger models, but as these terms contribute fairly little to the overall loss, the model does not prioritize learning them exactly. In contrast, longer inputs contribute more to the loss, and the model expends more capacity to make many of these have very low discrepancy; but because there may be many challenging (or even “bad”) long sequences, there is also a thicker tail of longer examples which the model does not learn at all.
7.2 TV Tails↩︎
For all but the smallest models, the modal outcome (widest part of the violin plot) is for the LLM to learn the ENTD almost perfectly (TV \(\approx 0\)). However, across model scales, there is a noticeable “bump” of evaluation samples with close to maximum discrepancy (TV = 1) with the ENTD, even at the end of training (refer to Figure 4). In this section, we examine whether there are any consistent properties which distinguish such samples from the rest of the data, which may explain why the model fails to learn them.
In Table 1, we examine some aggregate statistics of evaluation samples in the most- and least-converged tails of the TV distributions, and compare them to the statistics of the overall evaluation sample. The general trends are consistent across model scales. For two of the statistics—input length and ENTD entropy—the high- and low-TV tails tend to lie at the extremes of the evaluation sample: both tails contain longer than average inputs and lower than average ENTD entropy. On the other hand, input frequency in the corpus and LLM entropy have more monotone relationships with TV, with median corpus frequency clearly decreasing with TV and entropy of the LLM’s NTD clearly increasing with TV.
| Pythia-70M | Pythia-410M | Pythia-2.8B | |||||||
| Lo TV | All | Hi TV | Lo TV | All | Hi TV | Lo TV | All | Hi TV | |
| 29.88 | 14.65 | 20.13 | 32.82 | 14.65 | 17.31 | 33.53 | 14.65 | 16.52 | |
| 16 | 8 | 12 | 32 | 8 | 12 | 32 | 8 | 12 | |
| 2.74e5 | 3.03e6 | 5,457 | 2.5e5 | 3.03e6 | 1,264 | 2.52e5 | 3.03e6 | 507 | |
| 21,354 | 3,573 | 4 | 22,391 | 3,573 | 2 | 22,380 | 3,573 | 2 | |
| 0.13 | 3.80 | 3.88 | 8.61e-3 | 3.14 | 4.81 | 2.00e-3 | 2.75 | 4.49 | |
| 0.06 | 4.12 | 3.98 | 5.07e-3 | 3.24 | 4.94 | 7.90e-4 | 2.64 | 4.61 | |
| 0.03 | 1.59 | 0.09 | 4.67e-3 | 1.59 | 0.14 | 3.15e-3 | 1.59 | 0.20 | |
| 4.04e-3 | 0.24 | 0.00 | 3.69e-3 | 0.24 | 0.00 | 2.99e-3 | 0.24 | 0.00 | |
4pt
We then manually inspected examples from the min- and max-TV tails. Specific examples illustrating each of the categories discussed below can be found in Table 2. The “most converged” (low final TV) examples have very low entropy for the next-token distribution and a high count for the input prefix, meaning that the evaluation sample appears many times in the corpus (median occurrence around 20k) with the exact same completion almost all of the time. Upon inspection, we find that these examples frequently correspond to common headings in code (e.g. standard LaTeX preambles, HTML headers, ...), licensing agreements, or enumerations. By and large, these strings can be viewed as deterministic templates.
Examples in the high-TV tail tended to fall into several interpretable categories. One category was rare inputs, i.e., example strings which occur once or twice in the training corpus and therefore have a very peaked ENTD. Many of these examples are still reasonable natural text, and discrepancies on these inputs are likely due to noise in the estimation of the ENTD, i.e., a finite sample effect where the smoothing from the inductive bias of the transformer actually results in a better estimation of the next-token distribution than the empirical counts.
A second category of high-deviation examples contained frequent prefixes (count above \(10^4\)) but with high empirical entropy. These cases include genuinely uncertain mid-sentence continuations or variable web boilerplate. Here, it is worth emphasizing the difference in what we are measuring by considering the TV distance to the ENTD, vs. measuring the model’s cross-entropy loss on a given example. If we were measuring via loss, then it would be unsurprising that the model performs “worse” on high-entropy examples, because the minimum achievable loss is higher when the ground-truth distribution has higher entropy. However, the minimum TV distance is always 0, regardless of the entropy of the ENTD. This means that not only does the model have greater difficulty making predictions given higher true uncertainty, it also has greater difficulty even approximating the true uncertain distribution! In fact, among the natural statistics that we checked and for inputs with high counts (at least \(10^4\) occurrences in the corpus), empirical NTD entropy had the highest Spearman correlation with final TV (\(\rho = 0.80\)). While the softmax bottleneck is not enough to explain this, we posit that this fact is likely due to some other aspect of the transformer architecture’s inductive bias which forces the model to learn a “narrower” distribution than the real text in cases of genuinely high entropy. This fact may partially explain the observed homogeneity/diversity collapse of LLM-generated text [16].
Conversely, there was a category of frequent inputs with very low ENTD entropy in the high-TV tail. These are templated or spam-like strings (often occurring many times in a single document) whose fixed continuation the model declines to commit to, spreading its mass instead (model entropy of \(4.5-6.5\) against an ENTD entropy of \(\approx 0\)).
| Best fit: deterministic templated / code / markup | ||||
|---|---|---|---|---|
| Count | TV | \(H_{\mathrm{ENTD}}\) | \(H_{\mathrm{model}}\) | Input Text |
| 22,310 | 0.00 | 0.01 | 0.01 | )$\documentclass[12pt]{minimal} |
| 507 | 0.00 | 0.00 | 0.01 | ://java.sun.com/j2se/1.5.0/docs/api/java/lang/Object... |
| 361 | 0.00 | 0.00 | 0.00 | : {r: 0, g: 0, b: 0, a: 1}\n m_OnHover:\n m_Ba... |
| 3,215 | 0.00 | 0.00 | 0.00 | . \n 451. \n 452. \n 453. \n 454. \n 455. \n 456. \n ... |
| Worst fit: frequent (count \(\ge 10^4\)), high-entropy ENTD, model commits to a small number of completions | ||||
| Count | TV | \(H_{\mathrm{ENTD}}\) | \(H_{\mathrm{model}}\) | Input Text |
| 29,729 | 0.99 | 6.08 | 0.08 | now live! Get 30% to 60% off our most popular products today! |
| 20,246 | 0.99 | 6.03 | 0.10 | ...or intrusion, then please contact the editor here. If |
| 38,340 | 0.97 | 5.91 | 0.16 | ...this page to help users provide their email addresses. You |
| 13,975 | 0.95 | 6.04 | 1.50 | articles, please go to: www.TorontoStarReprints.com |
| Worst fit: frequent (count \(\ge 10^4\)), near-deterministic ENTD, templated/spam text, model does not commit | ||||
| Count | TV | \(H_{\mathrm{ENTD}}\) | \(H_{\mathrm{model}}\) | Input Text |
| 16,869 | 1.00 | 0.01 | 4.53 | LIKE A FEVER!!!! WE NEED YOU GUYS!! COME PLEASE! |
| 27,284 | 1.00 | 0.01 | 6.54 | spoke to Pat Joseph about his memoirs in 2014. “Jacobs’s |
| 20,772 | 1.00 | 0.01 | 6.00 | will not be shared between programs within our |
| 35,448 | 1.00 | 0.00 | 4.67 | link to the entire object, |
| Worst fit: rare (count \(\le 2\)), tiny-sample empirical | ||||
| Count | TV | \(H_{\mathrm{ENTD}}\) | \(H_{\mathrm{model}}\) | Input Text |
| 2 | 1.00 | 0.00 | 3.35 | ,’ the tower’s observatory deck, At.mosphere present... |
| 1 | 1.00 | 0.00 | 3.21 | a silly step can prevent a silly mistake |
| 2 | 1.00 | 0.00 | 5.00 | company literally has rolls of the stuff; its engineers plan to |
| 2 | 1.00 | 0.00 | 5.85 | ...of the state capital Austin. Caldwell County Sheriff Daniel |
8pt
In addition to simply pattern matching on example strings from the tails, we also analyzed the stability of the tails over training to see whether or not these are “structurally sound” or just artifacts of a particular snapshot. The results are shown in Figure 7, revealing that there is a combination of “structurally challenging" examples which are consistently present, but a number of other examples which move in and out of the tail as the model learns. We also conducted a sensitivity analysis for the tail cutoff threshold and find qualitatively similar results. The comparison plots can be found in Appendix 13.
8 Evolution of Individual NTDs↩︎
While the distribution plots in Figure 2 give more information than a scaling law, they are still somewhat coarse. For instance, they cannot differentiate between the setting where the discrepancy for all examples slowly and steadily decreases relatively uniformly across examples (which would indicate that the model learns small amounts of generalizable information from each sample, and learns how to correctly respond on a given sample by aggregating this information across training); the setting where there are sudden decreases in a few samples at a time (which would indicate that the model learns by essentially memorizing one sample at a time); or some combination of the two or other more complicated patterns. In this section, we examine the change in TV for individual examples over the course of training to shed light on this picture.
Figure 8: Individual TV trajectories for a small (70M, left) and large (2.8B, right) model. The trajectories exhibit much greater oscillation in the smaller model, aligning with the findings of other recent work [17]. For the larger model, examples at the extreme ends of the discrepancy distribution (TV\(\approx\) 0 or 1) tended to have more stable trajectories, while more oscillation is observed in the middle of the TV range..
Results stratified by evaluation example count and its TV at the final training step are shown in Figure 8 for a small model (70M, left) and a large model (2.8B, right). As expected for such a complicated system, the individual trajectories are a combination of many distinct behaviors, but we observe a few general trends.
First, the TV trajectories for the small model are in general much more oscillatory than they are for the large model. This accords with the recent work of [17], who showed that gradients of distinct samples exhibit much greater interference in smaller models.
For the large model, there are two categories of points which have exceptionally oscillatory TV trajectories: inputs with low counts but which nevertheless have low to medium final TV; and inputs with higher counts but with high final TV. In the former case, these low-count and low/medium final TV examples tend to have a point-mass ENTD (i.e., a single completion in the corpus) but are nevertheless predictable, for instance, code stack traces or enumeration formats. This has two consequences. First, because the ENTD is a point mass, the TV curve is just \(1-p_t(c)\), where \(p_t(c)\) is how much mass the model places on the single completion \(c\) at step \(t\). Since this is only dependent on a single token probability, there is no averaging effect over tokens, making it easier for variation to occur. Second, because these exact sequences are rare, it is unlikely that the model can memorize them exactly; but because they are predictable, generalizable information from other sequences keeps \(p_t(c)\) fluctuating in a medium range. This is in contrast to the low-count but high-TV examples, which have low variability: the model does not memorize these either, but their single observed completion is not strongly predictable from related text. Even when the input is perfectly natural, the observed continuation is only one of many plausible ones, so the model’s mass on it stays near zero, where fluctuations are necessarily negligible and the TV remains pinned close to 1.
One plausible explanation for oscillations among the high-count, high final TV examples may be related to the discussion of the TV tails from the previous section, where we found that many high-count but high-TV examples corresponded to examples with empirically high next-token entropy. It seems in these cases, potentially due to architecture constraints, the model is forced to choose a “narrower” (i.e., lower entropy) approximation to the “flatter” (i.e., higher entropy) empirical distribution. For such examples, the narrow approximation that any given example pushes the transformer towards may vary significantly between the many examples, thus leading to an unstable next-token distribution for this input over the course of training and consequently generating an oscillatory TV curve. Additional trajectories for different input lengths can be found in Appendix 14.
9 Discussion↩︎
In this work, we studied the extent to which LLM predictions can be explained by the empirical next-token distribution of the data, which is the global minimizer of the training loss and therefore ostensibly what training should push the model towards. While this does indeed occur on average over the training samples—a fact which is essentially equivalent to the existence of neural scaling laws—a more fine-grained look at the discrepancies of individual samples reveals that while the empirical distribution models many inputs well, there is also a long tail of examples which it does not. We identified several aspects of the LLM stack which could contribute to discrepancies between the LLM and the ENTD, including training context window alignment, randomness in the optimization procedure, and representational limitations of the transformer architecture.
Surprisingly, we found that transformers exhibit a strong positional invariance, negating the batch-alignment effect. We also found that the empirical distribution of natural data seems to be amenable to low-rank logit approximation, and the learned transformer exhibits stronger smoothing of the empirical distribution than what could be attributed solely to the softmax bottleneck effect. The randomness of the optimization procedure also seems to have some role in the deviation, but it becomes less pronounced as the model grows in scale.
9.0.0.1 Limitations and Future Work
These models are below frontier scale, and there is always a risk that our conclusions will not hold or will be qualitatively dissimilar at the frontier. Nevertheless, we have some confidence that our conclusions are relevant given that they hold across nearly two orders of magnitude model parameters. Extending our results to the larger models in the Pythia suite, as well as to other open-source models (e.g., the Olmo family [18]), is ongoing.
Another key limitation is that our strict ENTD analysis is only applicable to sequences belonging to the training set. The \(\infty\)-gram model can be extended to sequences not seen during training (indeed, the authors provide this functionality in the official release via backoffs to substrings of the input which are contained in the training set, and other works such as [11] use these backoffs to analyze strings not belonging to the training set). However, as the backoff procedure has no principled connection to the model’s training, we chose to limit our scope to training examples only, rather than weaken the connection between the LLM and the object we are using to explain it. Nevertheless, explaining how the model produces its outputs on unseen inputs is clearly a desirable goal for future work.
An orthogonal issue is that our analysis focuses on the model’s predicted distribution for a single next token. However, LLMs’ usefulness comes not only from predicting one single token into the future, but rather from chaining these distributions together to generate long sequences of new tokens. As the NTD is the atomic component of these generations, and this much simpler object is still far from completely understood, we chose to focus on the NTD in this work. However, using these atomic results to say something meaningful about the long generations of an LLM seems challenging in its own right.
10 Related Work↩︎
There have been numerous attempts to draw a more precise link between LLMs and their training data. Some works have attempted to trace model behavior to individual datapoints using approximate methods such as influence functions [19]–[21] or by finding training sequences which match part of the output [22]. In contrast, our work considers distributional properties of the data, rather than individual points. Previous work suggests that some model behaviors may not be attributable to individual points, but rather aggregate statistics of many data (e.g. [23]).
Other works have tested how training data statistics impact a model’s performance on specific tasks or behaviors [24]–[29] or on heuristic measures of generalization [30]. [31] studied the connection between training data statistics and the internals/representations learned by a transformer (rather than the final output distribution), focusing on early stages in training.
The work most similar in spirit to ours is [11], which also explored the agreement between LLM predictions and \(n\)-gram-based models (such as the ENTD). However, there are several important differences. Most of their experiments use a 160M-parameter transformer on TinyStories, which contains 480M tokens. The largest tested model has 1.4B parameters on a 4.4B-token training set. They include experiments which measure the TV distance between the model output and the ENTD for 7-gram sub-sequences of the input and for the fully trained model. In contrast, we explore a much wider range of model sizes/compute scales, as well as a wider range of \(n\)-grams, and more in-depth discussion of where the discrepancy between the model and the ENTD may emerge from.
In summary, with a few exceptions, most of these studies have proceeded via heuristic measurements of the data or by interrogating the models with specific, human-intelligible tasks. Our approach attempts to take the bare mathematical basis of LLM pretraining—training the model to perform next-token prediction via the cross-entropy loss—at face value to determine how far these quantitative facts can take us.
11 Evaluation Sample Statistics↩︎
| \(k\) | \([10^{0},\, 10^{1})\) | \([10^{1},\, 10^{2})\) | \([10^{2},\, 10^{3})\) | \([10^{3},\, 10^{4})\) | \([10^{4},\, 10^{5})\) | \([10^{5},\, 10^{6})\) | \([10^{6},\, \infty)\) | Total |
|---|---|---|---|---|---|---|---|---|
| 1 | 2 | 0 | 18 | 75 | 503 | 2,102 | 3,761 | 6,461 |
| 2 | 560 | 1,247 | 1,887 | 2,050 | 2,142 | 2,169 | 2,773 | 12,828 |
| 4 | 1,405 | 726 | 1,789 | 2,031 | 2,051 | 2,038 | 1,125 | 11,165 |
| 6 | 1,864 | 484 | 1,671 | 1,983 | 2,045 | 1,761 | 329 | 10,137 |
| 8 | 1,985 | 485 | 1,589 | 1,953 | 2,036 | 1,411 | 203 | 9,662 |
| 12 | 2,040 | 489 | 1,549 | 1,942 | 1,986 | 834 | 110 | 8,950 |
| 16 | 2,064 | 482 | 1,547 | 1,958 | 1,935 | 532 | 77 | 8,595 |
| 32 | 2,085 | 365 | 1,621 | 1,965 | 1,912 | 273 | 52 | 8,273 |
| 64 | 2,080 | 290 | 1,675 | 1,961 | 1,736 | 136 | 14 | 7,892 |
| Total | 14,085 | 4,568 | 13,346 | 15,918 | 16,346 | 11,256 | 8,444 | 83,963 |
4pt
| \(k\) | \([1,\, 3)\) | \([3,\, 10)\) | \([10,\, 30)\) | \([30,\, 100)\) | \([100,\, 300)\) | \([300,\, 1,000)\) | \([1,000,\, \infty)\) | Total |
|---|---|---|---|---|---|---|---|---|
| 1 | 5 | 70 | 197 | 772 | 2,407 | 2,312 | 4,542 | 10,305 |
| 2 | 764 | 659 | 743 | 1,355 | 1,734 | 1,931 | 3,604 | 10,790 |
| 4 | 1,774 | 313 | 580 | 1,152 | 1,492 | 1,585 | 854 | 7,750 |
| 6 | 1,997 | 188 | 601 | 1,023 | 1,269 | 737 | 192 | 6,007 |
| 8 | 2,065 | 116 | 509 | 1,081 | 1,009 | 369 | 94 | 5,243 |
| 12 | 2,077 | 131 | 583 | 985 | 490 | 115 | 67 | 4,448 |
| 16 | 2,086 | 157 | 611 | 899 | 261 | 89 | 63 | 4,166 |
| 32 | 2,097 | 159 | 714 | 639 | 130 | 45 | 44 | 3,828 |
| 64 | 2,096 | 205 | 703 | 435 | 60 | 5 | 11 | 3,515 |
| Total | 14,961 | 1,998 | 5,241 | 8,341 | 8,852 | 7,188 | 9,471 | 56,052 |
4pt
12 Confirmation of the Null Batch Alignment Result↩︎
Figure 9 shows the discrepancies when the model is trained for 1000 steps on two batches (500 steps on each batch). In this case, we get the expected result, namely, the model has overfit only to the batch-aligned sequences. Even with a single token offset, the gap widens significantly, and with an offset of 1024 tokens (exactly half the context window size) the median gap is close to maximal.
We also constructed the empirical next-token distributions via two methods: the infini-gram-based approach used for the rest of the experiments, as well as a “brute force” method, where we simply kept a hash table of next tokens following inputs seen during training. Notably, the brute-force index was constructed directly from the logged and independently verified training data order, without a suffix array, allowing us to directly verify that the sequences being fed to the model (or with the appropriate offset) are in fact the same as the ones being used to construct the index. In addition to the perfectly overlapping TV distributions in Figure 9, we also confirmed that the empirical indices constructed via these two approaches agreed exactly on all entries.
Figure 9: Batch-aligned vs. offset discrepancies after extended training on only two batches, and with the empirical indices computed via the infini-gram approach used in the other experiments (left) as well as a brute-force (but easily verifiably error-free) method for constructing the empirical indices (right). When the model is trained extensively on only two batches, the expected batch-aligned discrepancy gap emerges. The infini-gram-based indices also align perfectly with the brute-force method..
13 Tail Stability Sensitivity Analysis↩︎
The qualitative tail stability picture is fairly robust to the chosen cutoff. Refer to the plots in Figure 10. The shape of the curves is stable across thresholds.
14 Additional Trajectories↩︎
Figure 11: Individual TV trajectories for a small (left) and large (right) model, input length \(k=1-2\)..
Figure 12: Individual TV trajectories for a small (left) and large (right) model, input length \(k=4-8\)..
References↩︎
Pythia also contains several smaller models (as low as 14M parameters) and two larger scales, with 6.9B and 12B parameters. Due to the computational burden of running inference on these models, our efforts to extend our analysis to these larger models are still ongoing. We exclude some of the smaller models from the analysis since they clutter the figures and do not add much to our ultimate goal, which is to understand frontier models at scale.↩︎