The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation

The COUNT. Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation. By Ian Magnusson and two others, from the University of Washington and the Allen Institute for AI. Published on the fourteenth of October. Its code repository is worked out: about four hundred source files, with a readme and a licence.

How compute relates to performance in language models has been studied extensively. More recently, models are also being designed to scale their memory use. One way to do that is to add embeddings for n-grams. This scales memory separately from compute, because each pass only processes individual embedding rows, and those rows are much smaller than the storage taken up by the whole embedding table. The catch is in how those embeddings are made. Existing methods train them from scratch, and training needs fast reading and writing, which means expensive GPU memory. Meanwhile the ordinary memory of the host system is cheap, and these training methods leave it underutilized. The memory that is plentiful is not the memory the method can use.

COUNT, short for Counting Offline Unlocks N-gram Transformers, is built to scale on host memory instead. It builds n-gram embeddings without adding any learned parameters. What it stores on the host is not a table of trained vectors but a table of precomputed n-gram counts. Because nothing in that table is learned, nothing in it has to be written during training, and it can live in the cheaper memory.

The counts turn into an embedding through the model itself. For the current n-gram history, COUNT looks at the next tokens that most often follow it in the corpus. It then takes the language model's own unembeddings of those next tokens and weights them together using the corpus frequencies. The result is an embedding assembled from statistics on one side and the model's existing output vocabulary on the other. In pretraining, COUNT is added as an additional input embedding alongside the usual one.

The comparisons hold compute fixed. Added to the pretraining of language models with one billion parameters, COUNT improves validation loss as much as a model given fifty-seven percent more compute. A second comparison asks what to do when GPU memory is the limit. Using COUNT so that backbone parameters can fill GPU memory gives better performance than splitting that memory with additional learned n-gram embeddings.

The last result is about adaptation rather than efficiency. Since the table is just counts, a different table can be switched in. Counts drawn from a domain-specific corpus can replace the original at inference time, adapting the model to a new task without any training at all.

Two benefits come from one design decision. Replacing learned n-gram vectors with stored counts moves the memory burden off the GPU and onto the host, where memory is cheap. And because the counts are plain corpus statistics rather than parameters, they become something you can swap, which turns adaptation to a new domain into a change of table instead of a round of training.