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

1 INTRODUCTION

• Learned n-gram embeddings compete with backbone parameters for GPU memory because their weights and optimizer state must be read and written quickly.Count-based methods estimate continuations from corpus frequencies, but related neural-LM integrations fuse next-token distributions at the output rather than scaling dense input representations.Source: §1 · pp. 1–2
• COUNT uses precomputed corpus next-token statistics to weight and combine the LM’s unembeddings as an additional input embedding, without adding learned parameters.Its read-only count tables can be stored in host RAM and recomputed over new corpora for inference-time adaptation without parameter training.Source: §1 · p. 2
• 46–89% more compute is needed for a standard transformer to match COUNT’s language-modeling loss, and COUNT also improves task loss comparably to a 6–56% compute increase.Figure 1 illustrates the validation-loss comparison; the reported results also include task-loss comparisons.Source: §1 · p. 2, Figure 1 · p. 2

2.1 PRELIMINARIES

• An LM predicts the next token from a prefix, with token embeddings providing inputs to its first transformer block.The method denotes each token’s embedding as et and the first-block input as h(0)t.Source: §2.1 · pp. 2–3

2.2 COUNT

• For each stored history, COUNT retrieves its top-K next tokens and scores, then weights unembedding rows to form a dense COUNT embedding.The COUNT embedding is added to the token embedding as input to the first transformer block.Source: §2.2 · p. 3

2.3 N-GRAM COUNT-BASED SCORES

• COUNT can use count-based continuation scores computed offline; its experiments use stupid backoff, which backs off to shorter histories when a continuation is unseen.The backoff scores need not sum to one and are used without renormalization.Source: §2.3 · p. 3

3.1 BACKBONE LM

• Backbones range up to 1B parameters, with training token-to-parameter ratios D/N of 10, 20, and 40 on the Dolma 3 mix.The same corpus supplies the n-gram counts.Source: §3.1 · p. 3

3.2 BUILDING N-GRAM TABLES

• N-gram tables are built from counts through order six, with pruning, history thresholds, top-K continuations, and score quantization controlling table size.The largest table is 797 GB for 370M–1B backbones; smaller backbones use 17.1 GB tables.Source: §3.2 · p. 4

3.3 LEARNED N-GRAM EMBEDDINGS

• The learned n-gram baseline uses jointly trained tables and projections, with bigram and trigram embeddings, four hash tables per order, and embedding amplification.Input-embedding parameters are set to equal the remaining learned parameters; amplification is also applied to COUNT but not the standard-transformer baseline.Source: §3.3 · p. 4

3.4 EVALUATION

• Evaluation aggregates validation loss over 11 sources, task loss over 74 configurations, and accuracy over 41 multiple-choice configurations.Figure 2 plots task loss against training compute and downstream accuracy for the same runs.Source: §3.4 · pp. 4–5, Figure 2 · p. 5

4.1 SCALING HOST STORAGE SAVES GPU COMPUTE

• 46–89% more compute matches COUNT's LM-loss gains, while 6–56% more compute matches its task-loss gains over standard transformers.COUNT's downstream accuracy is higher than the standard transformer at matched compute at almost all budgets.Source: §4.1 · p. 5

4.2 USING HOST RAM FOR N-GRAMS FREES GPU RAM FOR BACKBONE PARAMETERS

• At equal learned-parameter budgets, COUNT leads LM loss across 60M–1B backbones; downstream accuracy follows the same trend with more noise.The bottom-row frontier occasionally includes a standard transformer or learned n-gram embedding configuration.Source: Figure 3 · p. 6, §4.2 · p. 6
• COUNT's input embeddings give larger generalist gains than late fusion; inference-only late fusion remains an option when pretraining cannot change.Late fusion combines count-based next-token distributions with neural predictions at the output without adding learned n-gram parameters.Source: §4.2 · p. 6

5.1 ADAPTING TO A NEW DOMAIN WITHOUT TRAINING

• At λ = 0.25, in-domain task loss improves on all three domains; accuracy is mixed but increases on average for math and law.Each Figure 4 point represents one subtask at one training budget.Source: Figure 4 · p. 7
• The domain-specific math, biomedical, and law tables are substantially smaller than the 797.487 GB general-domain table.Their respective sizes are 83.398, 81.723, and 80.151 GB.Source: §5.1 · p. 7

5.2 SHRINKING THE TABLE TO FIT A RAM-CONSTRAINED DEPLOYMENT

• On 1B COUNT checkpoints, table storage can shrink by more than an order of magnitude while performance remains better than a standard transformer.The smaller table can be substituted at inference time without training because the tables share a next-token-score format.Source: §5 · p. 7, §5.2 · p. 8

