Optimizing n-grams

Joint work with Zikun Wang and Michael Pimble.

Introduction

A simple language model is the n-gram ^[1]. Suppose we have a history of text $h$ and we want to predict the next word $w$. If $h$ appears in a large enough reference corpus, we can estimate the probability of $w$ by counting the number of times $w$ follows $h$, divided by the number of times $h$ appears.

Let $w = \text{``pickleball"}$ and $h = \text{``I enjoy playing"}$:

$$P(w \mid h) = \frac{\text{count}(\text{``I enjoy playing pickleball"})}{\text{count}(\text{``I enjoy playing"})}$$

But what if the corpus doesn't contain the history at all? A history like "Terry Tao and Lionel Messi invest in alien food startup" has almost certainly never been written before, so both counts are zero. A typical solution is to keep only the last few tokens of the history to approximate the probability. For example, keeping one token

$$P(w \mid h) \approx P(w_n \mid w_{n-1}),$$

results in a bigram model, since it counts pairs of tokens. Keeping two gives a trigram model, and so on, hence the name n-gram.

On their own, n-grams are not competitive language models. However, they still are useful next to neural models: infini-gram ^[2] delivers exact n-gram counts from massive corpora, and Deepseek's recent work uses n-gram lookup as a form of memory in large language models ^[3].

Problem

Let's say we now want to query a text for an arbitrary length n-gram, we can do a naive linear scan over the text to search for the exact phrase. But new query/n-gram we would like to search for would take a pass over $O(n)$ tokens before finding our exact sequence. We can do better. Using a data structure called a suffix array, we can support exact n-gram queries over a large tokenized corpora ^[2].

Solution

A suffix is a contiguous set of items (e.g. tokens) starting with a fixed index. For example, "great animals" and "are great animals" are suffixes of "Pandas are great animals". Note that a contiguous set of items ending with a fixed index is called a prefix. After collecting all the suffixes we can store the suffixes using a tree data structure.

We can collect all suffixes and arrange them as a tree. Each suffix that shares preceding tokens can be represented using the same path to the root of the tree. This compresses the size of the representation so that a repeated text can be stored using smaller space.

Walking $n$ tokens down from the root lands on the subtree of every suffix that starts with that n-gram and the number of leaves below is its count. In practice nobody stores the tree. Sorting the suffixes lexicographically and keeping only their starting positions gives us the suffix array, and the subtree becomes a contiguos interval of the array. A query is then a binary search for an interval.

There are indexes that take less space and have better theoretical bounds, such as the FM-index used by infini-gram mini ^[4]. They win on size, but each search step becomes a rank query into a compressed structure, and in practice that is slower than a plain binary search over sorted suffixes on modern hardware.

Improving Suffix Trees

Like other tree data structures, a practical representation of a suffix trees involves mapping it to an easily indexable array called a suffix array.

A major drawback of suffix arrays is the size of the stored index. On a 5 GB sample of the C4 dataset ^[5], using the Llama 2^[6] tokenizer, a tokenized suffix array of 4-grams reaches more than 6 GB in size. As the size of the corpus grows it becomes impractical to store the entire suffix array in-memory. Therefore, we would often want to store the suffix array on disk and retrieve from the suffix array during each query.

This data movement between main memory and disk can cause unwanted latency if we are performing a lot of queries. Additionally, the vocabulary of the Llama 2 tokenizer is about 32,000 tokens. That means for an specific n-gram with a vocabulary size of $|V|$ we must make in the worst case $n \times \log_2(|V|)$ comparisons (i.e. retrievals from disk) for every query if we binary search over each level (practically, it makes more sense to binary search over the entire suffix array making the worst case complexity $n \times log_2(|C|)$ where $|C|$ is the corpus size).

This means that most of the time, the bottleneck latency comes from the binary-search over the vocabulary/corpus, not the n-gram length.

Obviously, this is very expensive.

A natural question we can ask is whether we can cache some of the n-grams in-memory to speed up our on-disk n-gram search. Naively, we can cache the disk-pointers to our larger n-gram with an in-memory 1-gram or 2-gram. However, under this model each n-gram branch is of equal importance. In simple terms, this assumes that each in-memory n-gram is uniformly likely to appear in the on-disk n-gram, the power laws of natural language tells us that this isn't the case. Given the limited in-memory storage capacity, can we optimize our cache to exploit the actual distribution?

Optimized Suffix Arrays

Given our memory budget, let's store all bigrams then use the remaining budget to refine only those suffixes with the largest suffix array intervals. At query time the system probes for the longest available in-memory suffix to cut down on the number of binary searches.

Experimental Results

Let's test this hypothesis out on some real data.

You can check out the code here.

Using a small subset of the C4 dataset ^[4], we can test the unigram, bigram, and trigram caches against the optimized n-gram to see if the theoretical benefits are substantial.

We can see that increasing the n-gram memory cache monotonically decreases the step/comparison cost for the on-disk reads. More importantly, the greedy strategy outperforms all uniform n-gram strategies.

Some simple tests comparing latency and memory size between unigram, bigram, and trigram to the optimized suffix array show that the optimized version reduces latency and memory size!

Conclusion

Random disk reads are a substantial cost of an n-gram query. We have shown that an optimized cached n-gram mechanism can save significant disk access time by reducing the number of search steps required to get to the memory location on disk, thereby reducing the latency of a n-gram query.

Citation

@article{li2026ngram,
title = "Optimizing n-grams",
author = "Li, P., Pimble, M., Wang, Z."
year = "2026",
month = "September",
url = "https://peter.bio/blog/n-grams"
}