COUNT on RAM: A Different Way to Scale Language Models

TL;DR for operators

Improving a language model does not necessarily require proportionally more GPU-resident parameters. COUNT stores corpus-derived continuation statistics in host RAM and converts them into transformer input representations using existing model weights.

In controlled experiments, COUNT delivers validation-loss improvements equivalent to 46–89% more standard-transformer training compute, according to fitted scaling laws. An optimized implementation retains more than 93% of baseline training throughput, even with a count table approaching 798 GB.

The same statistical memory can be replaced or mixed with specialized corpus statistics without further parameter training. For teams developing their own language models, this introduces a new infrastructure tradeoff: additional host RAM and lookup operations in exchange for better predictive performance and more flexible adaptation. These results are not yet evidence of equivalent production cost savings.

The memory budget beyond the GPU

A language-model team wants better performance, but another round of GPU expansion is difficult to justify. Increasing the number of learned parameters consumes accelerator memory, while extending training requires additional computation. Meanwhile, the servers may have substantial CPU-accessible memory available.

Could that memory contribute directly to the model’s predictions rather than merely storing datasets?

Ian Magnusson, Noah A. Smith, and Pang Wei Koh investigate this question in The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation.¹

Their result is striking: for a nominal 1B-parameter model, COUNT reaches validation loss comparable to a standard transformer using approximately 57% more training compute, as estimated from the paper’s fitted scaling relationships.

The improvement comes from statistical information extracted before training, rather than a separately learned collection of n-gram embedding parameters. Understanding how those statistics enter the transformer explains both the performance gain and its infrastructure requirements.

How frozen word statistics become neural representations

COUNT begins with a familiar statistical operation: counting which tokens follow particular sequences in a text corpus.

For example, a corpus may reveal that certain words frequently follow a phrase such as “the interest rate.” These frequencies provide information about plausible continuations. Longer histories can be more informative, but they are also sparser.

COUNT handles missing combinations using stupid backoff. When a particular continuation has not been observed after a longer history, the method considers a shorter history and discounts its score by a factor of 0.4. The resulting scores need not sum to one.

The central innovation is what happens next.

A transformer already has an unembedding matrix: learned vectors used to convert internal representations into scores over possible output tokens. COUNT reuses those vectors to construct an additional representation of the current context.

At each token position, the system:

1. Retrieves scored continuation candidates from several matching history lengths.
2. Weights the corresponding unembedding vectors by their continuation scores.
3. Averages these contributions across history lengths.
4. Adds the resulting vector to the ordinary token embedding before transformer processing.

In the primary training configuration, COUNT uses six history orders and retrieves eight continuations per order. The count tables themselves are computed offline and remain frozen.

This distinction matters. The transformer still learns its ordinary parameters, but it can use information from a much larger collection of corpus statistics without training an independent embedding vector for every stored sequence.

The placement of this information also matters. COUNT inserts statistical representations before the transformer layers, allowing the network to process and potentially correct them. In the reported comparisons, this input-level integration outperforms additive and multiplicative methods that combine count-based and neural predictions only at the output. The appendix’s late-fusion experiments serve as an architectural comparison rather than a separate scaling result.

The compute-equivalent gains are substantial, but modeled

The authors compare standard transformers, COUNT, and learned n-gram embedding alternatives using matched OLMo-core configurations.

The experimental ladder covers seven backbone sizes from 60M to nominal 1B, with token-to-parameter training ratios of 10, 20, and 40. Models are pretrained using Dolma 3, then evaluated across eleven language-model validation sources and seventy-four downstream task-loss configurations.

The principal findings are:

These percentages require careful interpretation.

The authors fit separate loss-scaling relationships for standard transformers and COUNT, then estimate the training compute each would require to reach equivalent loss. The conventional approximation uses backbone parameter count and training tokens, but excludes COUNT’s additional CPU lookup and embedding operations.

Consequently, a 57% compute-equivalent advantage is not a measured 57% reduction in training expenditure.

