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

AI summary

Overview

Research area: Natural language processing; specifically efficient language model pretraining, n-gram memory architectures, and scaling laws.

Technical level: Intermediate. Readers need comfort with transformer input embeddings, unembedding/output heads, training compute budgets, and loss scaling laws.

Scope: The paper introduces COUNT, a method for building n-gram input embeddings for language models from offline corpus counts stored in host RAM, and evaluates it across 60M–1B parameter backbones, multiple token-to-parameter ratios, and 74 downstream task configurations.

What This Paper Is About

Language models can be improved by adding n-gram embeddings to their input, which scales memory separately from compute because each forward pass touches only a few embedding rows rather than the whole table. Existing methods learn those n-gram embeddings from scratch, so their weights and optimizer state must live in expensive GPU memory and compete with the backbone model's parameters. COUNT instead builds n-gram representations from precomputed corpus counts kept in cheap host RAM, adding no learned parameters, and its count tables can be swapped at inference time to adapt to a new domain without any training.

Key Contributions

1. COUNT, a parameter-free n-gram embedding method. COUNT stores a frozen table of the top-K most frequent next tokens for stored n-gram histories, where the table is derived offline from corpus counts using stupid backoff. At each position it uses those frequency scores to weight and combine rows of the backbone LM's own unembedding matrix into an additional input embedding, which is summed with the ordinary token embedding. No learned parameters are added, and the large table lives in host RAM rather than GPU memory.

2. A systematic scaling study from 60M to 1B parameters. The authors pretrain standard transformers, COUNT models, and learned n-gram embedding baselines at token-to-parameter ratios D/N ∈ {10, 20, 40} on allenai/dolma3_mix-6T-1025-7B, fitting compute-optimal scaling laws per method and comparing them on LM validation loss, task loss, and downstream accuracy.

3. Demonstration of training-free domain adaptation by table swapping. Because COUNT tables map histories to scores over a shared vocabulary, a domain-specific table can be interpolated with the pretraining table at inference time. The authors test this on math, biomedical, and law domains, and separately test swapping in much smaller tables to fit RAM-constrained deployments.

