Can Transformers Learn n -gram Language Models? (2024)

Abstract: Much theoretical work has described the ability of transformers to represent formal languages. However, linking theoretical results to empirical performance is not straightforward due to the complex interplay between the architecture, the learning algorithm, and training data. To test whether theoretical lower bounds imply learnability of formal languages, we turn to recent work relating transformers to n -gram language models (LMs). We study transformers’ ability to learn random n -gram LMs of two kinds: ones with arbitrary next-symbol probabilities and ones where those are defined with shared parameters. We find that classic estimation techniques for n -gram LMs such as add- \lambda smoothing outperform transformers on the former, while transformers perform better on the latter, outperforming methods specifically designed to learn n -gram LMs.

## introduction a large body of work has investigated the ability of transformers (vaswani et al., 2017) to represent formal languages (strobl et al., 2023). such results tell us what languages transformers can represent, but not how well they can learn them from data. existing work has tested the learning abilities of transformers as classifiers mapping strings to language membership decisions (e.g., bhattamishra et al., 2020; delétang et al., 2023). however, language models (lms) are not classifiers of strings but rather distributions over them. in this light, a recent line of work has argued for a more direct evaluation of lms as distributions (svete and cotterell, 2023; nowak et al., 2023; svete et al., 2024a,b; borenstein et al., 2024; nowak et al., 2024). in the pursuit to understand transformers as probability distributions, n-gram lms have recently emerged as a useful playground for analyzing transformers' interpretability (liu et al., 2024; voita et al., 2024), learning behavior (edelman et al., 2024; chen et al., 2024), in-context learning (olsson et al., 2022; arora et al., 2024; akyürek et al., 2024; nguyen, 2024), and representational capacity (svete and cotterell, 2024). svete and cotterell (2024), in particular, show that transformers can encode any n-gram lm, lower-bounding their representational capacity and suggesting that encoding n-gram statistics is indeed a feasible mechanism for transformers to process language. however, being able to represent n-gram statistics is insufficient to rely on them in practice. for that, learning them is necessary. this motivates the question: given that transformers can represent n-gram lms, how good are they at learning them? answering this brings us closer to a practical understanding of transformers, showing whether n-gram statistics are practically relevant for their behavior and illuminating what theoretical results on neural architectures imply about their ability to learn formal languages. we study how well transformers can learn ngram lms. on the one hand, we find that they generalize well on n-gram lms whose next-symbol probabilities depend on a linear combination of features of the n-gram symbols (representationbased n-gram lms), outperforming baselines such as count-based techniques and models with hand-crafted features. on the other, we find that transformers perform worse than count-based techniques on n-gram lms with arbitrary next-symbol probabilities. this illuminates transformers' inductive biases: it shows that they struggle to approximate even simple probability distributions whose next-symbol probabilities are arbitrary while being good at learning representation-based lms. ## preliminaries we begin by introducing some preliminaries. an alphabet σ is a finite, non-empty set of symbols. the kleene closure σ ˚of the alphabet σ is the set of all strings of the symbols in σ. the length of the string y " y 1 . . . y t p σ ˚, denoted by |y| " t , is the number of symbols it contains. we will use y j i def " y i ¨¨¨y j to denote substrings between positions i and j inclusive. a language model p is a probability distribution over σ ˚. two lms p and q are weakly equivalent if p pyq " q pyq for all y p σ ˚. many lms define p pyq autoregressively: assumption 2.1. the n-gram assumption states that the conditional probability of the symbol y t given y ăt only depends on n ´1 previous symbols: we will refer to n as the order of the n-gram lm and y t´1 t´n`1 as the history of y t . we will denote a length-pn ´1q history as y n whenever its position t in the string is unimportant. additionally, we assume that the histories are padded with the beginning-of-string symbol bos r σ and denote σ def " σ y tbosu for convenience. representation-based n-gram lms. the nextsymbol probabilities p n py | y n q of n-gram lms can be arbitrary. this requires storing op|σ| n´1 q parameters and does not capture the intuitive notion that similar n-grams define similar nextsymbol probabilities. this can be addressed by representation-based n-gram lms, which define p n py | y n q as a function of a representation of y n . definition 2.1. an n-gram lm p n is representation-based if there exist a representation function h : σ n´1 ñ r d and an output matrix e p |σ| ˆd for some d p n such that a transformer lm computes the p py t | y ăt q using the self-attention mechanism, which builds the representation by attending to different symbols in the preceding string (vaswani et al., 2017; radford et al., 2019). @xcite this can intuitively be connected to the way n-gram lms compute the next-symbol probabilities by attending to the last n ´1 symbols. transformers can represent n-gram lms. svete and cotterell (2024, thms. 3.1 and 3.2) show that, for any n-gram lm, there exists a weakly equivalent transformer lm. they describe intuitive mechanisms in which n ´1 heads (thm. 3.1) or n ´1 layers (thm. 3.2) attend to the previous n ´1 symbols using sparse attention that computes attention weights with the sparsemax function (martins and astudillo, 2016). the use of sparse attention is particularly noteworthy since it departs from the standard soft-attention transformers usually used in practice (vaswani et al., 2017). while the theory does not consider soft-attention transformers concretely, it does suggest that being able to assign no attention to irrelevant tokens, which is impossible with soft attention, could be beneficial. ## learnability of n-gram lms at a high level, we study how well transformers can learn n-gram lms. we do that by training transformers on strings sampled from randomly generated n-gram lms and measuring a notion of distance between the trained and the ground-truth lm. ## experimental results and discussion this section presents the experimental results. all tables contain the mean and standard deviation of the estimate of the kl divergence z d kl between the ground truth and trained lms estimated on the test dataset. @xcite negative values indicate an approximation error of the (low) d kl . ## conclusion we compare the performance of transformers to that of classical n-gram estimation techniques and hand-crafted baselines at learning artificially generated n-gram lms. transformers show a good inductive bias towards learning representation-based n-gram lms but fare worse than classic estimation techniques at learning general n-gram lms, underlining the importance of parameter sharing in the ground-truth model for transformers to learn it well. moreover, the impact of the number of heads and layers on the performance agrees with existing theoretical results. the better performance of sparse-attention transformers motivates further investigation into how sparse attention could be used when learning formal languages with transformers.