**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**.

![n-gram](https://files.peter.bio/uploads/2026/09/9d571f87-ddd6-4cd3-bd92-9d2a9ea63350.png)

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 [^6]. 

## 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.

![suffix array](https://files.peter.bio/uploads/2026/08/53d7cc59-1abc-43d1-8e5e-1b1dca360e11.png)

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 [^5]. 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 [^3], using the Llama 2[^4] 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.

![optimized-suffix-array](https://files.peter.bio/uploads/2026/09/7270febf-ea8e-496d-aff0-c4a95eca205f.png)

## Experimental Results

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

You can reference the code [here](https://github.com/lipet2k/n-gram).

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

![steps by length](https://files.peter.bio/uploads/2026/09/91524867-0d8a-454c-b134-f2b3afa24b9e.png)

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!

![search cost versus memory](https://files.peter.bio/uploads/2026/09/089ba7d3-f7cb-42bb-a339-cd37f3432d54.png)

## 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"
}
```

[^1]: Jurafsky, D. and Martin, James H. (2026). Speech and Language Processing: An Introduction to Natural Language Processing, Computational Linguistics, and Speech Recognition, with Language Models. https://web.stanford.edu/~jurafsky/slp3/
[^2]: Liu, J., Min, S., Zettlemoyer, L., Choi, Y., Hajishirzi, H. (2024). Infini-gram: Scaling unbounded n-gram language models to a trillion tokens. https://arxiv.org/pdf/2401.17377
[^3]: Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., Liu, P. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. https://arxiv.org/abs/1910.10683
[^4]: Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., and Bhosale, S., and others (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. https://arXiv.org/abs/2307.09288
[^5]: Xu, H., Liu, J., Choi, Y., Smith, N., Hajishirzi, H. (2025). Infini-gram mini: Exact n-gram Search at the Internet Scale with FM-Index. https://arxiv.org/abs/2506.12229
[^6]: Cheng, X., et al. (2026). Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models. https://arxiv.org/pdf/2601.07372

Optimizing n-grams

The transformer archiecture in language models has popularized the **self-attention mechanism**. This mechanism consists of three key components: 

1. **Key**: A label for the previous tokens, describes what they contain
2. **Value**: The actual information or context that is stored
3. **Query**: The question being asked about the preceding text, this information is not cached

For example, let's say that the preceding context consists of the phrase "My name is Bob", tokenized by word. 

$$
\begin{array}{|c|c|c|}
\hline
\textbf{Order} & \textbf{Tokens} & \textbf{KV Cache} \\
\hline
\text{1} & \text{My} & \boxed{My}\; \\
\hline
\text{2} & \text{My name} & \boxed{My}\;\boxed{name} \\
\hline
\text{3} & \text{My name is} &  \boxed{My}\;\boxed{name}\;\boxed{is} \\
\hline
\text{4} & \text{My name is Bob} & \boxed{My}\;\boxed{name}\;\boxed{is}\;\boxed{Bob} \\
\hline
\end{array}
$$

Each token attends strictly to itself and its preceding tokens. This is enabled by an **attention mask**, which determines which tokens are attended to. 

Large language model inference consists of two stages: the **prefill** and the **decode** stages. The prefill stage processes the context and creates the initial KV-cache, while the decode stage autoregressively generates the output token based on the previous context. During the prefill stage where the KV-Cache is generated, the attention weight for each token is generated by attending to all previous tokens. This would be computationally expensive if we had to recompute from scratch each attention weight. Instead, we can use the KV-Cache to speed up the process. The intermediary key and value cache can be stored and later used to calculate these attention weights. By storing the key and value for previous tokens, we no longer have to restart the attention weight calculation for every new token.

![self-attention](https://files.peter.bio/uploads/2026/05/4cc5349a-15f8-475e-bb5b-b82f8cf0e5ac.png)

Even with these optimizations, inference times are still bottlenecked by the prefill phase. There are different methods for speeding up inference, one obvious optimization is not to attend to all previous tokens, only the ones that are most relavent to the current token (caculated through an **attention score**), this is called **sparse attention**. 

\[
\text{Scaled Attention Score} = \frac{Q \times K^T}{\sqrt{d_k}}
\]

A simple but effective method of sparse attention we will investigate is called **block attention**. In block attention, we chunk the context into separate chunks (these could be separate documents or turns in a chat window). Within each chunk, each token attends to only itself and other tokens.

![block-attention](https://files.peter.bio/uploads/2026/05/7ae0561f-1a3f-43f5-abed-be75d3bf4ce2.png)

Block attention be immensely useful in conjunction with augmented generated retrieval. However, block attention does come with some deficits. One of which being that distant relationships across block boundaries are degraded. To fix this post-processing steps are often required. 

KV-Cache Reuse in Large Language Models

Hello world,

I study problems in theoretical computer science and large-scale data systems with applications to machine learning. My current work sits at the intersection of statistical learning theory and machine learning systems.

My intuition is that theoretical models can extend to empirical results and that informative answers come from abstracting experimental conclusions into concrete theorems and proofs. The problems I find the most interesting connect simple theoretic ideas to practical applications. 

As proofs and experiments become much cheaper to produce, the qualifying bar for *useful/interesting* research increases substantially. Additionally, the *interpretation* of the research results becomes much more valuable. With this in mind, these are some of the problems that I find interesting: which patterns emerge in large-scale systems? Can these scaling behaviors be abstracted into predictable properties? Which properties of the data and model determine the efficiency and efficacy of the system?

One of the many joys of computer science research is the compatibility of theory and experimentation. Theoretical ideas can be cheaply tested (relatively speaking) on widely accessible hardware without bureaucratic processes. This solves an important problem: too many theoretical ideas yield zero practical applications, and many experiments produce unexplainable results. The ability to condense asymmetric amounts of information and connect these observations to mathematical foundations provides computer scientists with a unique and powerful ability!

Candidly, I expect some of my conjectures to be incorrect and revised. These posts are merely ideas that I am still trying to formalize and build into deeper theories. The overall goal is to make the research available to the wider scientific community. Nonetheless, I will attempt to make these ideas as well developed as possible (and limit my reliance on LLMs).

These notes are written for other researchers interested in TCS and large-scale systems, though I hope they remain accessible to anyone curious about learning theory and why machine learning behavior. If something here is useful (or wrong), I would be glad to hear about it and feel free to [reach out](https://peter.bio).

Thanks for reading!

Blog

2026

Optimizing n-grams

KV-Cache Reuse in Large Language Models

Hello World