4. Released artifacts. The authors release n-gram counts of the full Dolma 3 pretraining corpus, table-building and training code, and pretrained checkpoints for COUNT and learned n-gram baselines (code at https://github.com/IanMagnusson/COUNT).

Main Findings

• COUNT buys performance with host RAM instead of compute. At a fixed compute budget, adding COUNT as an extra input embedding for pretraining 1B-parameter LMs improves validation loss as much as an LM with 57% more compute (abstract). Across the study, the authors report that standard transformers would need 46–89% more compute to reach the same language modeling loss as COUNT, and 6–56% more compute to match downstream task loss. The gains come from storing up to 797 GB of n-gram frequencies in host RAM.

• COUNT is Pareto-optimal against learned n-gram embeddings on parameter budget. In comparisons of loss against total learned parameters — a hardware-agnostic proxy for GPU memory footprint — the Pareto frontier mostly follows COUNT because it adds no learned parameters. Learned n-gram embeddings can reach comparable performance but only by adding learned parameters that cannot be easily offloaded during training and that contend with the backbone for GPU memory.

• COUNT beats output-level fusion of counts. Late fusion, which combines count-based next-token distributions with neural predictions at the output, also keeps tables in host RAM without competing for GPU memory, but constructing input embeddings with COUNT gives larger generalist gains. The authors also find that activating late fusion during pretraining does not improve on purely inference-time late fusion.

• Inference-time domain adaptation improves in-domain results without training. Interpolating a domain-specific table with the general table at weight λ = 0.25 — the value that decreased loss for all three domains — improved in-domain task loss on math, biomedical, and law. Improvements were larger for math and law; biomedical results were mixed. Accuracy results were more mixed but still increased on average for math and law. For reference, at D/N = 20, COUNT with the generalist table reduces task loss by 0.010 BPB relative to the standard 1B transformer.

• Tables can shrink by more than an order of magnitude at inference time. Replacing a 1B COUNT checkpoint's table with general-domain tables of 459.222, 142.037, 68.458, or 17.114 GB preserved better performance than a standard transformer, giving graceful degradation for RAM-constrained deployments.

• Throughput overhead is small. At the 1B backbone scale, COUNT retains 93.1–97.8% of standard transformer training throughput, 95.5–95.8% of prompt prefill throughput, and 95.9–96.7% of autoregressive decoding throughput. Three optimizations recover most of the gap: compiled multithreaded host lookup, prefetching lookups while the GPU processes the preceding batch, and a fused embedding reduction. For the learned baseline, full-table optimizer updates ran at 4.4% of standard transformer training throughput at 190M with a 17.1 GB table, while updates restricted to retrieved rows ran at 96.3%; the 1B learned baseline trains at 98.3%.

• Implementation details that matter. COUNT's table stores K = 32 candidates per history, while training uses K = 8 and J = 6, with embedding amplification applied to the combined input. The primary general-domain table stores the top-32 continuations for every history of order one through six seen at least 265 times, after dropping order-4 to order-6 n-grams with fewer than seven occurrences, and occupies 797.487 GB. A 17.1 GB table serves the 60M, 100M, and 190M backbones; the 797 GB table serves the 370M, 600M, 760M, and 1B backbones.

• The unembedding matrix is the best basis. Appendix B.7 compares using the input token embedding versus the unembedding matrix as the next-token basis, and finds the unembedding performs best.

Methodology in Plain English

The researchers start from a simple idea: instead of learning what n-grams mean, count how often they occur. They scan the full Dolma 3 pretraining corpus and, for every short sequence of tokens (up to order six), record which tokens most often come next and how often. These counts are compressed into a table that lives on the host machine's ordinary RAM.

At training time, when the model reads a token, it looks up the longest matching history in the count table. The table returns a short list of likely next tokens with frequency scores. COUNT then takes the model's own unembedding matrix — the same matrix used to turn hidden states into next-token predictions — and pulls out the rows for those candidate tokens, weighting them by their scores and averaging them into a single dense vector. That vector is added to the ordinary token embedding before the first transformer block. Because the unembedding matrix already exists, nothing new is learned.

The authors compare this setup against two baselines: a standard transformer, and a learned n-gram embedding method (Over-Encoding as implemented in LongCat, using bigram and trigram embeddings with four hash tables per order and table width d/8, with a 1:1 ratio of learned input-embedding parameters to backbone parameters). All arms share the same OLMo-core backbone ladder and training recipe, trained on allenai/dolma3_mix-6T-1025-7B with 8,192-token sequences and the Dolma 3 tokenizer.

Evaluation uses three aggregates, all weighted equally per configuration: LM validation loss (mean token cross-entropy in nats across 11 Paloma-based validation sources: C4, Dolma Books, Dolma Common Crawl, Dolma peS2o, Dolma Reddit, Dolma Stack, Dolma Wiki, ICE, M2D2 S2ORC, Pile, and WikiText-103); task loss (mean bits per byte of reference answers across 74 downstream task configurations); and downstream accuracy (mean accuracy across the 41 multiple-choice configurations). The authors fit scaling laws of the form L(N, D) = E + A/N^α + B/D^β and select the allocation minimizing loss at each compute budget, with compute defined as C = 6ND.

For domain adaptation, they build math (83.398 GB, from 10.66B tokens of Dolmino Math), biomedical (81.723 GB, from 95.39B tokens of olmOCR health), and law (80.151 GB, from 12.34B tokens of olmOCR crime-and-law) tables, then interpolate them with the general table using weight λ ∈ {0, 0.25, 0.5, 1}, selecting the eight highest-scoring candidates without renormalizing.

Why This Matters

Impact on research. The paper reframes n-gram memory as a resource problem rather than a modeling problem. By showing that corpus statistics can supply useful input embeddings with zero learned parameters and no GPU footprint, it opens a route to scaling language models along a memory axis that is cheap and largely unused. It also connects modern LM pretraining back to classical count-based language modeling, and gives a systematic scaling-law comparison across model sizes, training budgets, and a large downstream suite.

Real-world applications:

• Deploying capable models on hardware with limited GPU memory, since COUNT's tables can be stored on host RAM and even compressed by more than an order of magnitude at inference time while still beating a standard transformer.
• Adapting a generalist model to a specialized domain (math, biomedicine, law, or other narrow corpora) by swapping or interpolating in a domain count table, with no training run.
• Ingesting data that cannot be released as raw text, since an n-gram count table is an aggregate statistic rather than a corpus — though the authors caution that high-order counts over rare contexts can still reveal fragments of the source.
• Serving as a lightweight prior for domains with too little data to train on, since n-gram frequencies can be estimated from far less text than pretraining requires.

Industry relevance. High-bandwidth GPU memory is the scarcest and most expensive resource in training and serving. COUNT shifts part of the scaling burden to cheaper, abundant host RAM and keeps GPU memory free for backbone parameters, which directly affects hardware procurement and serving costs. The measured throughput retention (93.1–97.8% for training at 1B scale) indicates the approach is practical rather than purely theoretical, and the released checkpoints, tables, and code lower the barrier to adoption.

Future Directions

• Pushing the throughput recovery further. The authors note that further optimizations would "not have far to go" to recover both COUNT and learned n-gram embeddings to full standard transformer throughput, suggesting headroom in host lookup, prefetching, and fused reduction.

• Extending training-free domain adaptation. The interpolation weight was swept over only four values and a single weight (0.25) decreased loss for all three domains tested; the full grid and selection procedures for picking λ per deployment remain open.

• Smaller and rarer corpora as count sources. The conclusion argues that count-based priors can likely be fit from corpora too small for parametric learning, which would extend COUNT to low-data settings beyond the math, biomedical, and law tables tested here.

• Editing and auditing the prior directly. Because each row lists next-token candidates and scores for a history, the count table can in principle be read, corrected, or extended by editing rows rather than retraining — a capability the paper raises but does not evaluate.

Target Audience

This paper is most useful to language model pretraining researchers and engineers working on scaling, efficient architectures, or memory-augmented models; to practitioners who need domain adaptation without a training budget; and to readers interested in how classical count-based language modeling can be reintegrated with modern neural LMs. Readers should already understand transformer embeddings, token-to-parameter ratios, and basic scaling-law fitting to get the most from the results.

Authors’ 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.