• Training throughput is measured on 8 H200 GPUs, while inference prefill and decode are measured on one H200 after warm-up.Training rates use median tokens per second from optimization steps 50–100.Source: §6 · p. 8
• COUNT retains 93.1–97.8% of standard transformer training throughput, 95.5–95.8% for prompt prefill, and 95.9–96.7% for autoregressive decoding.These measurements are at the 1B backbone scale.Source: §6 · p. 8
• Compiled multithreaded host lookup, lookup prefetch, and fused embedding reduction are used to recover most standard-transformer throughput.Prefetching overlaps lookups with GPU processing of the preceding batch; fused reduction gathers and sums candidate vectors without storing them separately.Source: §6 · pp. 8–9
• For the learned n-gram baseline, sparse updates raise training throughput from 4.4% to 96.3% of standard-transformer throughput at 190M with a 17.1 GB table.The 1B learned baseline trains at 98.3%; host-offloaded inference throughput is similar to COUNT.Source: §6 · p. 9

7.1 N-GRAMS IN NEURAL LANGUAGE MODELS

• Neural-LM n-gram methods include hashing, lower-order token representations, inference-time token-embedding combinations, and auxiliary or hidden-state-conditioned models.The cited approaches address the intractability of independently representing all |V|^n n-grams (Roy et al., 2022; Huang et al., 2025; Liu et al., 2026; Zhou et al., 2026; Qu et al., 2026; Yu et al., 2025; Ding et al., 2026; Tseng & De Sa, 2026; Cheng et al., 2026).Source: §7.1 · p. 9
• COUNT represents count-based next-token distributions in the neural LM’s shared embedding space, unlike late-fusion methods that combine predictions at the output.Late-fusion work includes context-datastore or Infini-gram interpolation and adding n-gram log probabilities to output logits (Khandelwal et al., 2020; Liu et al., 2025b; Li et al., 2022).Source: §7.1 · p. 9

7.2 TOKENIZATION AND N-GRAMS

• Tokenization-focused n-gram work changes segmentation or makes multiword n-grams individual LM tokens, a different motivation from COUNT.The cited examples explore further segmentation or relaxed segmentation (Clark et al., 2022; Pagnoni et al., 2025; Liu et al., 2025a).Source: §7.2 · pp. 9–10

7.3 RETRIEVAL-AUGMENTED LANGUAGE MODELS

• Retrieval-augmented LMs share COUNT’s motivation of complementing parametric knowledge with corpus information, but feed token sequences into the model, increasing context use and compute.COUNT stores counts rather than the raw corpus, making its data store much smaller than the source documents.Source: §7.3 · p. 10

8 CONCLUSION

• COUNT provides n-gram memory in cheaper host RAM rather than GPU memory, leaving GPU memory for backbone parameters without adding learned parameters.The conclusion also notes that high-bandwidth GPU memory is scarce and expensive, while host RAM is cheaper per byte and largely uncontended by the model.Source: §8 · p. 10
• With only a small amount of domain-specific data, COUNT fits count-based priors that steer a generalist pretrained model toward better in-domain performance.The conclusion suggests such priors may also be useful with corpora too small for parametric learning.Source: §8 · p. 10
• COUNT priors can be read, corrected, or extended by editing rows directly, rather than retraining the model.Each row lists next-token candidates and scores for a history.Source: §8 · p. 10

AI USE STATEMENT

• The authors used generative AI tools to implement methods and for several other disclosed tasks, reviewed the assisted work, and take responsibility for the final content.The first author manually reviewed the final method implementation, and all authors revised the paper multiple times.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

• COUNT tables summarize continuations rather than documents, making them less sensitive and more compact to share, though rare high-order contexts can still reveal source fragments.The ethics statement says such tables should be released with the same care as other derived statistics; experiments use publicly available pretraining data and evaluation suites.Source: ETHICS STATEMENT · pp. 10–11, §8 · p. 10

REPRODUCIBILITY STATEMENT

• The paper specifies its training recipe, n-gram table construction, evaluation and timing procedures, and releases counts from the full Dolma 3 corpus, code, and pretrained checkpoints.Appendix B contains additional ablation and inference-time sweep results.Source: REPRODUCIBILITY STATEMENT · p. 11

A.1 OLMO-CORE

• The 1B configuration has 1,279,395,840 backbone parameters under the paper’s budget convention, which excludes input embeddings and added n-gram parameters but includes the output head.Token budgets use this exact parameter count and are rounded to full optimizer steps.Source: §A.1 · p. 17

A.2 N-GRAM TABLES

