The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation
COUNT tackles the GPU-memory bottleneck of scaling language-model embeddings by replacing learned n-gram tables with precomputed corpus counts stored on cheap host RAM. It constructs input embeddings via a weighted sum of the model's own unembedding vectors, requiring no extra learned parameters and enabling training-free domain adaptation at inference. On 60M--1B-parameter backbones, the authors demonstrate consistent perplexity and downstream gains without increasing GPU memory pressure.
COUNT is a clever engineering solution that delivers measurable pretraining efficiency gains for models up to 1B parameters. The core idea--reframing n-gram memory as read-only host-resident count tables rather than trainable GPU tensors--is well motivated and cleanly executed. However, the work stops at 1B, leaving open whether host-memory bandwidth and the relative signal of n-gram counts remain favorable at larger scales. Moreover, headline compute-equivalence claims rely on fitted scaling laws that are only sparsely sampled.
The technical design is elegant: stupid-backoff scores index the model's unembedding basis $B$, yielding dense input embeddings $f_t = \frac{1}{J}\sum_{j=1}^J\sum_{k=1}^K S(w_k \mid h_{t,j})B[w_k]$ without extra parameters. Throughput measurements are rigorous, showing COUNT retains over 93% of standard transformer training throughput thanks to prefetching and fused embedding reductions. The Pareto analysis convincingly demonstrates that learned n-gram embeddings cannot match COUNT's loss-per-parameter tradeoff because they bloat GPU memory with trainable weights.
The experimental scope is limited to 1B-parameter backbones, so the paper offers no evidence that host-resident n-gram tables remain effective or practical when scaled to multi-billion-parameter models where embedding dimensions and context lengths are larger. Compute-saving claims are derived from fitted scaling laws $\widehat{L}(N,D)=E+\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}}$ over just seven model sizes and three token budgets, yet uncertainty around these fits is not quantified.
Domain-adaptation gains are uneven: biomedical tasks show mixed results, and the method assumes a conventional unembedding matrix, which may not transfer neatly to encoder-decoder or differently tied architectures.
The comparisons against learned n-gram embeddings (Over-Encoding / LongCat) are fair and carefully controlled, including embedding amplification and sparse optimizer updates. COUNT also outperforms output-side late-fusion counterparts such as additive and multiplicative interpolation of stupid-backoff scores. Still, the downstream evaluation leans heavily on likelihood-based BPB rather than discrete accuracy or instruction-following metrics; the authors acknowledge that discrete metrics are noisy for small models, but this choice limits direct comparability to standard instruction-tuned benchmarks. The evidence supports the headline efficiency claims, though generalization to larger-scale regimes remains extrapolation.
Reproducibility is a clear strength. The authors release the full 797 GB Dolma 3 n-gram counts, training code, pretrained checkpoints, and exhaustive appendices specifying table construction, model ladders, and optimization recipes. Throughput benchmarks are similarly well-documented with exact hardware and batch-size settings. One omission is the lack of reported host-system specifications (CPU core counts, memory bandwidth) for the lookup benchmarks, which could affect whether independent researchers achieve the same latency hiding on different hardware.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.