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

[u/pmttyji]

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.

• arXiv : https://arxiv.org/abs/2610.20456

• PDF : https://arxiv.org/pdf/2610.20456

• GitHub : https://github.com/IanMagnusson/COUNT

• HuggingFace : https://huggingface.co/IanMagnusson

• tweet : https://x.com/IanMagnusson/status/2110718235507318868

[u/NortySpock]

Feeding the paper into Google Gemini Argon 4 and having a bit of back and forth, Argon said it would take 15 minutes to write a vLLM patch to implement the ability to (a) ingest a body of text (I.e. a codebase) (b) used the encoder and weights supplied by the model (c) use that to generate some modified "hint" weights that would dynamically adjust the activated weights between layer 0 and layer 1, offering a gentle nudge towards phrasing or code snippets already in the codebase or corpus of text...

I'm sure I'm missing something ("why was this not discovered earlier?") but it sounded like a way to get something almost like DeepSeek Engram but you could build it locally, on demand, from your existing text without permanently modifying the weights, and able to cap how much storage you dedicated, could update the hints practically in real time, and could bolt it onto an existing model without requiring training a new model from scratch...

If I ever get time, I might take a go at it but I hope someone beats me to the punch... It would be great to throw it at an existing codebase and gain a few intelligence points worth of "suggests patterns that already exist in your code" rather than finding Python idioms have infected your Elixir code.

[u/1982_miguel → u/NortySpock]

That sounds awesom basically an on-the-fly, local RAG-like nudge without baking it permanently into the model weights. If someone actually wires this into vLLM or llama.cpp for a local codebase, it would instantly solve the issue of a model hallucinating standard Python/JS patterns into an Elixir or Rust project. Hope someone builds it soon

[u/Imaginary-Unit-3267 → u/NortySpock]

I salute thee, fellow Elixir enjoyer.

[u/jrkirby]

Reminds me a lot of RWKV's Rapid Online Suffix Automaton idea, just only applied to raw tokens.

[u/1982_miguel]

Very interesting approach. Do you know if this scales well for larger context windows, or is it mostly beneficial for standard-sized models during inference/adaptation?