• The primary general-domain n-gram table stores 32 continuations per qualifying history and occupies 797.487 GB.It covers histories of orders one through six seen at least 265 times, with an additional occurrence floor for orders four through six.Source: §A.2 · p. 18
• Domain tables use Math, Health, and Crime and Law corpora, with orders through six, a count floor of one, and a 0.4 backoff factor.The source subsets contain 10.66B, 95.39B, and 12.34B tokens, respectively; biomedical and law history-count thresholds were set to fit within the math table’s 83.398 GB.Source: §A.2 · p. 18, Table 3 · p. 18

A.3 EVALUATION BREAKDOWN

• Task loss averages 74 per-configuration BPB values, while the accuracy aggregate covers 41 multiple-choice configurations.Validation loss averages token cross-entropy within each of 11 sources before averaging across sources.Source: §A.3 · p. 18
• ARC, MMLU, CommonsenseQA, PIQA, and SocialIQA use ranked-choice and completion prompts; other benchmarks use completion only.Source: §A.3 · p. 19
• The evaluation spans knowledge and science, code, math and numerical reasoning, commonsense, reading and cloze, and symbolic and pattern reasoning.Subject subsets, splits, prompt formats, and language variants count separately; each configuration has equal weight.Source: Figure 8 · p. 19, §A.3 · p. 19

A.4 THROUGHPUT CONFIGURATIONS

• Throughput measurements cover training, prefill, and decode, with reported rates in tokens per second per executing H200 and peak GPU memory in a separate table.COUNT tables remain in host RAM; inference uses D/N = 40 checkpoints on C4 validation prompts, and lookup and host–GPU transfer are timed.Source: §A.4 · p. 19
• The throughput setup uses one H200 GPU, batch sizes 1 and 8, 8,192-token prompts, COUNT’s 797.5 GB table, and K = 8 per order for orders 1–6.Timing uses three warmup batches and 96 timed batches at batch 1 or 12 at batch 8.Source: §A.4 · p. 20

B.1 LATE FUSION

• The late-fusion variants use only the longest matching history, rather than averaging orders, because the neural LM cannot correct noisy low-order predictions in late fusion.The compared variants use the same tables and K = 8.Source: §B.1 · p. 20
• Additive fusion interpolates probabilities using top-K stupid-backoff scores while capping retained mass at one; multiplicative fusion uses a normalized product.The multiplicative variant retains at least ε = 10^-4 total mass for non-candidates, distributed uniformly.Source: §B.1 · p. 21
• COUNT achieves better generalist performance than either additive or multiplicative late fusion, while training-free adaptation is shared by count-based methods.COUNT combines that adaptation capability with better pretrained generalist performance.Source: §B.1 · p. 21, Figure 9 · p. 21
• At 190M, pretraining late-fusion arms with learned scalar weights does not improve performance over inference-only fusion, so other scales use inference-only fusion.Source: §B.1 · p. 21

B.2 FULL THROUGHPUT RESULTS

• Tables 7–9 report training, prefill, and decode throughput settings, with rates in tokens per second per executing H200 and peak GPU memory reported separately.COUNT tables are host-resident in all measurements; the backbone unembedding matrix is GPU-resident for every method.Source: §B.2 · p. 21, Table 7 · p. 22, Table 8 · p. 23

B.3 TABLE SWAP INTERPOLATION ON LOSS

• The table-swap interpolation grid reports task-loss changes at four interpolation weights, averaged across subjects within MATH, LabBench, and the MMLU biomedical group.Source: Figure 12 · p. 24

B.4 INFERENCE-TIME K AND J VARIATION

• Larger departures from COUNT’s trained K and J settings lose most of its improvement over the standard transformer, although task accuracy is especially tolerant.The sensitivity sweeps change K or J without retraining; pretraining each configuration may retain more of the improvement.Source: §B.4 · p. 24, Figure 13 · p. 24

B.5 TRAINING WITH FEWER N-GRAM ORDERS

• Training COUNT with orders one through six gives lower language-model and task loss at every budget than training with orders one through three.The comparison uses the same table and K; the authors therefore train with all six orders.Source: §B.5 · p. 24, Figure 15 · p. 25

B.6 ABLATING EMBEDDING AMPLIFICATION √

• Embedding amplification lowers COUNT’s language-model and task loss at every budget, while its effects on standard transformers vary by scale and include higher loss at 760M and 1B.For standard-transformer comparisons, the authors therefore omit amplification.Source: §B.6 · pp. 24, 26, Figure 16 · p. 26, Figure 17 · p. 26

B.7 ABLATING THE COUNT EMBEDDING BASIS

• At 1B, the unembedding basis gives lower language-model and task loss and higher accuracy than the token-embedding basis at all three training budgets.The token-embedding basis still lowers language-model loss relative to the standard transformer, but has worse task loss and accuracy.Source: §B.7 · p. 26, Figure 18 · p. 26