← Blog

Language modeling with 30–86 bytes of context memory per token

Work in progress · Updated 29 October 2026

The problem

A language model's context can cost more memory than its weights. Full attention keeps a separate key and value for every token at every global-attention layer. Longer documents and larger batches keep multiplying that cache. Sparse attention reduces how much the model reads, but often still stores every key and value.

Can we shrink what is stored as well as what is read, while preserving useful prediction? This project explores a shared memory of about 30 bytes per token at 155M parameters, growing to 86 bytes at 963M.

The idea: a compact record that attention can still use

The central question is how little a model needs to remember about each token for attention to stay useful. Instead of a full key and value per token at every attention layer, each token keeps one small, fixed-size record, learned during training and shared by the layers that read long-range memory. At 155M parameters it takes about 30 bytes.

Making the record small is only half the task: reading it must also stay affordable. A search step decides which parts of the history to read in full, so the work per token stays bounded as the context grows. The record also lets the model reuse earlier text directly, which turns out to matter a great deal for prediction (more below).

The experiments below test the combined system: how much memory and work it saves, and how much predictive quality survives that compression.

How much memory does it save?

At 155M parameters, the record and its index use about 30 bytes per token for the whole model. Full attention in the same two global layers uses 1,536 bytes: about 51 times as much cache. As the model grows, additional key groups increase the store to 58 and 86 bytes per token.

That saving translates into serving capacity. A separate 412M-parameter configuration on one A40 accommodates 256 sequences of 1.05M tokens within 48 GB, including bf16 parameters and execution workspace. The comparison below uses models of approximately 415M parameters.

Mamba-3 keeps a fixed-size recurrent state rather than token-addressable history. Our method retains a record for every token. Context limits are the reported fitting points on the tested doubling grid; retaining a history does not by itself establish reliable retrieval across its full length.

How fast does it run?

On the same A40, the 412M serving configuration is compared with a roughly 415M DeepSeek-V4.1-style model. At 1.05M context, prefill is 12.1 times faster and decode is 4.9 times faster. The gap grows with context:

These are the study's implementations, not timings of released pretrained models. The DeepSeek-V4.1 implementation stores compressed entries in FP4, uses PyTorch selection, and uses custom prefill attention with a bf16 working copy; the comparison reflects these implementation choices as well as the architectures.

Our measured full-model prefill stays near 9.9 microseconds per prompt token from roughly 1M to 537M tokens; decode reaches 3.5 ms per token at 537M. Restricting the direct scan keeps the read manageable, although tree descent and summary coverage still add work. The DeepSeek-V4.1 prefill measurements stop at 8.39M, so no measured speedup is claimed against it at 537M.

Training work is lower too: at 155M parameters and a 65K-token window, 0.48 GFLOP/token versus 0.83 for DeepSeek-V4.1, about 42% less. The full comparison follows below. These counts estimate training as three times forward work, excluding recomputation; they are separate from measured throughput.

Does the smaller memory preserve prediction?

Six models of about 155M parameters train on the same tokenizer, data order, optimizer recipe, and 1B-token budget, mixing books, scientific papers, and code. Ours is lower from the start of training and stays lower:

On held-out test text, the table pairs perplexity with training work. Lower is better in every numeric column.

Perplexity falls by 29% on books, 35% on papers, and 48% on code against the best baseline in each domain. These are architecture-style implementations trained for this comparison, with one seed per model. Much of the gain comes from the model reusing tokens it has stored: retrained without that ability, its mean perplexity across the three domains rises from 16.6 to 44.0, worse than three of the five baselines. So the result reflects the whole design, not the compact memory alone.

A separate controlled comparison isolates the memory itself: direct token reuse is turned off in every model, and only the two global-attention layers differ. There, the compact memory stays within the declared ±0.05-nat equivalence margin of full attention in every domain, using paired 90% intervals across test documents.

Does it improve with scale?

Prediction improves in every domain as width and depth grow. Comparisons with the other architectures currently cover only the smallest size.

What comes next

Three things come next. The first is making the model use distant information as reliably as nearby text. The compact memory lets it keep much longer histories; now I want to improve how it finds and uses the parts that matter, and test that on tasks beyond predicting the next word.

The second is faster training. Every result above uses AdamW. In an earlier version of this architecture, a second-order optimizer that approximates the loss curvature with small per-layer factors reached AdamW's loss in 0.63 to 0.81 of the wall-clock time, over the first 150 updates of a 40M-parameter model. The plan is to bring it to the current design and check whether that lead holds over a full training run and at larger sizes.

The third is training at the context length the memory was built for. Today the model trains on 65K-token windows and is served at contexts up to 537M tokens. The goal is training itself over a context of one billion tokens on a single GPU; an earlier version of the architecture already fit one exact training update over a 100-million-token context into 27.5 GiB of GPU memory.