Understanding Transformers via N-gram Statistics

Summary

The paper studies how transformer-based large language models (LLMs) use context when predicting the next token by approximating these predictions with N-gram-based statistical rules. The authors propose a method to describe transformer predictions using simple N-gram rules and study how well these rules approximate the predictions. They claim their key findings include a new method to detect overfitting without a holdout set, insights into the progression from simple to complex statistical rules during training, a criterion for when transformer predictions align with N-gram rules, and an understanding of how well transformers can be approximated by increasingly complex N-gram rulesets. The experiments (primarily carried out on the TinyStories dataset and validated on Wikipedia data) seem provide possible insights into the statistical nature of LLM behavior.

Strengths

- The proposed method for detecting overfitting without needing a holdout set is new and interesting for optimization applications. - The paper includes visualizations and concrete examples. - Although focused on specific datasets, the methods and insights could potentially be scaled to other domains.

Weaknesses

- The study provides descriptive approximations without offering explanations into why certain rules work, not going deep enough in understanding the transformer behavior. - The paper doesn't properly explore how fine-tuning on different datasets might affect the effectiveness of the proposed methods. - The study mainly considers context lengths of up to seven tokens, which may not fully capture the long-range dependencies that transformers are capable of handling.

Questions

- How do you plan to address the computational complexity associated with selecting the optimal N-gram rule for each context? - Do you have any plans to transition from descriptive approximations to more explanatory models that can predict why certain rules work? - Have you considered additional evaluation metrics that might capture other aspects of model performance and approximation quality?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

- The computational overhead required to apply and test the N-gram rule sets at inference time could be significant. - The proposed methods might not easily adapt to new or evolving datasets without substantial re-calibration or re-training, limiting their long-term applicability in dynamic environments - While N-gram rules provide a way to approximate transformer behavior, the results might not always be easily interpretable by humans.

Summary

The authors use n-gram statistics and regular expression templates to study how well they describe the predictions of Transformers-based models. They craft rules that vary in context length and/or the number of marginalised context variables to predict the next word. They use their framework to study overfitting (and memorisation) of increasingly complex template rules. They discover a counter-intuitive phenomenon: longer n-gram rules can be less predictive of a model's behaviour than shorter; an insight they name model-variance trade-off.

Strengths

The idea of defining templates that apply both to the data distribution and the model's generation is good and sound. I appreciate the idea of having a mathematical methodology as it states clearly what is the objective and how rules are defined (despite a few things I do not fully understand and I reckon may be wrong; see next sections). Findings are noticeable: 7-gram rules better approximate a model's behaviour than 8-grams (but please refer to the next section for. a caveat on this finding). The model-variance explanation of why some rules are better than others is intriguing (but see weaknesses).

Weaknesses

I am concerned about the evaluation, in particular Eq. 10. You measure how closely the dataset adheres to the template rules and compare it to your model’s behaviour. With such a small dataset, the fact 7-gram rules better approximate the distribution of data may be caused by a considerable drop in the overall frequency (and thus, variance) of 8-grams (vs. 7-grams). See a question about that. In summary, my biggest concern is whether the best-rule is identified by the model-variance trade-off or the intrinsic larger variance of 8- vs. 7-gram rules (i.e., 8-grams are more difficult to predict because you don't have enough data, so the 7-gram is optimal). With larger datasets and different models, results can vary. In fact, another major concern is that you only have two datasets and one model to prove your hypothesis. Furthermore, it is not clear if you take inspiration from Chinchilla's architecture (it seems to me you do not use a pre-trained model). Otherwise, it would be nice to test this hypothesis with smaller, purely next-word-prediction trained models (even gpt-2 or Pythia would be good). Other concerns (in descending order of importance): Figure 2: low R-squared may mean the regression failed and not lack of "correlation". Furthermore, you mention 7-grams but some plots are for 8grams (see y-axis of figs (a,b)). While I understand why, the plot is named 7-grams, so I was a bit confused at the beginning. It's better to show less in this case or make it clear also in the caption that the top-plots refer to 8grams. (minor) Equation 10 seems to me wrong: p and q are distributions, you can’t subtract them (unless i stands for p(t=v_i |C), but that has to be clarified). If not, probably you want to say something like Divergence(p || q). On a side note, I think this is one of the cases where an asymmetric rule makes sense. You want to measure the amount of “information” a model lacks to turn p (its prediction) into q (the ground truth). (minor) I am a bit confused by the notation in Section 4 (see questions, though the general idea is clear enough).

Questions

In order of importance: (major) Can you give a clear definition of model-variance in terms of predictive capabilities of the model and when compared to an n-gram rule? (major) Have you checked if there is a large drop in the variance of 7 vs. 8 grams in your dataset? I suspect that causes the 7-gram rules to be optimal and not the model-variance trade-off. In other words, the variance of 8-grams in your dataset is so large (because you don't have enough data), that the architecture you use cannot approximate it. (minor) What if the optimal rule you find only correlates with a model's behaviour (e.g., due to the size of the dataset)? Can you give some evidence that a model is actually using n-gram rules to predict the next token? (minor) Eq. 10: what is i? Is p(t=v_i | C), with v_i the i-th token in the vocabulary, or the i-th run of the model? (minor) Line 111, the expression +-*+ should produce C_{-5}*C_{-1} and not C_{-4}*C_{-1}, or am I missing something in your notation? Is it for the left padding you mention earlier? (minor) Eq. 6 (left hand side): why don’t you have C_{-4} at the denominator?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

See previous points.

Summary

The authors use 160M parameter models trained on the TinyStories dataset (artificially generated short stories, 480M tokens, made up of "vocabulary consisting of about 1500 basic words") to study transformers by comparing their predictions to the predictions of n-gram models. This leads to a few observations [quoted from the paper here]: 1. (Approximation-Variance Tradeoff) We observe an “approximation-variance tradeoff", which roughly states that next token LLM predictions that have low variance (across different training runs) tend to be well-approximated by N-gram rules. (Section 5) 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. (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. (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. 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)

Strengths

LMs are everywhere but we don't really know how they work and how the learn. Figuring this out might lead to stronger, more efficient models. I usually don't like interpretability research because most methods employed just don't make sense- but I *love* the approach taken by this paper of using n-gram models to analyze transformer models. These results seem straightforward but I'm pretty sure that this paper is the first to ever present them.

Weaknesses

There's a huge weakness in this paper- they analyze extremely small (160M param) models trained on a total toy dataset (TinyStories). TinyStories is a very small dataset made up of short stories that were generated by an LM. The vocabulary in this dataset is limited to 1.5k basic words. All of these things together mean that none of these results might transfer towards bigger, realistic LMs like LLaMA 3 70B. If this paper would have been written with LLaMA 3 70B as the model being used I would have argued for a strong accept, but because the models used here are so tiny I'm less excited about this contribution. Second thing- the authors put out a lot of what seems like 'intermediate analysis'- observations they noticed, but they don't explain what the conclusion from that observation is or why it's interesting. For example fig 2. Third- why is the 'Overfitting Criterion' interesting? Is anyone overfitting their LMs? Would anyone need this in practice?

Questions

1. is the Approximation-Variance Tradeoff true also for predictions that have low entropy?

Rating

7

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

n/a

Thank you for your response.

I agree that N-gram models are inherently simple and interpretable due to their reliance on frequency-based statistics; however, my concern is that while individual N-grams are indeed easy to understand, the challenge arises when these N-grams are aggregated and interact within a complex model like a transformer. Anyway, I appreciate the time you spent writing this respose.