Counting n-grams boosts memory use and speeds up language model training
The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation
Computation and LanguageMachine Learning
Summary
Language models use big tables to predict text, but these tables take a lot of memory and computing power. The authors show a way called COUNT that uses pre-stored counts of word groups (n-grams) on cheap computer memory instead of learning new embeddings from scratch. This saves expensive GPU memory and can improve model quality without extra computing cost. COUNT also allows changing these counts easily to adapt a model to new topics without retraining.
What this means in practice
- •For machine learning engineers: Use precomputed n-gram counts stored in host memory to improve large language model training efficiency under GPU memory limits.
- •For nlp system builders: Adapt language models at inference time to new domains by swapping in domain-specific n-gram count tables without retraining.
Authors
Ian Magnusson, Noah A. Smith, Pang Wei Koh
Abstract
The scaling relationship of compute and performance for language models (LMs) has been extensively studied; more recently, models are being designed to scale memory use, as well. In particular, adding embeddings for n-grams scales memory separately from compute use, because each pass only processes individual embedding rows, which are much smaller than the storage used for the whole embedding table. However, existing methods train embeddings from scratch, which needs expensive GPU memory for fast reading and writing. Meanwhile, the ordinary memory of the host system is cheap and underutilized by these training methods. To scale on host memory, we introduce COUNT (Counting Offline Unlocks N-gram Transformers), which builds n-gram embeddings without adding learned parameters by storing precomputed n-gram counts on host. We use the corpus frequencies to weight together the LM's own unembeddings of the next tokens most often following the current n-gram history. At a fixed compute budget, COUNT can be added as an additional input embedding for pretraining 1B-parameter LMs to improve validation loss as much as an LM with 57% more compute. If GPU memory is limited, using COUNT to let backbone parameters fill GPU memory gives better performance than splitting memory with additional learned n-gram embeddings. Beyond efficiency, COUNT allows switching in a different table of n-gram counts from a domain-specific corpus at inference time to adapt to a new task without any training.