The architecture comparison supplies another result. COUNT occupies most of the measured Pareto frontier for language-model loss against total learned parameters. Learned n-gram embeddings sometimes achieve comparable performance, but require additional trainable parameters and associated GPU-memory resources.

Downstream multiple-choice accuracy generally improves as well, although the results are noisier than loss measurements. Better predictive loss is the more consistent finding.

Nearly 798 GB of statistics, without nearly 798 GB of GPU memory

COUNT’s largest general-domain count table occupies approximately 797.5 GB of host RAM. The smaller primary table is approximately 17.1 GB.

That is a substantial memory commitment, but the largest table does not need to reside on the GPU. Instead, optimized CPU lookup, prefetching, caching, and fused embedding reduction limit the overhead of incorporating its contents into GPU computation.

For the nominal 1B configuration, the authors report:

Training benchmarks use eight NVIDIA H200 GPUs; inference benchmarks use one H200.

In one training comparison using six history orders and the 797.5 GB table, COUNT processes 52,008 tokens per second per executing H200, versus 55,836 for the standard transformer. Peak allocated GPU memory is approximately 61.0 GiB for COUNT and 60.8 GiB for the standard model, compared with 63.4 GiB for the GPU-resident learned n-gram baseline.

These measurements indicate that large host-resident statistical tables need not impose proportionally large GPU-memory or throughput penalties under the tested implementation.

The authors also replace the original table at inference time with alternatives more than an order of magnitude smaller while preserving an advantage over the standard transformer. This establishes a deployment tradeoff between available host RAM and predictive performance, rather than requiring the largest training-time table everywhere.

Domain adaptation without changing model weights

COUNT’s external statistical memory creates another capability: modifying the corpus-derived information available to a pretrained model without changing its parameters.

The authors build specialized count tables for mathematics, biomedicine, and law. At inference time, they interpolate the original general-domain continuation scores with domain-specific scores.

With a mixture of 75% general-domain and 25% domain-specific scores, aggregate in-domain task loss improves across all three evaluated domains without further gradient-based training.

Accuracy responds less consistently. Mathematics and law improve on average, while biomedical results are mixed.

This distinction matters for specialized applications. A lower task loss indicates better probabilistic predictions under the evaluation metric, but does not guarantee more correct answers on every domain-specific task.

The procedure is also conditional on having a compatible COUNT-pretrained model and count tables using the same tokenization and scoring conventions. It is not a training-free upgrade that can simply be attached to any existing transformer.

For a team maintaining a specialized model, however, table interpolation creates an additional experimentation route: compare several corpus-derived statistical priors against a fixed model checkpoint before committing to more parameter training.

The infrastructure decision, and its remaining boundaries

The business inference from these results is a change in how model-development teams can allocate resources.

Instead of comparing only larger backbones, longer training runs, and learned embedding expansions, teams can examine a fourth option: host-resident statistical memory.

For an organization developing its own specialist language model, that comparison should account for predictive loss, domain accuracy, GPU allocation, host-RAM capacity, memory-access throughput, and the effort required to construct and maintain count tables. A latency-sensitive service may value a smaller table differently from an offline analytical workload.

Several boundaries remain material. The controlled experiments cover one backbone family through nominal 1B scale. The throughput results depend on optimized H200 implementations. Independent production-scale cost evidence is absent, and larger-model scaling remains untested. In addition, high-order counts may expose fragments of source text, so aggregate statistics should not automatically be treated as privacy-preserving.

COUNT therefore establishes a credible architectural option, not a universal replacement for conventional scaling or fine-tuning.

Its broader contribution is to show that useful language-model memory can be constructed from corpus statistics, stored outside the GPU, and connected to learned representations without requiring an equivalent expansion of trainable parameters.

For model builders with substantial host RAM and constrained accelerator budgets, that is a hardware-and-architecture tradeoff worth measuring.

Cognaptus: Automate the Present, Incubate the Future.