COUNT: a 797 GB read-only n-gram table in host RAM buys a 1B LM the loss of 57% more compute — zero learned parameters
UW/AI2's COUNT builds n-gram embeddings from precomputed corpus counts in host RAM, adds no learned parameters, and swaps tables at inference for training-free domain adaptation.
Researchers at the University of Washington and the Allen Institute for AI introduce COUNT (Counting Offline Unlocks N-gram Transformers), a pretraining method that adds n-gram memory to a transformer with zero learned parameters: it stores precomputed corpus counts in host RAM and uses them to weight the model's own unembedding rows. On a 1B-parameter backbone with a 797 GB count table, the authors report validation loss equivalent to a standard transformer trained with 57% more compute. The same mechanism doubles as a training-free domain-adaptation knob for math, biomedicine, and law.
What shipped
COUNT is an additional input embedding for a transformer backbone, built from a frozen, read-only table of n-gram statistics. At each position, the model looks up the current n-gram history in a table that maps each stored history to its K most frequent next tokens and their count-based scores. Those scores weight rows of the LM's own unembedding matrix; the weighted combination is summed with the token embedding as the input to the first transformer block:
f_t = (1/J) Σ_j Σ_k S(w_k | h_t,j) · B[w_k], h_t^(0) = f_t + e_t
where B is the unembedding matrix and S is a stupid-backoff score (paper, Eq. 1–2). Scores are the count ratio C(h,w)/C(h) when the continuation is observed, and a 0.4-discounted backoff to a shorter history when it is not; they are used without renormalization. Because the unembedding head already learns representations of the predicted next tokens, reusing it as the basis adds no learned parameters — the table is read, never trained.
The table is built offline by counting n-grams through order six over the full Dolma 3 pretraining corpus (paper §3.2). Each stored history maps to its 32 most frequent next tokens (training uses 8), with scores quantized to 8 bits per order (GitHub README, dataset card).
The flagship table is 797 GB, sized to saturate the host RAM on the authors' compute nodes. A 17 GB variant covers the 60M–190M backbones in training; a 142 GB variant is a smaller swap-in table for inference. Domain tables for math (83 GB), biomedicine (82 GB), and law (80 GB) are built from domain subsets of Dolma 3 (dataset card).
The headline result is a compute-equivalence: at a fixed compute budget, adding COUNT to a 1B-parameter LM improves validation loss as much as training a standard transformer with 57% more compute (abstract). Across the 60M–1B backbone ladder, standard transformers would need 46–89% more compute to match COUNT's language-modeling loss and 6–56% more to match downstream task loss (paper intro). COUNT also reports Pareto-optimal loss-vs-learned-parameter tradeoffs against learned n-gram embeddings, and larger gains than output-level fusion of corpus counts (the kNN-LM-style approach) (paper intro).
Everything is released. The GitHub repo (Apache-2.0) has training code and a fork of OLMo-core with the COUNT and learned-n-gram embedding layers. The n-gram tables (ODC-By, derived from Dolma 3) and COUNT checkpoints from 60M to 1B at three token budgets are on Hugging Face.
Why it matters
The interesting part is what COUNT decouples. Learned n-gram embeddings scale a model's memory footprint, but their weights and optimizer state must live in GPU memory during training, competing with the backbone for scarce memory. COUNT's table is read-only, so it sits in cheap, underutilized host RAM while backbone parameters fill the GPU. The authors' framing: "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" (abstract). That is a different scaling axis — memory that scales without touching the learned parameter budget or the compute per step.
The second payoff is training-free adaptation. Because the count tables are computed offline, they can be recomputed over a new corpus and swapped in at inference. The paper reports that a 1B COUNT checkpoint combining its general table with a domain-specific one (a 0.25/0.75 domain/general interpolation, keeping the eight highest-scoring candidates) improves task loss in math, biomedicine, and law. Accuracy improves on average for math and law, with mixed results for biomedical (paper). No fine-tuning, no gradient steps — the embedding table becomes the adaptation knob.
For engineers, this is runnable today: the release includes the count tables, training code, and checkpoints, so the 57% compute-equivalence claim can be reproduced rather than taken on faith. The architectural lesson is broader than the paper's 1B scale: a read-only, host-resident memory store can be bolted onto a transformer without touching its learned parameter budget, and swapped at inference time.
Caveats
- The 57% figure is a fitted scaling-law equivalence in training FLOPs, not a wall-clock speedup. The paper measures throughput separately and reports COUNT retains roughly 93–98% of standard-transformer training throughput at 1B under its optimized setup (paper).
- Everything is demonstrated at 60M–1B parameters. Whether host-memory scaling holds at larger scales is untested.
- Domain-swap accuracy is mixed even where task loss improves — the adaptation story is stronger on loss than on accuracy.
- This is a newly posted preprint with no independent replication yet.
- Two reproducibility caveats from the authors: the released 190M checkpoints are a rerun (the originals were lost), and the released code is the simple implementation, "not the optimized runtime behind the paper's throughput measurements" (GitHub README).